From d670466dde1b2a804646a637256f791dc3d3f84b Mon Sep 17 00:00:00 2001 From: Parth Mannan Date: Tue, 3 Feb 2026 18:07:39 -0800 Subject: [PATCH] Bug fix --- megatron/core/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/megatron/core/utils.py b/megatron/core/utils.py index 9ef95285a2b..036589b209a 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -2094,8 +2094,8 @@ def get_thd_batch_on_this_cp_rank( max_seqlen_kv=int(max_seqlen[0].item()), ) - cp_size = get_context_parallel_world_size() if cp_size is None else cp_size - cp_rank = get_context_parallel_rank() if cp_rank is None else cp_rank + cp_size = parallel_state.get_context_parallel_world_size() if cp_size is None else cp_size + cp_rank = parallel_state.get_context_parallel_rank() if cp_rank is None else cp_rank if cp_size > 1: # slice batch along sequence dimension for context parallelism assert tex is not None and is_te_min_version("1.10.0"), ( "Please update Transformer Engine to >= 1.10 to use "