fix(slime): pin numpy < 2 so Megatron init does not abort the train worker

[littlemex]

Purpose

Fourth fix in the epic to run the SLIME GRPO test case end-to-end on B300 / H200 (tracked by awslabs#1163). This merges into the epic branch epic/enable-slime-grpo-on-b300-h200, not into main directly; the epic is what goes to upstream in awslabs#1164.

With walls 1, 2 and 5 fixed, training now starts, the rollout HTTP server binds, and the train worker launches. The next wall is at Megatron initialization: the train worker aborts before the first step.

Changes

Megatron-LM refuses numpy 2.x. At init it asserts:

assert np.__version__.startswith("1."), "Megatron does not support numpy 2.x"

(the assert is in slime's slime/backends/megatron_utils/initialize.py L66, enforcing NVIDIA/Megatron-LM#1563). But the image ends up with numpy 2.x: sglang[all] — installed just before requirements.txt in slime.Dockerfile L160-L168 — pulls numpy 2.x transitively (numpy is a dependency of sglang, transformers, datasets, accelerate, scipy, and many others). So MegatronTrainRayActor dies during init with AssertionError: Megatron does not support numpy 2.x.

The fix is a one-line pin: add numpy<2 to requirements.txt.

Why requirements.txt specifically (placement matters): slime.Dockerfile installs in this order — sglang[all] (which introduces numpy 2.x), then requirements.txt, then the sgl-router wheel and slime itself (both effectively --no-deps). Putting the pin in requirements.txt means the requirements install downgrades numpy back to 1.x after sglang pulled 2.x, and nothing after it reintroduces 2.x. Placing the pin earlier (before sglang) would be undone; placing it in a separate later RUN would also work but adds a layer for no benefit.

I verified the downgrade is clean: pip install "numpy<2" resolves to numpy 1.26.4 with no dependency conflicts that block it (a pip install --dry-run reports only Would install numpy-1.26.4). tifffile prints an advisory that it prefers numpy>=2.1, but it is a transitive image dependency not used on the SLIME/Megatron/SGLang GRPO path, and both torch and megatron.core import cleanly under numpy 1.26.4.

Implementation

• requirements.txt (adds numpy<2 with a comment) — https://github.com/littlemex/awsome-distributed-training/blob/b8ec49f/3.test_cases/pytorch/slime/requirements.txt
• Commit (pin) — a3a488d
• Commit (comment hardening) — b8ec49f

Alternatives considered (and why this is the most elegant)

The constraint is fixed and external: Megatron-LM does not support numpy 2.x (NVIDIA/Megatron-LM#1563), and this is a downstream test case that cannot change Megatron or SGLang. So the only decision is where and how to honor numpy<2. I considered each option:

1. Pin numpy<2 in requirements.txt (chosen). One line, at the last dependency step that touches numpy, so it downgrades the 2.x that sglang[all] pulled and nothing after re-introduces it. It is exactly what upstream SLIME's own docker/Dockerfile does (pip install "numpy<2"), so it mirrors the sanctioned approach and stays close to upstream.

2. A separate trailing RUN pip install "numpy<2" after the SLIME install. Functionally equivalent and slightly more order-robust (survives a future reorder or a new numpy-pulling install), but it adds a Docker layer and a second place that states the same constraint. Rejected as the default because it is strictly more machinery for no current benefit; the ordering is already correct and documented in the pin's comment. (If the team later wants order-independence for free, this is the drop-in upgrade — keeping <2, not ==.)

3. Pin an exact numpy==1.26.4. Rejected: it would diverge from upstream SLIME's open-ended <2 and needlessly freeze the image on one 1.x release. <2 lets pip pick the best 1.x the rest of the set resolves to (currently 1.26.4) and is what Megatron actually requires (startswith("1.")).

4. Patch/relax the Megatron assert, or install numpy 2.x-compatible Megatron. Rejected: that is changing the upstream training backend's supported configuration from a test case — out of scope, unverifiable here, and against the "install upstream unmodified where possible" principle. The assert is correct; the image should satisfy it, not defeat it.

5. Force SGLang onto numpy 1.x from the start / avoid pulling 2.x. Not actionable: numpy 2.x arrives transitively through many packages (sglang, transformers, datasets, accelerate, scipy, ...). Constraining every producer is fragile; correcting once, after they have all resolved, is the clean seam. sglang 0.5.12.post1 itself declares plain numpy with no lower bound, so pinning <2 does not fight it.

Option 1 is the most elegant: one line, upstream-identical, at the correct layer, with no new Docker layer and no second source of truth.

Portability (B300 / other GPUs)

This pin is GPU-generation- and CUDA-independent, so it does not risk B300. numpy is a pure-CPU package: numpy 1.26.4 ships no native extension that links a CUDA runtime (verified: zero of its .so files reference libcudart/libcuda), and Megatron's guard is a pure version string check (np.__version__.startswith("1.")), not a GPU or CUDA-major concern. The same numpy<2 is required and equally harmless on H100, H200 and B300 — unlike the wall-5 .so selection, there is no per-generation behavior here to break.

Test Plan

This test case runs on Amazon EKS, so verification is done with kubectl exec into a Ray worker pod. The check below is copy-paste-safe (no <placeholder> tokens, no shell comment lines inside a command, no Ray job id required) and shows both the bug and the fix in one run. Set the namespace once:

export NAMESPACE=<your-namespace>

The check reports the effective numpy version and whether Megatron's init assert (np.__version__.startswith("1.")) would pass, and confirms torch/megatron import cleanly:

WPOD=$(kubectl -n "$NAMESPACE" get pod -l ray.io/group=gpu-workers -o jsonpath='{.items[0].metadata.name}')
kubectl -n "$NAMESPACE" exec -i "$WPOD" -c ray-worker -- python3 -W ignore - <<'PY'
import numpy
ok = numpy.__version__.startswith("1.")
print("numpy version =", numpy.__version__)
print("Megatron init assert np.__version__.startswith('1.') =>", "PASS" if ok else "FAIL (Megatron does not support numpy 2.x)")
import torch, megatron.core
print("import torch/megatron.core => OK")
PY

Expected with the pin (numpy resolved to 1.x, so the assert passes):

numpy version = 1.26.4
Megatron init assert np.__version__.startswith('1.') => PASS
import torch/megatron.core => OK

Without the pin the same check reports numpy version = 2.1.0 and FAIL, which is exactly the AssertionError: Megatron does not support numpy 2.x the train worker hits at init.

Test Results

Verified on hardware (H200, 2× p5en.48xlarge, CUDA 13.0):

• The image builds cleanly from this slime.Dockerfile: during the build the requirements.txt install step uninstalls numpy 2.1.0 and installs numpy 1.26.4 (the pin takes effect after sglang[all] pulled 2.x), and the build completes and pushes.
• A container started from that freshly built image (no manual change) reports numpy 1.26.4; Megatron's init assert passes; import torch and import megatron.core succeed.

	before	after numpy<2
numpy version	2.1.0	1.26.4
import torch / import megatron.core	OK	OK
Megatron init	AssertionError: Megatron does not support numpy 2.x	passes
train worker	aborts at init	reaches Megatron init and proceeds into the rollout/train/weight-sync cycle

Blast radius / risk

• Low. One added constraint in requirements.txt. numpy 1.26.4 is the last 1.x release and is the version Megatron expects.
• pip resolves the rest of the pinned set against numpy 1.x without conflict (verified via dry-run and by importing torch/megatron/sglang).
• No CLI/API or recipe changes.

Checklist

• I am working against the epic branch (not main directly).
• The change is minimal (one pinned constraint, with a comment explaining the ordering).
• External dependencies are pinned (numpy<2 resolves to 1.26.4).
• Verified on hardware (H200 / CUDA 13): numpy is 1.x, Megatron init passes, the train worker proceeds.