compat megatron-core main branch

[Jintao-Huang]

No description provided.

[gemini-code-assist]

Code Review

This pull request refactors imports in megatron_lm_utils.py by moving get_default_save_sharded_strategy and get_default_load_sharded_strategy to local scopes. It also updates _undo_attention_load_balancing in convert_utils.py to pass packed_seq_params. Feedback suggests dynamically checking the signature of _undo_attention_load_balancing before passing the third argument to prevent a TypeError on older versions of Megatron-Core.

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.

[gemini-code-assist]

In older versions of Megatron-Core, _undo_attention_load_balancing only accepts two arguments (tensor and cp_size). Passing a third argument (packed_seq_params) will raise a TypeError on those versions. To maintain backward compatibility with older Megatron-Core versions, we should dynamically check the function's signature before passing the third argument.

Suggested change

	mg_logits = _undo_attention_load_balancing(mg_logits, args.context_parallel_size,
	mg_inputs.get('packed_seq_params'))
	import inspect
	if 'packed_seq_params' in inspect.signature(_undo_attention_load_balancing).parameters:
	mg_logits = _undo_attention_load_balancing(mg_logits, args.context_parallel_size,
	mg_inputs.get('packed_seq_params'))
	else:
	mg_logits = _undo_attention_load_balancing(mg_logits, args.context_parallel_size)