Guarantee masked assistant tokens in diffusion training

README.md

@@ -82,16 +82,22 @@ This is a genuinely **unexplored frontier** in the literature:
 ## Running Training
 
 ```bash
+# CPU smoke: Stage 1 projector path
+python code/train_production.py --stage 1 --epochs 1 --batch_size 1 --grad_accum 1 --num_workers 0 --max_samples 1 --dry_run_batches 1
+
+# CPU smoke: Stage 2 subset path
+python code/train_production.py --stage 2 --resume_from ./vil-dlm-output/stage1_best --dataset_configs ai2d,aokvqa --epochs 1 --batch_size 1 --grad_accum 1 --num_workers 0 --max_samples 8 --dry_run_batches 1
+
 # Stage 1: projector-only alignment
 python code/train_production.py --stage 1 --require_cuda --epochs 1 --batch_size 8 --grad_accum 4
 
 # Stage 2: full-model finetune on the balanced Cauldron mix
 python code/train_production.py --stage 2 --require_cuda --epochs 3 --batch_size 2 --grad_accum 16
 
-# Stage 3a: build the Gemma teacher candidate bank from a Stage 2 checkpoint
+# Stage 3a: build the Gemma teacher candidate bank from a Stage 2 checkpoint (GPU only)
 python code/train_production.py --stage 3a --require_cuda --resume_from ./vil-dlm-output/stage2_best --teacher_batch_size 2
 
-# Stage 3b: sparse KD training from the cached teacher bank
+# Stage 3b: sparse KD training from the cached teacher bank (GPU only)
 python code/train_production.py --stage 3b --require_cuda --resume_from ./vil-dlm-output/stage2_best --epochs 2 --batch_size 2 --grad_accum 16
 
 # Cheap validation gate for any stage
@@ -99,6 +105,7 @@ python code/train_production.py --stage 1 --require_cuda --dry_run_batches 1 --m
 ```
 
 Training now saves checkpoints locally by default. Add `--push_to_hub` only when you want to publish artifacts.
+CPU sessions should stop after the Stage 2 subset smoke test. Stage 3 requires a CUDA GPU because Gemma 4 teacher-bank preparation uses quantized multimodal teacher inference.
 
 ### Hardware Requirements
 - **Stage 1**: A10G (24GB) or T4 (16GB) — only projector gradients (~7M params)

code/train_production.py

@@ -126,11 +126,24 @@ class MDLMScheduler:
     def __init__(self, mask_token_id: int) -> None:
         self.mask_token_id = mask_token_id
 
-    def add_noise(self, input_ids: torch.Tensor, t: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    def add_noise(
+        self,
+        input_ids: torch.Tensor,
+        t: torch.Tensor,
+        eligible_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         batch, length = input_ids.shape
         mask_ratio = 1.0 - torch.cos(t * math.pi / 2)
         mask_ratio = mask_ratio.unsqueeze(1).expand(batch, length)
         mask = torch.rand(batch, length, device=input_ids.device) < mask_ratio
+        if eligible_mask is not None:
+            eligible_mask = eligible_mask.bool()
+            mask = mask & eligible_mask
+            missing_mask = (mask.sum(dim=1) == 0) & (eligible_mask.sum(dim=1) > 0)
+            for batch_idx in torch.nonzero(missing_mask, as_tuple=False).flatten():
+                eligible_positions = torch.nonzero(eligible_mask[batch_idx], as_tuple=False).flatten()
+                chosen = eligible_positions[torch.randint(eligible_positions.numel(), (1,), device=input_ids.device)]
+                mask[batch_idx, chosen] = True
         noisy_ids = input_ids.clone()
         noisy_ids[mask] = self.mask_token_id
         return noisy_ids, mask
@@ -209,7 +222,8 @@ class ViLDLM(nn.Module):
             loss_mask = attention_mask
 
         t = self.scheduler.sample_timesteps(batch_size, device)
-        noisy_ids, noise_mask = self.scheduler.add_noise(input_ids, t)
+        eligible_mask = (loss_mask > 0) & (attention_mask > 0)
+        noisy_ids, noise_mask = self.scheduler.add_noise(input_ids, t, eligible_mask=eligible_mask)
         inputs_embeds, full_attention_mask = self.prepare_multimodal_inputs(
             pixel_values=pixel_values,
             input_ids=noisy_ids,
@@ -218,7 +232,7 @@ class ViLDLM(nn.Module):
         outputs = self.lm(inputs_embeds=inputs_embeds, attention_mask=full_attention_mask)
         text_logits = outputs.logits[:, self.num_patches :, :]
 
-        active_mask = noise_mask.float() * loss_mask.float()
+        active_mask = noise_mask.float() * eligible_mask.float()
         if active_mask.sum() == 0:
             loss = torch.tensor(0.0, device=device, requires_grad=True)
         else: