LLama3.1-405B Seeing Functionality issue on both GB300 and GB200.

[rsalagame-nvidia]

LLama3.1-405B Seeing Functionality issue on both GB300 and GB200. These functionality issues will be addressed according to Nemo team in 26.02 nemo container.

GB300 issue:

[rank3]: torch.AcceleratorError: CUDA error: operation failed due to a previous error during capture
[rank3]: Search for `cudaErrorStreamCaptureInvalidated' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
[rank3]: CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
[rank3]: For debugging consider passing CUDA_LAUNCH_BLOCKING=1
[rank3]: Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions. 

GB200:

[rank112]: torch.distributed.DistBackendError: NCCL error in: /opt/pytorch/pytorch/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:3694, unhandled cuda error (run with NCCL_DEBUG=INFO for details), NCCL version 2.28.3
[rank112]: ncclUnhandledCudaError: Call to CUDA function failed.
[rank112]: Last error:
[rank112]: Cuda failure 2 'out of memory'
[rank125]: Traceback (most recent call last):
[rank125]:  File "/lustre/fsw/infra_rd_gsw/rsalagame/llama3.1-405B-70B-8B-GB300/test-env-25.11.01.rc3/workloads/pretrain_llama3.1/Megatron-Bridge/scripts/performance/run_script.py", line 60, in <module>

[rsalagame-nvidia]

log-infra_rd_gsw-gsw.pretrain_llama31_405b_nvfp4_gpus128_tp4_pp8_cp2_vp4_ep1_mbs1_gbs64_564068_0_GB300.txt

[rsalagame-nvidia]

log-infra_rd_gsw-gsw.pretrain_llama31_405b_nvfp4_gpus128_tp4_pp8_cp2_vp4_ep1_mbs1_gbs64_1364660_0_GB200.txt

[malay-nagda]

GB300 issue is a copy of this issue- #1582. We will track that issue in #1582 and not here.

GB200 OOM is resolved with this PR- NVIDIA/TransformerEngine#2535

[rsalagame-nvidia]

Will wait for 26.02.rcx containers to be released (rc3 will have this fix) and test with those with 405B NVFP4 datatype to confirm successful run

[rsalagame-nvidia]

@nv-mollys @malay-nagda

When testing with 26.02.rc1 nemo container, still seeing issues on GB300, GB200, and H100 chips.

when testing with variant 2:
GB300:
FP8_CS: Seeing issue where training stops at the 30 iteration mark (log attached: log-infra_rd_gsw-gsw.llama31_405b_pretrain_128gpu_gb300_fp8_cs_724174_0.txt)
nvfp4: Seeing torch.distributed.DistBackendError issue (log attached: log-infra_rd_gsw-gsw.llama31_405b_pretrain_128gpu_gb300_nvfp4_724207_0.txt)

log-infra_rd_gsw-gsw.llama31_405b_pretrain_128gpu_gb300_nvfp4_724207_0.txt

log-infra_rd_gsw-gsw.llama31_405b_pretrain_128gpu_gb300_fp8_cs_724174_0.txt

)

GB200:
FP8_CS: torch.OutOfMemoryError: CUDA out of memory. (log attached: log-infra_rd_gsw-gsw.llama31_405b_pretrain_128gpu_gb200_fp8_cs_1504564_0.txt)
nvfp4: torch.OutOfMemoryError: CUDA out of memory. (log attached: log-infra_rd_gsw-gsw.llama31_405b_pretrain_128gpu_gb200_nvfp4_1504597_0.txt)

log-infra_rd_gsw-gsw.llama31_405b_pretrain_128gpu_gb200_fp8_cs_1504564_0.txt

)

H100
FP8_CS: gets cancelled after 30 iterations of running. Seeing a similar log to the above logs with same error.

log-infra_rd_gsw-gsw.llama31_405b_pretrain_128gpu_gb200_nvfp4_1504597_0.txt

)

@malay-nagda should I be using variant 1 here for llama3.1-405B? Would using this variant alleviate these issues as well as the ones above?