diff --git a/difflet/backends/neuron/hunyuan_video/transformer.py b/difflet/backends/neuron/hunyuan_video/transformer.py index 74ee7ba7..126dfbb7 100644 --- a/difflet/backends/neuron/hunyuan_video/transformer.py +++ b/difflet/backends/neuron/hunyuan_video/transformer.py @@ -189,7 +189,9 @@ def get_example_inputs(self) -> tuple[torch.Tensor, ...]: pooled_projections, guidance)`` at the static shape: the latents ``[B, C, F', H/8, W/8]``, the per-sample timestep, the right-padded Llama embeddings ``[B, text_seq_len, text_embed_dim]`` with an int64 mask, the CLIP pooled projection and the guidance - (scale x 1000). Exactly the per-step shapes the orchestrator sends, so the warm-up + (scale x 1000); the timestep in fp32, as diffusers passes it, the guidance in ``dtype`` + (diffusers builds it in the transformer's dtype). Exactly the per-step shapes the + orchestrator sends, so the warm-up compiles the only block graphs the run needs. Under CP or SP every rank receives the whole latent and the module scatters it. """ @@ -200,7 +202,7 @@ def get_example_inputs(self) -> tuple[torch.Tensor, ...]: return ( torch.zeros(batch, int(cfg.in_channels), cfg.latent_frames, cfg.latent_height, cfg.latent_width, dtype=dt), - torch.ones(batch, dtype=dt), + torch.ones(batch, dtype=torch.float32), torch.zeros(batch, text, int(cfg.text_embed_dim), dtype=dt), mask, torch.zeros(batch, int(cfg.pooled_projection_dim), dtype=dt), diff --git a/difflet/backends/neuron/wan/transformer.py b/difflet/backends/neuron/wan/transformer.py index d2921fd8..d57434d0 100644 --- a/difflet/backends/neuron/wan/transformer.py +++ b/difflet/backends/neuron/wan/transformer.py @@ -118,7 +118,8 @@ def teacache_mod_input(self, hidden_states: torch.Tensor, timestep: torch.Tensor timestep.to(self.device), self._teacache_text) def get_example_inputs(self) -> tuple[torch.Tensor, ...]: - """``(latents, timestep, text embeddings)`` at the static shape, in ``dtype``. + """``(latents, timestep, text embeddings)`` at the static shape, in ``dtype`` except the + timestep, which the DiT takes in fp32 as diffusers passes it. These are exactly the per-step shapes the orchestrator sends (two batch-1 forwards per CFG step), so the warm-up compiles the only block graph the run needs. With CFG @@ -136,7 +137,7 @@ def get_example_inputs(self) -> tuple[torch.Tensor, ...]: cfg.latent_width, dtype=self.dtype, ), - torch.zeros(batch, dtype=self.dtype), + torch.zeros(batch, dtype=torch.float32), torch.zeros(batch, int(cfg.text_seq_len), int(cfg.text_dim), dtype=self.dtype), ) diff --git a/difflet/models/hunyuan_video/neuron_application.py b/difflet/models/hunyuan_video/neuron_application.py index 4cf67e62..b4c73dd2 100644 --- a/difflet/models/hunyuan_video/neuron_application.py +++ b/difflet/models/hunyuan_video/neuron_application.py @@ -210,6 +210,10 @@ class _OnDeviceHunyuanVideoTransformer(OnDeviceDiT): probe's per-step signal: the orchestrator's fused branch. """ + # As diffusers passes it: the scheduler's timestep unrounded. bf16 holds only multiples of 4 + # between 512 and 1024, which shifts the sinusoidal timestep embedding at every step. + timestep_dtype = torch.float32 + def __init__(self, module, config: NeuronHunyuanVideoConfig, dtype, device, probe: _HunyuanVideoTeaCacheProbe | None = None) -> None: super().__init__(module, config, dtype, device) @@ -223,7 +227,12 @@ def __call__(self, bundle: HunyuanVideoDiTInputBundle) -> torch.Tensor: if not isinstance(bundle, HunyuanVideoDiTInputBundle): raise TypeError(f"expected a HunyuanVideoDiTInputBundle, got {type(bundle).__name__}") _check_dit_shapes(self.config, bundle) - return self._forward(bundle.as_model_inputs()) + # The timestep goes to the device in fp32, not through _forward's cast to the model + # dtype; the other inputs follow it already on the device, keeping the module's order. + hidden_states, timestep, *rest = bundle.as_model_inputs() + timestep = timestep.to(self.timestep_dtype).contiguous().to(self.device) + return self._forward((hidden_states,), + extra=(timestep, *(self._device_input(t) for t in rest))) def _device_input(self, tensor: torch.Tensor) -> torch.Tensor: if not tensor.is_floating_point(): @@ -233,8 +242,8 @@ def _device_input(self, tensor: torch.Tensor) -> torch.Tensor: def teacache_delta(self, bundle: HunyuanVideoDiTInputBundle) -> torch.Tensor: if self.probe is None: raise RuntimeError("no TeaCache probe: adaptive TeaCache (teacache_speedup) is off") - return self.probe.teacache_delta( - *(t.contiguous() for t in (bundle.hidden_states, bundle.timestep, + return self.probe.teacache_delta( # the probe keeps its model-dtype timestep + *(t.contiguous() for t in (bundle.hidden_states, bundle.timestep.to(self.dtype), bundle.pooled_projections, bundle.guidance))) @@ -399,7 +408,7 @@ def phase_seconds(self) -> dict[str, float]: def dit_input_contract(self) -> dict[str, dict[str, Any]]: """The Trainium contract's six tensors at the static shape (``dit_input_shapes``).""" - dtypes = {"encoder_attention_mask": torch.int64} + dtypes = {"encoder_attention_mask": torch.int64, "timestep": torch.float32} return {name: {"shape": shape, "dtype": dtypes.get(name, self.dtype)} for name, shape in dit_input_shapes(self.config).items()} diff --git a/difflet/models/hunyuan_video/pipeline.py b/difflet/models/hunyuan_video/pipeline.py index 1152768e..1bf969d1 100644 --- a/difflet/models/hunyuan_video/pipeline.py +++ b/difflet/models/hunyuan_video/pipeline.py @@ -233,11 +233,11 @@ def _denoise( prev_mod_handle: torch.Tensor | None = None for step_index, timestep in enumerate(timesteps): model_dtype = _component_dtype(self.transformer, self.dtype) - timestep_batch = _batch_timestep( + timestep_batch = _batch_timestep( # the DiT's own timestep dtype, if it declares one timestep, batch_size=latents.shape[0], device=latents.device, - dtype=model_dtype, + dtype=getattr(self.transformer, "timestep_dtype", model_dtype), ) model_bundle = HunyuanVideoDiTInputBundle( hidden_states=latents.to(dtype=model_dtype), diff --git a/difflet/models/wan/neuron_application.py b/difflet/models/wan/neuron_application.py index 96a73317..77568f28 100644 --- a/difflet/models/wan/neuron_application.py +++ b/difflet/models/wan/neuron_application.py @@ -224,13 +224,21 @@ class _OnDeviceTransformer(OnDeviceDiT): expert's adapter reports to the high-noise one's, so one hook times every DiT forward). """ + # As diffusers passes it: the scheduler's timestep unrounded. bf16 holds only multiples of 4 + # between 512 and 1024, which shifts the sinusoidal timestep embedding at every step. + timestep_dtype = torch.float32 + def __init__(self, module, config, dtype, device, hook_from=None): super().__init__(module, config, dtype, device) self.hook_from = hook_from def __call__(self, hidden_states, timestep, encoder_hidden_states) -> torch.Tensor: _check_dit_shapes(self.config, hidden_states, timestep, encoder_hidden_states) - return self._forward((hidden_states, timestep, encoder_hidden_states)) + # The timestep goes to the device in fp32, not through _forward's cast to the model + # dtype; the text states follow it already on the device, keeping the module's order. + timestep = timestep.to(self.timestep_dtype).contiguous().to(self.device) + return self._forward((hidden_states,), + extra=(timestep, self._device_input(encoder_hidden_states))) def _forward_hook(self): return self.forward_hook if self.hook_from is None else self.hook_from.forward_hook @@ -529,7 +537,7 @@ def dit_input_contract(self) -> dict[str, dict[str, Any]]: ), "dtype": self.dtype, }, - "timestep": {"shape": (batch,), "dtype": self.dtype}, + "timestep": {"shape": (batch,), "dtype": torch.float32}, "encoder_hidden_states": { "shape": (batch, int(cfg.text_seq_len), int(cfg.text_dim)), "dtype": self.dtype, diff --git a/tests/unit/backends/test_neuron_hv1_app.py b/tests/unit/backends/test_neuron_hv1_app.py index 5f90a430..afbe3659 100644 --- a/tests/unit/backends/test_neuron_hv1_app.py +++ b/tests/unit/backends/test_neuron_hv1_app.py @@ -190,6 +190,7 @@ def test_application_accepts_the_trainium_cli_kwargs(tiny_pipeline): assert app.shape == SHAPE contract = app.dit_input_contract() assert contract["encoder_attention_mask"] == {"shape": (1, TEXT), "dtype": torch.int64} + assert contract["timestep"]["dtype"] == torch.float32 assert contract["hidden_states"]["shape"] == (1, 4, 5, 8, 8) diff --git a/tests/unit/backends/test_neuron_hv1_dit.py b/tests/unit/backends/test_neuron_hv1_dit.py index 71da6bfb..ce5639d6 100644 --- a/tests/unit/backends/test_neuron_hv1_dit.py +++ b/tests/unit/backends/test_neuron_hv1_dit.py @@ -99,6 +99,7 @@ def test_example_inputs_are_the_static_shape(): assert latents.shape == (1, 16, 16, 40, 64) and latents.dtype == torch.bfloat16 assert text.shape == (1, 256, 4096) and pooled.shape == (1, 768) assert timestep.shape == guidance.shape == (1,) + assert timestep.dtype == torch.float32 # unrounded, as diffusers passes it assert torch.equal(guidance, torch.tensor([6000.0], dtype=torch.bfloat16)) # 6016 in bf16 assert mask.dtype == torch.int64 and mask.shape == (1, 256) and int(mask.sum()) == 64 diff --git a/tests/unit/backends/test_neuron_wan_w1.py b/tests/unit/backends/test_neuron_wan_w1.py index 3793aa95..659828ce 100644 --- a/tests/unit/backends/test_neuron_wan_w1.py +++ b/tests/unit/backends/test_neuron_wan_w1.py @@ -286,7 +286,7 @@ def test_example_inputs_match_the_contract(tiny): ) assert _shapes_and_dtypes(real.get_example_inputs()) == [ ((1, 16, 3, 60, 104), torch.bfloat16), - ((1,), torch.bfloat16), + ((1,), torch.float32), # the timestep unrounded, as diffusers passes it ((1, 512, 4096), torch.bfloat16), ] @@ -423,9 +423,11 @@ def __call__(self, *args): sent.append((latents, timestep, text)) assert len(received) == 2 for args, originals in zip(received, sent, strict=True): - for got, original in zip(args, originals, strict=True): - assert got.is_contiguous() and got.dtype is torch.bfloat16 - assert torch.equal(got, original.to(torch.bfloat16)) + # The timestep (index 1) arrives in fp32, as diffusers passes it; the rest in bf16. + for i, (got, original) in enumerate(zip(args, originals, strict=True)): + want = torch.float32 if i == 1 else torch.bfloat16 + assert got.is_contiguous() and got.dtype is want + assert torch.equal(got, original.to(want)) assert received[1][1].stride() == (1,) # the batch-2 timestep was materialised assert len(seconds) == 2 and all(s >= 0.0 for s in seconds)