diff --git a/fastvideo/models/dits/cosmos3.py b/fastvideo/models/dits/cosmos3.py index 71eaf214b..391e4ac0b 100644 --- a/fastvideo/models/dits/cosmos3.py +++ b/fastvideo/models/dits/cosmos3.py @@ -934,6 +934,44 @@ class Cosmos3VFMTransformer(BaseDiT): output["last_hidden_state"] = last_hidden_state return output + def reason_forward( + self, + inputs_embeds: torch.Tensor, + position_ids: torch.Tensor, + deepstack_embeds: list[torch.Tensor] | None = None, + visual_pos_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Reasoner (causal) backbone forward over one text(+image) sequence. + + All tokens are und (causal); the gen pathway is empty. Mirrors the + framework ``reasoner_forward``: run the und layers, optionally injecting + per-layer deepstack visual embeds at ``visual_pos_mask`` for the first + ``len(deepstack_embeds)`` layers (Qwen3-VL deepstack), then final + ``norm``. Returns post-norm hidden states ``[N, hidden]`` (apply + ``lm_head`` for logits). Used by image-conditioned reasoning, with the + text/vision embeds + ``position_ids`` produced by the multimodal prefill. + + Args: + inputs_embeds: ``[N, hidden]`` text embeds with image features + already scattered in at the image-token positions. + position_ids: ``[3, N]`` (or ``[N]``) mRoPE positions. + deepstack_embeds: per-deepstack-layer ``[n_visual, hidden]`` embeds. + visual_pos_mask: ``[N]`` bool mask of visual token positions. + """ + device = self.embed_tokens.weight.device + target_dtype = inputs_embeds.dtype + cos, sin = self.rotary_emb(position_ids, device=device, dtype=target_dtype) # [N, head_dim] + und = inputs_embeds + gen = inputs_embeds.new_zeros((0, self.hidden_size)) + cos_gen = cos.new_zeros((0, cos.shape[-1])) + sin_gen = sin.new_zeros((0, sin.shape[-1])) + for li, layer in enumerate(self.layers): + und, _ = layer(und, gen, cos, sin, cos_gen, sin_gen) + if deepstack_embeds is not None and visual_pos_mask is not None and li < len(deepstack_embeds): + und = und.clone() + und[visual_pos_mask] = und[visual_pos_mask] + deepstack_embeds[li].to(dtype=und.dtype) + return self.norm(und) # [N, hidden] + def _unpack_action( self, packed_preds: torch.Tensor, diff --git a/tests/local_tests/cosmos3/test_cosmos3_reasoning_parity.py b/tests/local_tests/cosmos3/test_cosmos3_reasoning_parity.py index 346782437..72abc0e39 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_reasoning_parity.py +++ b/tests/local_tests/cosmos3/test_cosmos3_reasoning_parity.py @@ -121,3 +121,39 @@ class TestCosmos3ReasoningParity: fv = cosmos3_generate_reasoner_text(dit, input_ids.tolist(), max_new_tokens=8) print(f"\n[greedy_reason seed={seed}] framework={fw} native={fv}") assert fv == fw, f"token mismatch: native={fv} framework={fw}" + + def test_deepstack_reasoner_forward_matches_framework(self): + """The deepstack reasoner backbone (image-conditioned reasoning's one new + native piece) matches the framework ``reasoner_forward`` given identical + prefill inputs (inputs_embeds + positions + per-layer deepstack embeds + + visual mask).""" + from cosmos_framework.model.vfm.mot.unified_mot import ReasonerKVCache + + vfm, dit = self._build() + hidden = dit.hidden_size + n_layers = dit.num_hidden_layers + seq = 10 + torch.manual_seed(5) + inputs_embeds = torch.randn(seq, hidden) + position_ids = torch.arange(seq).unsqueeze(0).expand(3, -1).contiguous() # [3, seq] + visual_pos_mask = torch.zeros(seq, dtype=torch.bool) + visual_pos_mask[2:6] = True # 4 "image" tokens + k = int(visual_pos_mask.sum()) + deepstack = [torch.randn(k, hidden) for _ in range(n_layers)] # one per layer + + model = vfm.language_model.model + cache = ReasonerKVCache.empty(num_layers=len(model.layers)) + with torch.no_grad(): + hid_fw = model.reasoner_forward( + input_ids=None, + inputs_embeds=inputs_embeds.unsqueeze(0), + position_ids=position_ids.unsqueeze(1), # [3, B=1, seq] + visual_pos_masks=visual_pos_mask.unsqueeze(0), + deepstack_visual_embeds=deepstack, + cache=cache, + )[0] # [seq, hidden] + hid_fv = dit.reason_forward(inputs_embeds, position_ids, deepstack, visual_pos_mask) # [seq, hidden] + assert hid_fw.shape == hid_fv.shape, f"shape fw={hid_fw.shape} fv={hid_fv.shape}" + mx, mn = _diffs(hid_fv, hid_fw) + print(f"\n[deepstack_reasoner] hidden max abs diff = {mx:.3e} mean abs diff = {mn:.3e}") + torch.testing.assert_close(hid_fv, hid_fw, atol=1e-4, rtol=1e-3)