[bugfix] fix MOSS-Audio deepstack_input_embeds initialization in PP (#47607)

Signed-off-by: Yan Ma <yan.ma@intel.com>
2026-07-06 17:15:50 +08:00
parent 26c754d847
commit 90ce3a09be
1 changed files with 4 additions and 0 deletions
@@ -862,6 +862,10 @@ class MossQwen3ForCausalLM(Qwen3ForCausalLM):
            batch_size, dtype, device
        )
        for layer_idx in self.deepstack_inject_layer_indices:
+            # Non-first PP ranks only receive DeepStack payloads for layers
+            # at or after their local start layer.
+            if layer_idx < self.model.start_layer:
+                continue
            intermediate_tensors[f"deepstack_input_embeds_{layer_idx}"] = torch.zeros(
                (batch_size, self.config.hidden_size),
                dtype=dtype,