NVIDIA / NVIDIA/Megatron-LM

Need TECudaGraphHelper support M4

Open
#3,135 2 comments 0 reactions 1 assignee Claimed by @yaoyu-33 View on GitHub
enhancement
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Is your feature request related to a problem? Please describe.**
Global pg groups, such as `_TENSOR_MODEL_PARALLEL_GROUP `, in `parallel_state.py` will be deprecated after M4.
So we need pass pg_collection while init TECudaGraphHelper, and use pg group in pg_collection later.

reproduce this issue, using MBridge with flag `dist.use_decentralized_pg=True` and `cuda_graph_impl: transformer_engine`.
for more informantion, refer to https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/2011.
```
0: [rank0]: Traceback (most recent call last):
0: [rank0]: File "/root/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/transformer/cuda_graphs.py", line 1690, in __init__
0: [rank0]: chunk_with_decoder = get_attr_wrapped_model(
0: [rank0]: ^^^^^^^^^^^^^^^^^^^^^^^
0: [rank0]: File "/root/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/utils.py", line 598, in get_attr_wrapped_model
0: [rank0]: raise RuntimeError(f"_get_attr_wrapped_model couldn't find attribute {attr}")
0: [rank0]: RuntimeError: _get_attr_wrapped_model couldn't find attribute decoder
0:
0: [rank0]: During handling of the above exception, another exception occurred:
0:
0: [rank0]: Traceback (most recent call last):
0: [rank0]: File "/lustre/fs1/portfolios/coreai/projects/coreai_devtech_all/users/shifangx/1.workspace/3.MBridge_
0: qwen3_vl_perf_m4/Scripts-MBridge/../Megatron-Bridge/scripts/performance/run_script.py", line 75, in
0: [rank0]: main()
0: [rank0]: File "/root/Scripts-MBridge/../Megatron-Bridge/scripts/performance/run_script.py", line 67, in main
0: [rank0]: pretrain(config=recipe, forward_step_func=forward_step_func)
0: [rank0]: File "/root/Megatron-Bridge/src/megatron/bridge/utils/decorators.py", line 39, in wrapper
0: [rank0]: return func(*args, **kwargs)
0: [rank0]: ^^^^^^^^^^^^^^^^^^^^^
0: [rank0]: File "/root/Megatron-Bridge/src/megatron/bridge/training/pretrain.py", line 92, in pretrain
0: [rank0]: _pretrain(state=state, forward_step_func=forward_step_func)
0: [rank0]: File "/lustre/fs1/po
0: rtfolios/coreai/projects/coreai_devtech_all/users/shifangx/1.workspace/3.MBridge_qwen3_vl_perf_m4/Megatron-Bridge/src/megatron/bridge/training/pretrain.py", line 134, in _pretrain
0: [rank0]: train(
0: [rank0]: File "/root/Megatron-Bridge/src/megatron/bridge/training/train.py", line 239, in train
0: [rank0]: cuda_graph_helper = TECudaGraphHelper(
0: [rank0]: ^^^^^^^^^^^^^^^^^^
0: [rank0]: File "/root/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/transformer/cuda_graphs.py", line 1695, in __init__
0: [rank0]: log_on_each_pipeline_stage(
0: [rank0]: File "/root/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/utils.py", line 869, in log_on_each_pipeline_stage
0: [rank0]:
0: tp_rank = parallel_state.get_tensor_model_parallel_rank()
0: [rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: [rank0]: File "/root/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/parallel_state.py", line 1577, in get_tensor_model_parallel_rank
0: [rank0]: return get_tensor_model_parallel_group().rank()
0: [rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: [rank0]: File "/root/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/parallel_state.py", line 1376, in get_tensor_model_parallel_group
0: [rank0]: _TENSOR_MODEL_PARALLEL_GROUP is not None
0: [rank0]: AssertionError: tensor model parallel group is not initialized
```

Tag the [@mcore-oncall](https://github.com/orgs/NVIDIA/teams/mcore-oncall)
to get oncall's attention to this issue.

**Describe the solution you'd like**
A clear and concise description of what you want to happen.

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**
Add any other context or screenshots about the feature request here.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.