B200/GB200 on NGC pytorch

[terrykong]

After #897, we need to ensure mcore is new enough to run on B200/GB200. This issue tracks this effort

[terrykong]

There is a change in-flight for megatron bridge support that bumps mcore #905 . We should check if that's a new enough mcore that works both for the NGC pytorch and default dockerfile

[guyueh1]

I have pipecleaned the GB200 GRPO run; currently the workflow is like this:

1. Build container from arm NGC pytorch base using docker/Dockerfile.ngc_pytorch
2. Start container interactively
3. If your code and container are built from different RL commits, you may need to re-sync the uv venv via

uv lock
UV_PROJECT_ENVIRONMENT=/opt/nemo_rl_venv uv sync --link-mode symlink --locked --inexact --extra vllm --extra mcore --extra automodel --all-groups --no-install-project --no-install-package torch --no-install-package torchvision --no-install-package triton --no-install-package nvidia-cublas-cu12 --no-install-package nvidia-cuda-cupti-cu12 --no-install-package nvidia-cuda-nvrtc-cu12 --no-install-package nvidia-cuda-runtime-cu12 --no-install-package nvidia-cudnn-cu12 --no-install-package nvidia-cufft-cu12 --no-install-package nvidia-cufile-cu12 --no-install-package nvidia-curand-cu12 --no-install-package nvidia-cusolver-cu12 --no-install-package nvidia-cusparse-cu12 --no-install-package nvidia-cusparselt-cu12 --no-install-package nvidia-nccl-cu12 --no-install-package vllm --no-install-package flash-attn --no-install-package transformer-engine --no-install-package transformer-engine-cu12 --no-install-package transformer-engine-torch --no-install-package numpy

Note that this only syncs part of the dependencies, as all packages specified in --no-install-package is installed in system python environment; if you want to update those packages you will need to use pip install.

4. Run test with python ./examples/run_grpo_math.py --config examples/configs/grpo_math_8B_megatron.yaml cluster.gpus_per_node=4This also depends on build: Use no-build-isolation to install deep_gemm to fix arm install #970.

[guyueh1]

Re-assign to @chtruong814 because the next step is track the merge of #905 and releasing ngc pytorch based container