Megatron dataset helper compilation fails inside `docker/Dockerfile` image (`No module named pybind11`), leading to NCCL init timeout on multi-node SFT
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
# Megatron dataset helper compilation fails inside `docker/Dockerfile` image (`No module named pybind11`), leading to NCCL init timeout on multi-node SFT
## Summary
When building the image with `docker/Dockerfile` and running a multi-node
Megatron SFT job (`examples/run_sft.py`), every `MegatronPolicyWorker` fails to
compile the Megatron-LM C++ dataset index builder because `pybind11` and
`python3-config` are missing from the runtime venv (`/opt/nemo_rl_venv`). Rank 0
dies while `compile_helpers()` runs, so the remaining ranks never receive the
NCCL unique id and every worker eventually crashes with a
`torch.distributed.DistBackendError` (TCPStore `wait timeout after 600000ms`).
## Environment
- Image: built from `docker/Dockerfile` (base `nvcr.io/nvidia/cuda-dl-base:25.05-cuda12.9-devel-ubuntu24.04`)
- Arch: aarch64 (GB200), Python 3.13.13, venv at `/opt/nemo_rl_venv`
- Backend: Megatron (`--extra mcore`), Megatron-Bridge + Megatron-LM submodules
- Model: NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16 (MoE, PP=4)
- Multi-node Ray cluster (4 nodes x 4 GPUs = 16 ranks)
- Entry point: `examples/run_sft.py`
## Steps to Reproduce
1. Build the image:
```bash
docker buildx build -f docker/Dockerfile --tag /nemo-rl:latest --push .
```
2. Launch a multi-node Megatron SFT run for a 30B MoE model (PP=4, 16 ranks).
3. Observe worker initialization.
## Actual Behavior
During `MegatronPolicyWorker.__init__` -> `setup_model_and_optimizer` ->
`initialize_megatron`, Megatron-LM tries to build its dataset helpers:
```
(MegatronPolicyWorker pid=19507) > compiling dataset index builder ...
(MegatronPolicyWorker pid=19507) make: Entering directory '/opt/nemo-rl/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/datasets'
(MegatronPolicyWorker pid=19507) make: python3-config: No such file or directory
(MegatronPolicyWorker pid=19507) /opt/nemo_rl_venv/bin/python3: No module named pybind11
(MegatronPolicyWorker pid=19507) helpers.cpp:12:10: fatal error: pybind11/pybind11.h: No such file or directory
(MegatronPolicyWorker pid=19507) 12 | #include
(MegatronPolicyWorker pid=19507) | ^~~~~~~~~~~~~~~~~~~~~
(MegatronPolicyWorker pid=19507) compilation terminated.
(MegatronPolicyWorker pid=19507) make: *** [Makefile:13: helpers_cpp] Error 1
(MegatronPolicyWorker pid=19507) ERROR:megatron.core.datasets.utils:_rank_utils.py:89: Failed to compile the C++ dataset helper functions
```
`compile_helpers()` in
`megatron/core/datasets/utils.py` calls `sys.exit(1)` on a non-zero return code,
so rank 0 (pid=19507) exits during setup. The other ranks then time out waiting
for the NCCL unique id from rank 0:
```
torch.distributed.DistBackendError: [9] is setting up NCCL communicator and
retrieving ncclUniqueId from [0] via c10d key-value store by key '0', but
store->get('0') got error: wait timeout after 600000ms, keys: /default_pg/0//cuda//0
```
which surfaces as:
```
ray.exceptions.ActorDiedError: The actor died because of an error raised in its
creation task, ray::lm_policy-2-1:MegatronPolicyWorker.__init__()
File ".../nemo_rl/models/megatron/setup.py", line 1177, in setup_model_and_optimizer
initialize_megatron(...)
File ".../megatron/bridge/training/initialize.py", line 154, in initialize_megatron
torch.distributed.barrier()
```
## Root Cause
The C++ compilation step needs two things that are not present in the built
image's runtime venv:
1. `pybind11` — not installed in `/opt/nemo_rl_venv`, so
`helpers.cpp` cannot find `pybind11/pybind11.h`.
2. `python3-config` — not on `PATH`, so the Makefile cannot resolve include
flags.
Because `compile_helpers()` hard-exits rank 0 on failure, the failure is not
local; it cascades into a cluster-wide NCCL timeout, making the true error easy
to miss in the noise.
## Expected Behavior
The Megatron dataset helpers should either be:
- pre-compiled during the image build, or
- compilable at runtime (i.e. `pybind11` and `python3-config` available in the
runtime venv/image),
so that Megatron SFT runs start without every rank timing out.
## Suggested Fixes (any of)
- Add `pybind11` to the mcore dependency group so it lands in `/opt/nemo_rl_venv`.
- Install the `python3-config` (`python3-dev`) tooling in the image, or ensure
the uv-managed python's `python-config` is on `PATH`.
- Pre-compile the Megatron dataset helpers at build time (e.g. run
`compile_helpers()` / `make -C .../megatron/core/datasets`) so no runtime
compilation is required.
## Notes
- This blocks multi-node Megatron SFT (`examples/run_sft.py`) on images built
from `docker/Dockerfile`.
- The 600s TCPStore timeout is a symptom, not the cause — the actionable failure
is the `pybind11` / `python3-config` missing during `compile_helpers()`.
Contributor guide
Assessment
This issue has not been assessed yet.