fix apply_rotary_pos_emb_thd missing cp_group#146
Conversation
There was a problem hiding this comment.
Code Review
This pull request updates _apply_rotary_pos_emb_thd in src/mcore_bridge/patcher.py to retrieve and pass the context parallel group (cp_group) when calling _origin_apply_rotary_pos_emb_thd. However, retrieving the context parallel group unconditionally when context parallel is not initialized (e.g., when cp_size == 1) will cause an AssertionError and crash the execution. It is recommended to retrieve the group conditionally only when cp_size > 1.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| cp_size = mpu.get_context_parallel_world_size() | ||
| cp_group = mpu.get_context_parallel_group() |
There was a problem hiding this comment.
Calling mpu.get_context_parallel_group() unconditionally when context parallel is not initialized (which is the case when cp_size == 1) will raise an AssertionError: context parallel group is not initialized in Megatron-LM. This will cause any run without context parallel to crash.\n\nWe should only retrieve the context parallel group if cp_size > 1.
cp_size = mpu.get_context_parallel_world_size()\n cp_group = mpu.get_context_parallel_group() if cp_size > 1 else None
https://github.com/NVIDIA/Megatron-LM/blob/core_v0.18.0/megatron/core/models/common/embeddings/rope_utils.py#L200-L201