Buckets:
| """Joint text/BSQ teacher-forcing objective for interleaved UniAR SFT.""" | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn.functional as F | |
| def unwrap_ar_model(model): | |
| """Return the UniAR base model from an optional PEFT wrapper.""" | |
| return model.get_base_model() if hasattr(model, "get_base_model") else model | |
| def predecessor_mask(target_mask: torch.Tensor) -> torch.Tensor: | |
| """Select each causal position immediately preceding a target position.""" | |
| result = torch.zeros_like(target_mask, dtype=torch.bool) | |
| result[:, :-1] = target_mask[:, 1:] | |
| return result | |
| def encode_visual_codes(ar_model, pixel_values: torch.Tensor, image_grid_thw: torch.Tensor) -> torch.Tensor: | |
| """Encode all images once into grouped multi-level 0/1 BSQ targets.""" | |
| ar_model = unwrap_ar_model(ar_model) | |
| vision = ( | |
| ar_model.model.visual | |
| if hasattr(ar_model, "model") and hasattr(ar_model.model, "visual") | |
| else ar_model | |
| ) | |
| device = next(vision.parameters()).device | |
| vision.eval() | |
| main, deep = vision( | |
| pixel_values.to(device=device, dtype=vision.dtype), | |
| grid_thw=image_grid_thw.to(device), | |
| bsq_only=True, | |
| direct_indices=True, | |
| ) | |
| spatial = vision.spatial_merge_unit | |
| bits = [main, *deep] | |
| levels = [value.reshape(-1, spatial, vision.config.bsq_dim) for value in bits] | |
| return torch.stack(levels, dim=1).to(torch.uint8) | |
| def joint_sft_loss( | |
| model, | |
| batch: dict, | |
| visual_codes: torch.Tensor, | |
| *, | |
| image_gen_id: int, | |
| visual_loss_weight: float = 1.0, | |
| marker_token_weight: float = 20.0, | |
| visual_input_flip_prob: float = 0.0, | |
| return_predictions: bool = False, | |
| ) -> dict: | |
| """Compute normalized language CE and visual bitwise CE for one trajectory.""" | |
| ar_model = unwrap_ar_model(model) | |
| device = ar_model.device | |
| input_ids = batch["input_ids"].to(device) | |
| attention_mask = batch["attention_mask"].to(device) | |
| labels = batch["labels"].to(device) | |
| image_grid_thw = batch["image_grid_thw"].to(device) | |
| generated_sequence_mask = batch["generated_sequence_mask"].to(device) | |
| generated_flat_mask = batch["generated_flat_mask"].to(device) | |
| visual_codes = visual_codes.to(device=device, dtype=torch.long) | |
| if not 0.0 <= visual_input_flip_prob <= 1.0: | |
| raise ValueError("visual_input_flip_prob must be in [0, 1]") | |
| flip_strength = visual_codes.new_tensor(0.0, dtype=torch.float32) | |
| input_codes = visual_codes | |
| if visual_input_flip_prob > 0.0: | |
| # UniAR samples one corruption strength per example, uniformly from zero | |
| # to the configured maximum. Corrupt only generated-state teacher inputs; | |
| # observation codes and clean visual targets remain unchanged. | |
| flip_strength = torch.rand((), device=device) * visual_input_flip_prob | |
| generated_bits = generated_flat_mask[:, None, None, None] | |
| flip = (torch.rand(visual_codes.shape, device=device) < flip_strength) & generated_bits | |
| input_codes = visual_codes ^ flip.to(visual_codes.dtype) | |
| vision = ar_model.model.visual | |
| backbone = ar_model.forward_interleaved( | |
| input_ids=input_ids, | |
| attention_mask=attention_mask, | |
| image_grid_thw=image_grid_thw, | |
| visual_codes=input_codes, | |
| ) | |
| hidden_prenorm = backbone.hidden_states | |
| shifted_labels = labels[:, 1:] | |
| text_positions = shifted_labels != -100 | |
| if not text_positions.any(): | |
| raise ValueError("trajectory has no supervised assistant text") | |
| text_hidden = hidden_prenorm[:, :-1][text_positions] | |
| text_logits = ar_model.project_text(text_hidden).float() | |
| text_targets = shifted_labels[text_positions] | |
| text_token_losses = F.cross_entropy(text_logits, text_targets, reduction="none") | |
| text_predictions = text_logits.argmax(-1) | |
| marker_mask = text_targets == image_gen_id | |
| text_weights = torch.ones_like(text_token_losses) | |
| text_weights[marker_mask] = marker_token_weight | |
| text_loss = (text_token_losses * text_weights).sum() / text_weights.sum() | |
| marker_loss = ( | |
| text_token_losses[marker_mask].mean() | |
| if marker_mask.any() | |
| else text_token_losses.new_tensor(float("nan")) | |
| ) | |
| visual_predict_positions = predecessor_mask(generated_sequence_mask) | |
| visual_logits_all = backbone.visual_logits | |
| vision_cfg = vision.config | |
| levels = 1 + len(vision_cfg.deepstack_visual_indexes) | |
| visual_logits = visual_logits_all[visual_predict_positions].reshape( | |
| -1, levels, vision.spatial_merge_unit, vision_cfg.bsq_dim, 2 | |
| ) | |
| visual_targets = visual_codes[generated_flat_mask] | |
| if visual_logits.shape[:-1] != visual_targets.shape: | |
| raise ValueError( | |
| f"visual alignment mismatch: logits={visual_logits.shape}, " | |
| f"targets={visual_targets.shape}" | |
| ) | |
| visual_loss = F.cross_entropy( | |
| visual_logits.reshape(-1, 2).float(), visual_targets.reshape(-1) | |
| ) | |
| visual_predictions = visual_logits.argmax(-1) | |
| result = { | |
| "loss": text_loss + visual_loss_weight * visual_loss, | |
| "text_loss": text_loss.detach(), | |
| "marker_loss": marker_loss.detach(), | |
| "visual_loss": visual_loss.detach(), | |
| "text_accuracy": (text_predictions == text_targets).float().mean().detach(), | |
| "visual_accuracy": (visual_predictions == visual_targets).float().mean().detach(), | |
| "marker_accuracy": ( | |
| (text_predictions[marker_mask] == text_targets[marker_mask]).float().mean().detach() | |
| if marker_mask.any() | |
| else text_loss.new_tensor(float("nan")) | |
| ), | |
| "text_tokens": int(text_targets.numel()), | |
| "visual_bits": int(visual_targets.numel()), | |
| "markers": int(marker_mask.sum()), | |
| "visual_input_flip_strength": flip_strength.detach(), | |
| } | |
| if return_predictions: | |
| result["visual_predictions"] = visual_predictions.detach() | |
| result["visual_targets"] = visual_targets.detach() | |
| return result | |
Xet Storage Details
- Size:
- 6.01 kB
- Xet hash:
- e2ec4aeaf24b0ade7d6a59fdbfb95b4c3270269c29ff341612af89ff4a144993
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.