Need TECudaGraphHelper support M4
- 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
Assessment
This issue has not been assessed yet.