[QUESTION] but the custom CUDA extension fused_weight_gradient_mlp_cuda module is not found

```python

[rank1]: Traceback (most recent call last):
[rank1]: File "/root/Megatron-LM/pretrain_gpt.py", line 233, in
[rank1]: pretrain(
[rank1]: File "/root/Megatron-LM/megatron/training/training.py", line 661, in pretrain
[rank1]: model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/root/Megatron-LM/megatron/training/training.py", line 1078, in setup_model_and_optimizer
[rank1]: model = get_model(model_provider_func, model_type)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/root/Megatron-LM/megatron/training/training.py", line 880, in get_model
[rank1]: model = build_model()
[rank1]: ^^^^^^^^^^^^^
[rank1]: File "/root/Megatron-LM/megatron/training/training.py", line 872, in build_model
[rank1]: model = model_provider_func(pre_process=pre_process, post_process=post_process)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/root/Megatron-LM/model_provider.py", line 69, in model_provider
[rank1]: return model_builder(args, pre_process, post_process, vp_stage)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/root/Megatron-LM/gpt_builders.py", line 76, in gpt_builder
[rank1]: model = GPTModel(
[rank1]: ^^^^^^^^^
[rank1]: File "/root/Megatron-LM/megatron/core/models/gpt/gpt_model.py", line 213, in __init__
[rank1]: self.output_layer = tensor_parallel.ColumnParallelLinear(
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/root/Megatron-LM/megatron/core/tensor_parallel/layers.py", line 914, in __init__
[rank1]: raise RuntimeError(
[rank1]: RuntimeError: ColumnParallelLinear was called with gradient_accumulation_fusion set to True but the custom CUDA extension fused_weight_gradient_mlp_cuda module is not found. To use gradient_accumulation_fusion you must install APEX with --cpp_ext and --cuda_ext. For example: pip install --global-option="--cpp_ext" --global-option="--cuda_ext ." Note that the extension requires CUDA>=11. Otherwise, you must turn off gradient accumulation fusion.
[rank0]: Traceback (most recent call last):
[rank0]: File "/root/Megatron-LM/pretrain_gpt.py", line 233, in
[rank0]: pretrain(
[rank0]: File "/root/Megatron-LM/megatron/training/training.py", line 661, in pretrain
[rank0]: model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/root/Megatron-LM/megatron/training/training.py", line 1078, in setup_model_and_optimizer
[rank0]: model = get_model(model_provider_func, model_type)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/root/Megatron-LM/megatron/training/training.py", line 880, in get_model
[rank0]: model = build_model()
[rank0]: ^^^^^^^^^^^^^
[rank0]: File "/root/Megatron-LM/megatron/training/training.py", line 872, in build_model
[rank0]: model = model_provider_func(pre_process=pre_process, post_process=post_process)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/root/Megatron-LM/model_provider.py", line 69, in model_provider
[rank0]: return model_builder(args, pre_process, post_process, vp_stage)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/root/Megatron-LM/gpt_builders.py", line 76, in gpt_builder
[rank0]: model = GPTModel(
[rank0]: ^^^^^^^^^
[rank0]: File "/root/Megatron-LM/megatron/core/models/gpt/gpt_model.py", line 213, in __init__
[rank0]: self.output_layer = tensor_parallel.ColumnParallelLinear(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/root/Megatron-LM/megatron/core/tensor_parallel/layers.py", line 914, in __init__
[rank0]: raise RuntimeError(
[rank0]: RuntimeError: ColumnParallelLinear was called with gradient_accumulation_fusion set to True but the custom CUDA extension fused_weight_gradient_mlp_cuda module is not found. To use gradient_accumulation_fusion you must install APEX with --cpp_ext and --cuda_ext. For example: pip install --global-option="--cpp_ext" --global-option="--cuda_ext ." Note that the extension requires CUDA>=11. Otherwise, you must turn off gradient accumulation fusion.
[rank0]:[W1023 04:57:37.057489583 ProcessGroupNCCL.cpp:1479] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
[rank1]:[W1023 04:57:37.088228054 ProcessGroupNCCL.cpp:1479] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
E1023 04:57:39.073000 14406 site-packages/torch/distributed/elastic/multiprocessing/api.py:874] failed (exitcode: 1) local_rank: 0 (pid: 14432) of binary: /root/miniconda3/envs/mega/bin/python3.11
Traceback (most recent call last):
File "/root/miniconda3/envs/mega/bin/torchrun", line 7, in
sys.exit(main())
^^^^^^
File "/root/miniconda3/envs/mega/lib/python3.11/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
return f(*args, **kwargs)
^^^^^^^^^^^^^^^^^^
File "/root/miniconda3/envs/mega/lib/python3.11/site-packages/torch/distributed/run.py", line 892, in main
run(args)
File "/root/miniconda3/envs/mega/lib/python3.11/site-packages/torch/distributed/run.py", line 883, in run
elastic_launch(
File "/root/miniconda3/envs/mega/lib/python3.11/site-packages/torch/distributed/launcher/api.py", line 139, in __call__
return launch_agent(self._config, self._entrypoint, list(args))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/miniconda3/envs/mega/lib/python3.11/site-packages/torch/distributed/launcher/api.py", line 270, in launch_agent
raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
============================================================
pretrain_gpt.py FAILED
------------------------------------------------------------
Failures:
[1]:
time : 2025-10-23_04:57:39
host : a9ba0af8b660
rank : 1 (local_rank: 1)
exitcode : 1 (pid: 14433)
error_file:
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
time : 2025-10-23_04:57:39
host : a9ba0af8b660
rank : 0 (local_rank: 0)
exitcode : 1 (pid: 14432)
error_file:
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================