debug

yeyu-nvidia · yeyu-nvidia · commit aa328ed57cba · 2025-10-03T15:07:08.000-07:00
Signed-off-by: Ye Yu &lt;yeyu@nvidia.com&gt;
diff --git a/modelopt/torch/speculative/plugins/megatron_eagle.py b/modelopt/torch/speculative/plugins/megatron_eagle.py
@@ -1404,6 +1404,8 @@ def pseudo_speculative_generate(
             )
             if self.config.sequence_parallel:
                 gathered_embedding = gather_from_sequence_parallel_region(eagle_inputs["embedding"])
+            else:
+                gathered_embedding = eagle_inputs["embedding"]
             if self.eagle_config.parallel_draft_step > 1:
                 # Replace dummy hidden_states with embedding for mask tokens
                 padded_hidden_states[

Original file line number	Diff line number	Diff line change
`@@ -1404,6 +1404,8 @@ def pseudo_speculative_generate(`
`1404`	`1404`	`)`
`1405`	`1405`	`if self.config.sequence_parallel:`
`1406`	`1406`	`gathered_embedding = gather_from_sequence_parallel_region(eagle_inputs["embedding"])`
	`1407`	`+ else:`
	`1408`	`+ gathered_embedding = eagle_inputs["embedding"]`
`1407`	`1409`	`if self.eagle_config.parallel_draft_step > 1:`
`1408`	`1410`	`# Replace dummy hidden_states with embedding for mask tokens`
`1409`	`1411`	`padded_hidden_states[`