Buckets:
| """Canonical full-GT text/bit replay loss for training modules.""" | |
| from __future__ import annotations | |
| import hashlib | |
| import math | |
| from dataclasses import dataclass | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from .labels import canonical_autonomous_text_labels | |
| from .modeling import model_autocast, unwrap_ar_model | |
| def select_noisy_sample_ids( | |
| sample_ids: list[str], noisy_fraction: float, *, seed: int | |
| ) -> frozenset[str]: | |
| if not 0.0 <= noisy_fraction <= 1.0: | |
| raise ValueError("noisy_fraction must be in [0, 1]") | |
| if len(sample_ids) != len(set(sample_ids)): | |
| raise ValueError("sample_ids must be unique") | |
| ranked = sorted( | |
| sample_ids, | |
| key=lambda sample_id: ( | |
| hashlib.sha256(f"{seed}\0{sample_id}".encode()).digest(), | |
| sample_id, | |
| ), | |
| ) | |
| return frozenset(ranked[: int(len(ranked) * noisy_fraction)]) | |
| def unpack_visual_codes( | |
| packed_or_item: torch.Tensor | dict, | |
| shape: tuple[int, ...] | list[int] | None = None, | |
| ) -> torch.Tensor: | |
| if isinstance(packed_or_item, dict): | |
| if shape is not None: | |
| raise ValueError("shape is implicit when unpacking a cache item") | |
| packed = packed_or_item["visual_codes_packed"] | |
| shape = packed_or_item["visual_codes_shape"] | |
| else: | |
| if shape is None: | |
| raise ValueError("packed visual codes require their original shape") | |
| packed = packed_or_item | |
| shape = tuple(int(value) for value in shape) | |
| values = np.unpackbits( | |
| packed.cpu().numpy(), | |
| axis=1, | |
| count=math.prod(shape[1:]), | |
| bitorder="little", | |
| ) | |
| return torch.from_numpy(values).reshape(shape).to(torch.uint8) | |
| def pack_visual_codes(codes: torch.Tensor) -> tuple[torch.Tensor, tuple[int, ...]]: | |
| if codes.ndim < 1 or codes.dtype != torch.uint8: | |
| raise ValueError("visual codes must be a uint8 tensor with a position axis") | |
| if codes.numel() and not bool(((codes == 0) | (codes == 1)).all()): | |
| raise ValueError("visual codes must contain only binary factors") | |
| shape = tuple(codes.shape) | |
| flattened = codes.detach().cpu().numpy().reshape(shape[0], -1) | |
| packed = np.packbits(flattened, axis=1, bitorder="little") | |
| return torch.from_numpy(packed), shape | |
| def _corrupt_generated_codes( | |
| codes: torch.Tensor, | |
| generated: torch.Tensor, | |
| *, | |
| probability: float, | |
| probability_mode: str, | |
| seed: int, | |
| update: int, | |
| sample_id: str, | |
| ) -> tuple[torch.Tensor, torch.Tensor]: | |
| if probability_mode not in {"fixed", "uniform"}: | |
| raise ValueError("probability_mode must be fixed or uniform") | |
| if not probability: | |
| return codes, torch.zeros_like(codes, dtype=torch.bool) | |
| value = int.from_bytes( | |
| hashlib.sha256(f"{seed}\0{update}\0{sample_id}".encode()).digest()[:8], | |
| "big", | |
| ) & ((1 << 63) - 1) | |
| generator = torch.Generator(device=codes.device).manual_seed(value) | |
| effective_probability = probability | |
| if probability_mode == "uniform": | |
| effective_probability = float( | |
| torch.rand((), device=codes.device, generator=generator).cpu() | |
| ) * probability | |
| flips = ( | |
| torch.rand(codes.shape, device=codes.device, generator=generator) | |
| < effective_probability | |
| ) | |
| flips &= generated[(slice(None),) + (None,) * (codes.ndim - 1)] | |
| return codes ^ flips.to(codes.dtype), flips | |
| class ReplaySFTOutput: | |
| loss: torch.Tensor | |
| text_loss: torch.Tensor | |
| bit_loss: torch.Tensor | |
| text_tokens: int | |
| visual_bits: int | |
| flipped_bits: torch.Tensor | |
| def replay_sft_loss( | |
| model, | |
| item: dict, | |
| *, | |
| image_gen_id: int, | |
| text_weight: float, | |
| bit_weight: float, | |
| image_gen_token_weight: float, | |
| noisy: bool, | |
| flip_probability: float, | |
| flip_probability_mode: str = "fixed", | |
| seed: int, | |
| update: int, | |
| ) -> ReplaySFTOutput: | |
| """Compute canonical SFT only on one complete dataset GT trajectory.""" | |
| ar = unwrap_ar_model(model) | |
| device = next(ar.parameters()).device | |
| input_ids_1d = item["input_ids"].to(device=device, dtype=torch.long) | |
| labels = canonical_autonomous_text_labels( | |
| input_ids_1d, | |
| item["labels"].to(device=device, dtype=torch.long), | |
| ).unsqueeze(0) | |
| input_ids = input_ids_1d.unsqueeze(0) | |
| attention = torch.ones_like(input_ids) | |
| grids = item["image_grid_thw"].to(device=device, dtype=torch.long) | |
| codes = unpack_visual_codes(item).to(device=device, dtype=torch.long) | |
| generated_flat = item["generated_flat_mask"].to(device=device, dtype=torch.bool) | |
| generated_sequence = item["generated_sequence_mask"].to( | |
| device=device, dtype=torch.bool | |
| ).unsqueeze(0) | |
| input_codes, flips = _corrupt_generated_codes( | |
| codes, | |
| generated_flat, | |
| probability=flip_probability if noisy else 0.0, | |
| probability_mode=flip_probability_mode, | |
| seed=seed, | |
| update=update, | |
| sample_id=item["sample_id"], | |
| ) | |
| positions, _ = ar.model.get_rope_index( | |
| input_ids, image_grid_thw=grids, attention_mask=attention | |
| ) | |
| visual_carriers = torch.zeros_like(generated_sequence) | |
| visual_carriers[:, :-1] = generated_sequence[:, 1:] | |
| kwargs = { | |
| "input_ids": input_ids, | |
| "attention_mask": attention, | |
| "image_grid_thw": grids, | |
| "visual_codes": input_codes, | |
| "position_ids": positions, | |
| "compute_visual_logits": True, | |
| } | |
| if hasattr(ar, "prepare_visual_features"): | |
| with model_autocast(model): | |
| kwargs["visual_features"] = ar.prepare_visual_features(input_codes) | |
| kwargs["visual_output_mask"] = visual_carriers | |
| with model_autocast(model): | |
| output = ar.forward_interleaved(**kwargs) | |
| shifted = labels[:, 1:] | |
| text_mask = shifted != -100 | |
| if not text_mask.any(): | |
| raise ValueError("GT replay trajectory has no supervised text") | |
| with model_autocast(model): | |
| text_logits = ar.project_text( | |
| output.hidden_states[:, :-1][text_mask] | |
| ).float() | |
| text_targets = shifted[text_mask] | |
| text_nll = F.cross_entropy(text_logits, text_targets, reduction="none") | |
| weights = torch.ones_like(text_nll) | |
| weights[text_targets == image_gen_id] = image_gen_token_weight | |
| text_loss = (text_nll * weights).sum() / weights.sum() | |
| visual_output = output.visual_logits | |
| if "visual_features" not in kwargs: | |
| visual_output = visual_output[visual_carriers] | |
| levels = 1 + len(ar.model.visual.config.deepstack_visual_indexes) | |
| spatial = ar.model.visual.spatial_merge_unit | |
| bits = ar.model.visual.config.bsq_dim | |
| visual_logits = visual_output.reshape(-1, levels, spatial, bits, 2).float() | |
| visual_targets = codes[generated_flat] | |
| if visual_logits.shape[:-1] != visual_targets.shape: | |
| raise ValueError("GT replay visual logits and clean targets differ") | |
| bit_loss = F.cross_entropy( | |
| visual_logits.reshape(-1, 2), visual_targets.reshape(-1), reduction="mean" | |
| ) | |
| return ReplaySFTOutput( | |
| loss=text_weight * text_loss + bit_weight * bit_loss, | |
| text_loss=text_loss, | |
| bit_loss=bit_loss, | |
| text_tokens=text_nll.numel(), | |
| visual_bits=visual_targets.numel(), | |
| flipped_bits=flips.sum(), | |
| ) | |
| __all__ = [ | |
| "canonical_autonomous_text_labels", | |
| "ReplaySFTOutput", | |
| "pack_visual_codes", | |
| "replay_sft_loss", | |
| "select_noisy_sample_ids", | |
| "unpack_visual_codes", | |
| ] | |
Xet Storage Details
- Size:
- 7.48 kB
- Xet hash:
- cc61b72b849f3ccd9adc950185e1c83e965e90af4d2e8735b870e140be3d75b7
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.