Missing import fix (#3241)

Author: parthmannan

megatron/core/utils.py

@@ -2082,8 +2082,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 "