From 9667d55ef1f7006faf92008a73134b7bc8436bab Mon Sep 17 00:00:00 2001 From: botaohu001 Date: Tue, 21 Jul 2026 07:23:47 +0000 Subject: [PATCH 1/2] test(odc): add unit/integration tests for ODC (PR #864) Add CPU-only tests that exercise the ODC feature wiring and gate the GPU/rocSHMEM-turbo parts behind clean skips. tests/runner/test_odc_run_launcher.sh (registered in run_all_tests.sh): Bash unit test for run_odc.sh launch/config wiring -- mori vs rocshmem backend env, the pad|nopad no-op arg, PYTHONPATH so `import odc` resolves, PRIMUS_TURBO_PATH prepend, extra KEY=VAL export, EXP/exp-name/MASTER_PORT plumbing, and that ODC feature switches are now config items, not env. tests/trainer/test_odc_megatron_trainer.py (+ .yaml fixture): - TestOdcRuntimeConfig: stdlib-only unit tests for the config-driven ODC runtime config (defaults, set_config override/None/unknown handling, the gda_pipe -> PRIMUS_TURBO_ODC_GDA_PIPE bridge, odc package re-exports). - TestOdcPatchWiring: before_train patch registration + enable_odc / use_torch_fsdp2 / enable_odc_lb_mini gating, LB-Mini aligned-vs-decoupled mode, the #856 device_id patch skip under ODC, and the trainer-config -> runtime-config bridge. Skips where torch is unavailable. - TestOdcConfigFixture: asserts the fixture declares the ODC FSDP2 path. - TestOdcMegatronTrainerE2E: FSDP2 pretrain smoke test that skips cleanly without ROCm GPU + Primus-Turbo rocSHMEM ops. Co-authored-by: Cursor --- tests/runner/run_all_tests.sh | 1 + tests/runner/test_odc_run_launcher.sh | 395 +++++++++++++++++++ tests/trainer/test_odc_megatron_trainer.py | 359 +++++++++++++++++ tests/trainer/test_odc_megatron_trainer.yaml | 108 +++++ 4 files changed, 863 insertions(+) create mode 100755 tests/runner/test_odc_run_launcher.sh create mode 100644 tests/trainer/test_odc_megatron_trainer.py create mode 100644 tests/trainer/test_odc_megatron_trainer.yaml diff --git a/tests/runner/run_all_tests.sh b/tests/runner/run_all_tests.sh index ee8a0737c..1a49a6927 100755 --- a/tests/runner/run_all_tests.sh +++ b/tests/runner/run_all_tests.sh @@ -40,6 +40,7 @@ TEST_SCRIPTS=( "$SCRIPT_DIR/test_primus_cli_slurm.sh" "$SCRIPT_DIR/test_primus_cli_container.sh" "$SCRIPT_DIR/test_primus_cli_direct.sh" + "$SCRIPT_DIR/test_odc_run_launcher.sh" ) # Run each test suite diff --git a/tests/runner/test_odc_run_launcher.sh b/tests/runner/test_odc_run_launcher.sh new file mode 100755 index 000000000..f42057179 --- /dev/null +++ b/tests/runner/test_odc_run_launcher.sh @@ -0,0 +1,395 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +# Unit tests for the ODC LB-mini launcher (primus/core/odc/rocshmem_runtime/ +# scripts/run_odc.sh). +# +# These are CPU-only / no-GPU tests: they exercise run_odc.sh's argument parsing +# and environment / PYTHONPATH wiring WITHOUT actually launching training. The +# real examples/run_pretrain.sh is stubbed by pointing PRIMUS_ROOT at a fake +# tree whose examples/run_pretrain.sh just dumps the environment run_odc.sh +# exported, so we can assert on it deterministically. +# +# What is asserted (real ODC launch/config wiring, not trivial always-true): +# * backend selection: `mori` (default) vs `rocshmem` sets ODC_P2P_BACKEND and +# the matching rocSHMEM / MORI infra env (heap size, bootstrap ifname, +# Triton cache policy). +# * the pad|nopad positional arg is a NO-OP (retained for backwards-compatible +# invocation): it does not change the backend or introduce any aligned-vs- +# decoupled env; only the echoed PAD token differs. +# * PYTHONPATH is wired so `import odc` resolves (odc_early shim + primus/core), +# and PRIMUS_TURBO_PATH is prepended when provided. +# * extra KEY=VAL trailing args are exported verbatim. +# * EXP / PRIMUS_EXP_NAME / MASTER_PORT plumbing. +# * ODC feature switches (enable_odc, odc_phase, enable_odc_lb_mini, ...) are +# NOT exported as env by the launcher -- they are now CONFIG items read from +# the EXP yaml. This guards the env->config migration from regressing. + +# Get project root (tests/runner/ -> tests/ -> repo root) +PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +RUN_ODC="$PROJECT_ROOT/primus/core/odc/rocshmem_runtime/scripts/run_odc.sh" + +# Test counters +TESTS_RUN=0 +TESTS_PASSED=0 +TESTS_FAILED=0 + +# Colors for output +GREEN='\033[0;32m' +RED='\033[0;31m' +YELLOW='\033[1;33m' +NC='\033[0m' # No Color + +# Scratch state (populated by run_odc_capture) +FAKE_ROOT="" +CAPTURE="" +STDOUT_LOG="" + +# --------------------------------------------------------------------------- +# Test helpers (mirrors tests/runner/test_primus_cli.sh conventions) +# --------------------------------------------------------------------------- +assert_pass() { + local test_name="$1" + ((TESTS_RUN++)) + echo -e "${GREEN}✓${NC} $test_name" + ((TESTS_PASSED++)) +} + +assert_fail() { + local test_name="$1" + local reason="${2:-}" + ((TESTS_RUN++)) + echo -e "${RED}✗${NC} $test_name" + if [[ -n "$reason" ]]; then + echo " Reason: $reason" + fi + ((TESTS_FAILED++)) +} + +# assert_line FILE EXACT_LINE TEST_NAME -- assert FILE contains an exact line. +assert_line() { + local file="$1" + local line="$2" + local test_name="$3" + ((TESTS_RUN++)) + if grep -qxF -- "$line" "$file"; then + echo -e "${GREEN}✓${NC} $test_name" + ((TESTS_PASSED++)) + else + echo -e "${RED}✗${NC} $test_name" + echo " Expected exact line: $line" + echo " In file: $file" + ((TESTS_FAILED++)) + fi +} + +# assert_contains HAYSTACK NEEDLE TEST_NAME +assert_contains() { + local haystack="$1" + local needle="$2" + local test_name="$3" + ((TESTS_RUN++)) + if echo "$haystack" | grep -qF -- "$needle"; then + echo -e "${GREEN}✓${NC} $test_name" + ((TESTS_PASSED++)) + else + echo -e "${RED}✗${NC} $test_name" + echo " Expected to contain: $needle" + ((TESTS_FAILED++)) + fi +} + +print_section() { + echo "" + echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}" + echo -e "${YELLOW}$1${NC}" + echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}" +} + +cap_value() { + # Echo the value of KEY from the capture file (or empty if absent). + local key="$1" + grep -m1 "^${key}=" "$CAPTURE" 2>/dev/null | cut -d= -f2- +} + +# --------------------------------------------------------------------------- +# Fixture: a fake PRIMUS_ROOT whose examples/run_pretrain.sh dumps the env that +# run_odc.sh exported, instead of launching real training. +# --------------------------------------------------------------------------- +setup_fake_root() { + FAKE_ROOT="$(mktemp -d)" + # Export CAPTURE once here (not inside the run subshell) so the stubbed + # run_pretrain.sh inherits it, while shellcheck does not flag a subshell-local + # modification (SC2030/SC2031). + export CAPTURE="$FAKE_ROOT/capture.env" + STDOUT_LOG="$FAKE_ROOT/stdout.log" + mkdir -p "$FAKE_ROOT/examples" + cat > "$FAKE_ROOT/examples/run_pretrain.sh" << 'STUB' +#!/bin/bash +# Stub standing in for the real trainer launch. Dump the env run_odc.sh set so +# the test can assert on it, then attempt `import odc` to prove PYTHONPATH wiring. +{ + for _k in ODC_P2P_BACKEND EXP PRIMUS_EXP_NAME MASTER_PORT TRITON_CACHE_DIR \ + MORI_SHMEM_HEAP_SIZE ROCSHMEM_HEAP_SIZE ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME \ + PYTHONPATH FUSED_LINEAR_CE GLOO_SOCKET_IFNAME NCCL_IB_DISABLE FOO \ + enable_odc ODC_ENABLE odc_phase enable_odc_lb_mini; do + if [[ -n "${!_k+x}" ]]; then + echo "${_k}=${!_k}" + else + echo "${_k}=" + fi + done + if command -v python >/dev/null 2>&1; then + if python -c "import odc" >/dev/null 2>&1; then + echo "IMPORT_ODC=ok" + else + echo "IMPORT_ODC=fail" + fi + elif command -v python3 >/dev/null 2>&1; then + if python3 -c "import odc" >/dev/null 2>&1; then + echo "IMPORT_ODC=ok" + else + echo "IMPORT_ODC=fail" + fi + else + echo "IMPORT_ODC=no-python" + fi +} > "$CAPTURE" +STUB +} + +teardown_fake_root() { + [[ -n "$FAKE_ROOT" && -d "$FAKE_ROOT" ]] && rm -rf "$FAKE_ROOT" +} + +# run_odc_capture [extra KEY=VAL ...] +# Optional overrides via env before the call: +# ODC_TURBO_PATH_OVERRIDE -> PRIMUS_TURBO_PATH +# ODC_MASTER_PORT_OVERRIDE -> MASTER_PORT +run_odc_capture() { + rm -f "$CAPTURE" "$STDOUT_LOG" + ( + # Start from a clean ODC-related env so absence assertions are meaningful. + unset PRIMUS_TURBO_PATH ODC_P2P_BACKEND ROCSHMEM_HEAP_SIZE TRITON_CACHE_DIR \ + MORI_SHMEM_HEAP_SIZE ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME MASTER_PORT \ + enable_odc ODC_ENABLE odc_phase enable_odc_lb_mini FOO + [[ -n "${ODC_TURBO_PATH_OVERRIDE:-}" ]] && export PRIMUS_TURBO_PATH="$ODC_TURBO_PATH_OVERRIDE" + [[ -n "${ODC_MASTER_PORT_OVERRIDE:-}" ]] && export MASTER_PORT="$ODC_MASTER_PORT_OVERRIDE" + export PRIMUS_ROOT="$FAKE_ROOT" + export PRIMUS_PACK_CACHE_DIR="$FAKE_ROOT/pack" + export TRAIN_LOG_DIR="$FAKE_ROOT/logs" + bash "$RUN_ODC" "$@" + ) > "$STDOUT_LOG" 2>&1 +} + +# ============================================================================ +# Test 1: mori backend (default) env wiring +# ============================================================================ +test_mori_backend_env() { + print_section "Test 1: MORI backend (default) env wiring" + run_odc_capture mori pad some/exp.yaml myexp + + assert_line "$CAPTURE" "ODC_P2P_BACKEND=mori" "mori sets ODC_P2P_BACKEND=mori" + assert_line "$CAPTURE" "MORI_SHMEM_HEAP_SIZE=8G" "mori sets MORI symmetric heap 8G" + assert_line "$CAPTURE" "TRITON_CACHE_DIR=/tmp/tcache_mori" "mori uses stable /tmp/tcache_mori" + assert_line "$CAPTURE" "ROCSHMEM_HEAP_SIZE=" "mori does NOT set ROCSHMEM_HEAP_SIZE" + assert_line "$CAPTURE" "ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME=" \ + "mori does NOT set ROCSHMEM bootstrap ifname" + assert_contains "$(cat "$STDOUT_LOG")" "P2P=mori" "launcher banner reports P2P=mori" +} + +# ============================================================================ +# Test 2: rocshmem backend env wiring +# ============================================================================ +test_rocshmem_backend_env() { + print_section "Test 2: rocSHMEM backend env wiring" + run_odc_capture rocshmem nopad other/exp.yaml exp2 + + assert_line "$CAPTURE" "ODC_P2P_BACKEND=rocshmem" "rocshmem sets ODC_P2P_BACKEND=rocshmem" + # Decimal-only heap parser: 8 GiB in raw bytes, NOT a K/M/G suffix. + assert_line "$CAPTURE" "ROCSHMEM_HEAP_SIZE=8589934592" "rocshmem sets 8 GiB heap in raw bytes" + assert_line "$CAPTURE" "ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME=lo" "rocshmem sets bootstrap ifname=lo" + + local triton + triton="$(cap_value TRITON_CACHE_DIR)" + if [[ "$triton" == /tmp/tcache_rocshmem_* ]]; then + assert_pass "rocshmem uses a fresh per-run Triton cache (/tmp/tcache_rocshmem_*)" + else + assert_fail "rocshmem uses a fresh per-run Triton cache" "got TRITON_CACHE_DIR=$triton" + fi + assert_contains "$(cat "$STDOUT_LOG")" "P2P=rocshmem" "launcher banner reports P2P=rocshmem" +} + +# ============================================================================ +# Test 3: pad|nopad positional arg is a no-op +# ============================================================================ +test_pad_arg_is_noop() { + print_section "Test 3: pad|nopad positional arg is a no-op" + + run_odc_capture mori pad e.yaml n + local pad_backend pad_triton pad_mori + pad_backend="$(cap_value ODC_P2P_BACKEND)" + pad_triton="$(cap_value TRITON_CACHE_DIR)" + pad_mori="$(cap_value MORI_SHMEM_HEAP_SIZE)" + local pad_banner; pad_banner="$(grep -o 'PAD=[^ ]*' "$STDOUT_LOG" | head -1)" + + run_odc_capture mori nopad e.yaml n + local nopad_backend nopad_triton nopad_mori + nopad_backend="$(cap_value ODC_P2P_BACKEND)" + nopad_triton="$(cap_value TRITON_CACHE_DIR)" + nopad_mori="$(cap_value MORI_SHMEM_HEAP_SIZE)" + local nopad_banner; nopad_banner="$(grep -o 'PAD=[^ ]*' "$STDOUT_LOG" | head -1)" + + if [[ "$pad_backend" == "$nopad_backend" && "$pad_triton" == "$nopad_triton" \ + && "$pad_mori" == "$nopad_mori" ]]; then + assert_pass "pad vs nopad produce identical backend/env wiring" + else + assert_fail "pad vs nopad produce identical backend/env wiring" \ + "pad=($pad_backend,$pad_triton,$pad_mori) nopad=($nopad_backend,$nopad_triton,$nopad_mori)" + fi + + if [[ "$pad_banner" == "PAD=pad" && "$nopad_banner" == "PAD=nopad" ]]; then + assert_pass "only the echoed PAD token reflects the arg" + else + assert_fail "only the echoed PAD token reflects the arg" \ + "pad_banner=$pad_banner nopad_banner=$nopad_banner" + fi + + # No aligned/decoupled A/B env should be introduced by the pad arg. + assert_line "$CAPTURE" "enable_odc_lb_mini=" "pad arg does not toggle any LB-mini env" +} + +# ============================================================================ +# Test 4: PYTHONPATH wiring for `import odc` +# ============================================================================ +test_pythonpath_wiring() { + print_section "Test 4: PYTHONPATH wiring for import odc" + run_odc_capture mori pad e.yaml n + + local pp; pp="$(cap_value PYTHONPATH)" + assert_contains "$pp" "primus/core/odc/odc_early" "PYTHONPATH includes the odc_early load-order shim" + if [[ "$pp" == *"/primus/core" ]]; then + assert_pass "PYTHONPATH ends with primus/core (parent of the odc package)" + else + assert_fail "PYTHONPATH ends with primus/core" "PYTHONPATH=$pp" + fi + + local imp; imp="$(cap_value IMPORT_ODC)" + if [[ "$imp" == "ok" ]]; then + assert_pass "import odc resolves via the wired PYTHONPATH" + elif [[ "$imp" == "no-python" ]]; then + assert_pass "import odc check skipped (no python interpreter on host)" + else + assert_fail "import odc resolves via the wired PYTHONPATH" "IMPORT_ODC=$imp" + fi +} + +# ============================================================================ +# Test 5: PRIMUS_TURBO_PATH is prepended to PYTHONPATH +# ============================================================================ +test_turbo_path_prepend() { + print_section "Test 5: PRIMUS_TURBO_PATH prepend" + ODC_TURBO_PATH_OVERRIDE="/fake/turbo/build" run_odc_capture mori pad e.yaml n + + local pp; pp="$(cap_value PYTHONPATH)" + if [[ "$pp" == /fake/turbo/build:* ]]; then + assert_pass "PRIMUS_TURBO_PATH is prepended to PYTHONPATH" + else + assert_fail "PRIMUS_TURBO_PATH is prepended to PYTHONPATH" "PYTHONPATH=$pp" + fi + assert_contains "$(cat "$STDOUT_LOG")" "TURBO_PATH=/fake/turbo/build" \ + "launcher banner reports the turbo build path" +} + +# ============================================================================ +# Test 6: extra KEY=VAL trailing args are exported verbatim +# ============================================================================ +test_extra_kv_exported() { + print_section "Test 6: extra KEY=VAL args exported verbatim" + run_odc_capture mori pad e.yaml n FOO=bar + + assert_line "$CAPTURE" "FOO=bar" "trailing FOO=bar is exported to the trainer env" +} + +# ============================================================================ +# Test 7: EXP / PRIMUS_EXP_NAME / MASTER_PORT plumbing +# ============================================================================ +test_positional_and_port_plumbing() { + print_section "Test 7: EXP / exp-name / MASTER_PORT plumbing" + + run_odc_capture mori pad path/to/exp.yaml my-run-name + assert_line "$CAPTURE" "EXP=path/to/exp.yaml" "EXP is set from the 3rd positional arg" + assert_line "$CAPTURE" "PRIMUS_EXP_NAME=my-run-name" "PRIMUS_EXP_NAME is set from the 4th positional arg" + assert_line "$CAPTURE" "MASTER_PORT=29600" "MASTER_PORT defaults to 29600" + + ODC_MASTER_PORT_OVERRIDE="29777" run_odc_capture mori pad e.yaml n + assert_line "$CAPTURE" "MASTER_PORT=29777" "MASTER_PORT is overridable via env" +} + +# ============================================================================ +# Test 8: ODC feature switches are config, not env +# ============================================================================ +test_switches_are_config_not_env() { + print_section "Test 8: ODC feature switches are config, not env" + run_odc_capture rocshmem pad e.yaml n + + # The launcher must NOT export any of the ODC feature switches: post-#864 + # they live in the EXP yaml config, read by the before_train patches. + assert_line "$CAPTURE" "enable_odc=" "launcher does not export enable_odc" + assert_line "$CAPTURE" "ODC_ENABLE=" "launcher does not export the legacy ODC_ENABLE env" + assert_line "$CAPTURE" "odc_phase=" "launcher does not export odc_phase" + assert_line "$CAPTURE" "enable_odc_lb_mini=" "launcher does not export enable_odc_lb_mini" +} + +# ============================================================================ +# Run all tests +# ============================================================================ +main() { + echo "╔══════════════════════════════════════════════════════════════╗" + echo "║ Unit Tests for the ODC LB-mini launcher (run_odc.sh) ║" + echo "╚══════════════════════════════════════════════════════════════╝" + + if [[ ! -f "$RUN_ODC" ]]; then + echo -e "${RED}✗ run_odc.sh not found at $RUN_ODC${NC}" + return 1 + fi + + setup_fake_root + trap teardown_fake_root EXIT + + test_mori_backend_env + test_rocshmem_backend_env + test_pad_arg_is_noop + test_pythonpath_wiring + test_turbo_path_prepend + test_extra_kv_exported + test_positional_and_port_plumbing + test_switches_are_config_not_env + + echo "" + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "Test Summary:" + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo " Total: $TESTS_RUN" + echo -e " Passed: ${GREEN}$TESTS_PASSED${NC}" + if [[ $TESTS_FAILED -gt 0 ]]; then + echo -e " Failed: ${RED}$TESTS_FAILED${NC}" + else + echo " Failed: 0" + fi + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + + if [[ $TESTS_FAILED -eq 0 ]]; then + echo -e "${GREEN}✓ All tests passed!${NC}" + return 0 + else + echo -e "${RED}✗ Some tests failed${NC}" + return 1 + fi +} + +main diff --git a/tests/trainer/test_odc_megatron_trainer.py b/tests/trainer/test_odc_megatron_trainer.py new file mode 100644 index 000000000..f589b1b22 --- /dev/null +++ b/tests/trainer/test_odc_megatron_trainer.py @@ -0,0 +1,359 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""ODC (On-Demand Communication) integration tests. + +Covers the feature merged in PR #864 (``feat: odc adapt``): + * ``primus/core/odc/runtime_config.py`` -- the config-driven runtime knobs the + ODC library reads instead of the former ``ODC_*`` env vars. + * the before_train patch wiring (``odc_torch_fsdp2_patches``, + ``odc_lb_mini_patches``, ``distributed_init_patches``) -- registration and + the ``enable_odc`` / ``use_torch_fsdp2`` / ``enable_odc_lb_mini`` gating. + * the ODC training-config fixture (``test_odc_megatron_trainer.yaml``). + * an end-to-end training smoke test on the ODC FSDP2 path. + +Hardware handling (matches the neighbor trainer tests): + * The pure config layer (``TestOdcRuntimeConfig``) is stdlib-only and runs + CPU-only, always -- it needs neither torch nor a GPU. + * The patch-wiring layer (``TestOdcPatchWiring``) needs ``torch`` importable + (NOT a GPU -- the patch modules import torch at module scope but only touch + CUDA at runtime); it skips cleanly where torch is absent. + * The end-to-end run (``TestOdcMegatronTrainerE2E``) needs ROCm GPUs + a + Primus-Turbo build with the rocSHMEM ops, so it skips with a clear reason + when the hardware / turbo stack is unavailable -- the same GPU-gating idea as + ``TestProjectionSimulate._require_2_gpus`` in test_megatron_trainer.py. +""" + +import dataclasses +import importlib +import importlib.util +import os +import sys +import unittest +from types import SimpleNamespace + +from tests.utils import PrimusUT, run_training_script + +# Repo root: tests/trainer/ -> repo root. +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +# The ``odc`` package lives at primus/core/odc/, imported as top-level ``odc`` +# with primus/core on the path (exactly what run_odc.sh wires into PYTHONPATH). +_PRIMUS_CORE = os.path.join(_REPO_ROOT, "primus", "core") +if _PRIMUS_CORE not in sys.path: + sys.path.insert(0, _PRIMUS_CORE) + +_HAS_TORCH = importlib.util.find_spec("torch") is not None + +_ODC_FIXTURE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_odc_megatron_trainer.yaml") + + +def _make_ctx(**params): + """Build a minimal PatchContext whose get_args(ctx) returns the given params. + + Mirrors the real call site: get_args reads ctx.extra['module_config'].params. + """ + from primus.core.patches import PatchContext + + module_config = SimpleNamespace(params=SimpleNamespace(**params)) + return PatchContext( + backend="megatron", + phase="before_train", + extra={"module_config": module_config}, + ) + + +class TestOdcRuntimeConfig(PrimusUT): + """CPU-only unit tests for primus/core/odc/runtime_config.py. + + This is the config-driven layer that the ODC before_train patch populates + (via ``odc.set_runtime_config(...)``) so the decoupled ODC primitives read + their tuning knobs from config, not os.environ. Stdlib-only -> no torch/GPU. + """ + + def setUp(self): + import odc + + # The runtime config is a process-wide singleton; reset it to defaults + # before each test so ordering never leaks state. + cfg = odc.get_runtime_config() + for f in dataclasses.fields(odc.OdcRuntimeConfig): + setattr(cfg, f.name, f.default) + os.environ.pop("PRIMUS_TURBO_ODC_GDA_PIPE", None) + + def test_defaults_match_documented_env_defaults(self): + # Defaults are byte-for-byte the previous ODC_* env defaults and match + # the trainer_base.yaml odc_* defaults; assert the contract explicitly. + import odc + + cfg = odc.get_runtime_config() + self.assertEqual(cfg.p2p_backend, "mori") + self.assertEqual(cfg.mori_init, "pg") + self.assertEqual(cfg.max_buffer_size, 64 * 1024 * 1024) # 67108864 + self.assertFalse(cfg.rocshmem_gda) + self.assertIsNone(cfg.rocshmem_lib) + self.assertEqual(cfg.gda_rs_blocks, 64) + self.assertEqual(cfg.gda_pipe, 1) + self.assertEqual(cfg.gda_defer_reduce, "auto") + self.assertEqual(cfg.gda_warmup_mode, "strided") + self.assertEqual(cfg.gda_stride_bytes, 65536) + + def test_set_config_applies_known_overrides(self): + import odc + + returned = odc.set_runtime_config( + p2p_backend="rocshmem", + rocshmem_gda=True, + gda_warmup_mode="hdp", + gda_rs_blocks=128, + ) + cfg = odc.get_runtime_config() + # set_config returns the same singleton it mutates in place. + self.assertIs(returned, cfg) + self.assertEqual(cfg.p2p_backend, "rocshmem") + self.assertTrue(cfg.rocshmem_gda) + self.assertEqual(cfg.gda_warmup_mode, "hdp") + self.assertEqual(cfg.gda_rs_blocks, 128) + + def test_set_config_ignores_none_and_unknown_keys(self): + import odc + + odc.set_runtime_config(p2p_backend="rocshmem") + # None must NOT clobber a previously-set value (keeps prior/default), and + # unknown keys are silently ignored -- this is what lets the patch forward + # every odc_* config item, unset ones as None, without wiping defaults. + odc.set_runtime_config(p2p_backend=None, not_a_real_field="boom") + cfg = odc.get_runtime_config() + self.assertEqual(cfg.p2p_backend, "rocshmem") + self.assertFalse(hasattr(cfg, "not_a_real_field")) + + def test_gda_pipe_bridged_to_turbo_env(self): + # gda_pipe is the one knob that must stay visible to the Primus-Turbo C++ + # device kernel (read via getenv), so set_config bridges it back to the + # PRIMUS_TURBO_ODC_GDA_PIPE env var. + import odc + + odc.set_runtime_config(gda_pipe=4) + self.assertEqual(os.environ.get("PRIMUS_TURBO_ODC_GDA_PIPE"), "4") + + def test_odc_package_reexports_runtime_config_api(self): + # The before_train patch depends on `import odc` exposing the runtime + # config API WITHOUT importing the heavy primitives (torch/triton/mori). + import odc + import odc.runtime_config as rc + + self.assertIs(odc.set_runtime_config, rc.set_config) + self.assertIs(odc.get_runtime_config, rc.get_config) + # Importing odc must not eagerly pull in the heavy primitives (whose + # import-time backend selection must run AFTER set_runtime_config). + self.assertNotIn("odc.primitives.scatter_accumulate", sys.modules) + + +@unittest.skipUnless(_HAS_TORCH, "ODC patch modules import torch at module scope") +class TestOdcPatchWiring(PrimusUT): + """CPU-only patch-registration / gating tests. + + Imports the ODC before_train patch modules (which need torch importable, but + NOT a GPU) and asserts that each patch is registered and that its condition + gates exactly on the ODC config items. No training is run. + """ + + def _import_patches(self, dotted: str): + try: + return importlib.import_module(dotted) + except Exception as e: # noqa: BLE001 -- torch present but megatron stack absent + self.skipTest(f"cannot import {dotted} ({type(e).__name__}: {e})") + + def test_odc_fsdp2_patch_registered_and_gated(self): + self._import_patches("primus.backends.megatron.patches.odc_torch_fsdp2_patches") + from primus.core.patches import PatchRegistry + + patch = PatchRegistry.get("megatron.fsdp.odc_torch_fsdp2") + self.assertIsNotNone(patch, "ODC FSDP2 patch must be registered") + self.assertEqual(patch.backend, "megatron") + self.assertEqual(patch.phase, "before_train") + + # Off by default (no enable_odc). + self.assertFalse(patch.applies_to(_make_ctx(use_torch_fsdp2=True))) + # enable_odc alone is not enough -- ODC requires the torch-FSDP2 path. + self.assertFalse(patch.applies_to(_make_ctx(enable_odc=True, use_torch_fsdp2=False))) + # Both set -> ODC integrates. + self.assertTrue(patch.applies_to(_make_ctx(enable_odc=True, use_torch_fsdp2=True))) + + def test_odc_lb_mini_patch_gated_independently_of_enable_odc(self): + mod = self._import_patches("primus.backends.megatron.patches.odc_lb_mini_patches") + from primus.core.patches import PatchRegistry + + patch = PatchRegistry.get("megatron.fsdp.odc_lb_mini") + self.assertIsNotNone(patch, "ODC LB-Mini patch must be registered") + + # LB-Mini serving is gated by enable_odc_lb_mini + use_torch_fsdp2 ALONE, + # orthogonal to the enable_odc comm switch. + self.assertFalse(patch.applies_to(_make_ctx(use_torch_fsdp2=True))) + self.assertFalse(patch.applies_to(_make_ctx(enable_odc_lb_mini=True, use_torch_fsdp2=False))) + self.assertTrue(patch.applies_to(_make_ctx(enable_odc_lb_mini=True, use_torch_fsdp2=True))) + # enable_odc is NOT required for LB-Mini to install. + self.assertTrue( + patch.applies_to(_make_ctx(enable_odc_lb_mini=True, use_torch_fsdp2=True, enable_odc=False)) + ) + + # ...but enable_odc selects the micro-batch alignment MODE: OFF -> ALIGNED + # (NCCL-safe same-count baseline), ON -> DECOUPLED (per-rank counts). + self.assertTrue(mod._lb_mini_aligned(SimpleNamespace(enable_odc=False))) + self.assertFalse(mod._lb_mini_aligned(SimpleNamespace(enable_odc=True))) + + def test_device_id_patch_skipped_under_odc(self): + # #856's eager-RCCL device_id injection must be SKIPPED when ODC is on + # (ODC drives gradient exchange over rocSHMEM P2P; the eager RCCL comms + # would serialize its XGMI copy streams). Assert the condition flips. + self._import_patches("primus.backends.megatron.patches.distributed_init_patches") + from primus.core.patches import PatchRegistry + + patch = PatchRegistry.get("megatron.distributed.init_process_group_device_id") + self.assertIsNotNone(patch, "device_id init patch must be registered") + # FSDP2 without ODC -> device_id patch applies. + self.assertTrue(patch.applies_to(_make_ctx(use_torch_fsdp2=True, enable_odc=False))) + # FSDP2 with ODC -> device_id patch is skipped. + self.assertFalse(patch.applies_to(_make_ctx(use_torch_fsdp2=True, enable_odc=True))) + + def test_populate_runtime_config_from_trainer_config(self): + # The bridge that copies odc_* trainer-config items into the ODC library + # runtime config at before_train. Exercises the real function end to end. + import odc + + mod = self._import_patches("primus.backends.megatron.patches.odc_torch_fsdp2_patches") + + cfg = odc.get_runtime_config() + for f in dataclasses.fields(odc.OdcRuntimeConfig): + setattr(cfg, f.name, f.default) + os.environ.pop("PRIMUS_TURBO_ODC_GDA_PIPE", None) + + ctx = _make_ctx( + odc_p2p_backend="rocshmem", + odc_rocshmem_gda=True, + odc_gda_warmup_mode="hdp", + odc_gda_pipe=3, + odc_gda_defer_reduce=1, + ) + mod._populate_odc_runtime_config(ctx) + + cfg = odc.get_runtime_config() + self.assertEqual(cfg.p2p_backend, "rocshmem") + self.assertTrue(cfg.rocshmem_gda) + self.assertEqual(cfg.gda_warmup_mode, "hdp") + self.assertEqual(cfg.gda_pipe, 3) + # defer_reduce is stringified by the bridge (config int -> library str). + self.assertEqual(cfg.gda_defer_reduce, "1") + # ...and gda_pipe is mirrored to the turbo C++ env. + self.assertEqual(os.environ.get("PRIMUS_TURBO_ODC_GDA_PIPE"), "3") + + +class TestOdcConfigFixture(PrimusUT): + """CPU-only assertions on the ODC training-config fixture. + + Verifies the fixture actually declares the ODC path so a real-GPU run would + exercise it (and documents the required co-settings), without needing the + heavy Primus/megatron config stack. + """ + + def _load_overrides(self): + try: + import yaml + except ImportError: + self.skipTest("PyYAML not installed") + self.assertTrue(os.path.exists(_ODC_FIXTURE), f"missing ODC fixture: {_ODC_FIXTURE}") + with open(_ODC_FIXTURE) as f: + cfg = yaml.safe_load(f) + return cfg["modules"]["pre_trainer"]["overrides"] + + def test_fixture_enables_odc_and_lb_mini_on_fsdp2(self): + ov = self._load_overrides() + # ODC comm on, full wiring depth, LB-Mini on -- the feature under test. + self.assertIs(ov["enable_odc"], True) + self.assertEqual(ov["odc_phase"], 2) + self.assertIs(ov["enable_odc_lb_mini"], True) + # ODC + LB-Mini require the torch-FSDP2 path (the patch condition). + self.assertIs(ov["use_torch_fsdp2"], True) + # A valid symmetric-memory backend must be selected. + self.assertIn(ov["odc_p2p_backend"], ("mori", "rocshmem")) + + def test_fixture_uses_odc_compatible_fsdp2_settings(self): + ov = self._load_overrides() + # ODC replaces the collective reduce-scatter, so these must be off (they + # would otherwise conflict with / be redundant to ODC's transport). + self.assertIs(ov["use_distributed_optimizer"], False) + self.assertIs(ov["overlap_grad_reduce"], False) + self.assertIs(ov["overlap_param_gather"], False) + # Keep the model tiny so a real-GPU run is cheap. + self.assertLessEqual(int(ov["train_iters"]), 5) + + +_GFX_TO_PLATFORM = {"gfx942": "MI300X", "gfx950": "MI355X"} + + +class TestOdcMegatronTrainerE2E(PrimusUT): + """End-to-end ODC training smoke test on the torch-FSDP2 path. + + ODC needs ROCm GPUs + a Primus-Turbo build carrying the rocSHMEM ops, so this + cannot run in plain CPU CI -- it skips with a clear reason when CUDA / turbo + are unavailable (same GPU-gating pattern as the neighbor trainer tests). + """ + + def _require_odc_hardware(self): + if not _HAS_TORCH: + self.skipTest("torch not available") + import torch + + if not torch.cuda.is_available() or torch.cuda.device_count() < 1: + self.skipTest("ODC needs a ROCm GPU; none available") + if importlib.util.find_spec("primus_turbo") is None: + self.skipTest("ODC rocSHMEM ops need a Primus-Turbo build (primus_turbo not importable)") + props = torch.cuda.get_device_properties(0) + arch = (getattr(props, "gcnArchName", "") or "").split(":")[0].strip() + platform = _GFX_TO_PLATFORM.get(arch) + if platform is None: + self.skipTest(f"unrecognized GPU arch {arch!r}; ODC validated on gfx942/gfx950") + return platform + + def test_odc_fsdp2_pretrain_smoke(self): + self._require_odc_hardware() + + env = os.environ.copy() + env["EXP"] = _ODC_FIXTURE + ut_log_path = os.environ.get("UT_LOG_PATH", "ut_out") + train_log_path = os.path.join(ut_log_path, "log.test_odc_megatron_trainer-fsdp2_smoke.txt") + env["TRAIN_LOG"] = train_log_path + # ODC's XGMI/rocSHMEM P2P bootstraps over the loopback iface on a single + # node; mirror the launcher's infra env. + env.setdefault("NCCL_SOCKET_IFNAME", "lo") + env.setdefault("GLOO_SOCKET_IFNAME", "lo") + + cmd = [ + "bash", + "./runner/primus-cli", + "direct", + "--log_file", + train_log_path, + "--", + "train", + "pretrain", + "--config", + _ODC_FIXTURE, + "--num_layers", + "2", + "--train_iters", + "3", + ] + stdout, _ = run_training_script( + tag="odc_fsdp2_smoke", cmd=cmd, train_log_path=train_log_path, env=env + ) + self.assertIn("Training completed.", stdout) + # ODC's before_train patch must have wired itself in. + self.assertIn("[ODC.torch_fsdp2]", stdout) + + +if __name__ == "__main__": + unittest.main(buffer=False) diff --git a/tests/trainer/test_odc_megatron_trainer.yaml b/tests/trainer/test_odc_megatron_trainer.yaml new file mode 100644 index 000000000..b96939ce5 --- /dev/null +++ b/tests/trainer/test_odc_megatron_trainer.yaml @@ -0,0 +1,108 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: &exp_name ${EXP:exp-odc-pretrain} +workspace: ./output + +# ODC (On-Demand Communication) integration test config. +# +# Modeled on tests/trainer/test_megatron_trainer.yaml but with the ODC feature +# turned on: rocSHMEM/XGMI point-to-point gradient reduction on the torch-FSDP2 +# path (enable_odc) plus LB-Mini sequence-length load balancing +# (enable_odc_lb_mini). A tiny model / few iters keeps a real GPU run cheap. +# +# NOTE: ODC needs ROCm GPUs + a Primus-Turbo build carrying the rocSHMEM ops, so +# a full training run cannot happen in plain CPU CI (the E2E test that consumes +# this fixture skips cleanly when CUDA / turbo are unavailable). ODC + LB-Mini +# also require use_torch_fsdp2=true (asserted by the config test), which is why +# the FSDP2 co-settings below (no distributed optimizer, no grad/param overlap) +# mirror the shipped examples/megatron/configs/*/deepseek1.5B-odc-lbmini.yaml. + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: ${PRIMUS_MODEL:deepseek_v2_lite}.yaml + overrides: + # log + wandb_project: "Primus_test_ODC" + disable_wandb: true + stderr_sink_level: DEBUG + + # debug: keep the model tiny for a cheap real-GPU run + num_layers: ${PRIMUS_NUM_LAYERS:2} + moe_layer_freq: ${PRIMUS_MOE_LAYER_FREQ:[0]*1+[1]*1} + + # optimizer: adam + moe_router_force_load_balancing: true + moe_router_dtype: null + log_avg_skip_iterations: 2 + log_avg_reset_interval: 5 + + # hyper parameters + train_iters: 3 + micro_batch_size: 1 + global_batch_size: ${PRIMUS_GLOBAL_BATCH_SIZE:8} + seq_length: ${PRIMUS_SEQ_LENGTH:4096} + max_position_embeddings: ${PRIMUS_MAX_POSITION_EMBEDDINGS:4096} + lr: 1.0e-5 + min_lr: 0.0 + lr_warmup_iters: 2 + lr_decay_iters: null + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: true + init_method_std: 0.008 + norm_epsilon: 1.0e-6 + + # === ODC (On-Demand Communication) === + # Master switch: route FSDP2 gradient reduce-scatter through ODC's + # rocSHMEM/XGMI point-to-point path. Requires use_torch_fsdp2=true. + enable_odc: true + # Full grad-routing wiring depth (hook train_step + optimizer.step). + odc_phase: 2 + # rocSHMEM symmetric-memory backend (validated single-node host/XGMI-IPC). + odc_p2p_backend: rocshmem + # Guard against ODC async-reduce non-deterministic grad spikes. + odc_grad_spike_threshold: 1000.0 + # LB-Mini sequence-length load balancing (orthogonal to enable_odc; with + # enable_odc=true it runs the DECOUPLED per-rank micro-batch mode). + enable_odc_lb_mini: true + lb_mini_cost_model: fit + lb_mini_max_token_len: 8192 + + # torch-FSDP2 path + ODC-compatible co-settings (mirror the shipped ODC + # example): ODC replaces the collective reduce-scatter, so distributed + # optimizer / grad+param overlap / grad-accum fusion are turned off. + use_torch_fsdp2: true + use_megatron_fsdp: false + use_distributed_optimizer: false + overlap_grad_reduce: false + overlap_param_gather: false + gradient_accumulation_fusion: false + + # parallel + tensor_model_parallel_size: ${PRIMUS_TP:1} + pipeline_model_parallel_size: ${PRIMUS_PP:1} + expert_model_parallel_size: ${PRIMUS_EP:1} + + # data + num_workers: 1 + mock_data: true + train_data_path: null + valid_data_path: null + test_data_path: null + + # fusion + moe_permute_fusion: false + moe_use_legacy_grouped_gemm: true + + # ckpt + finetune: false + auto_continue_train: false + load: null + save: null + save_interval: 20000 + disable_last_saving: true + ckpt_format: torch From 504504886ef90d5d9ff6d8d6a9b27553d86b4bb9 Mon Sep 17 00:00:00 2001 From: botaohu001 Date: Tue, 21 Jul 2026 09:36:30 +0000 Subject: [PATCH 2/2] test(odc): gate E2E on actual ODC turbo op, not just primus_turbo presence MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The E2E smoke needs Primus-Turbo built WITH the rocSHMEM ops (a stock/pip primus_turbo is compiled with -DDISABLE_ROCSHMEM and has neither op). Replace the coarse find_spec("primus_turbo") guard with a hasattr check for the op this single-node rocshmem fixture uses (odc_rocshmem_host), so a "GPU present but turbo lacks ODC ops" host skips cleanly instead of crashing mid-run. Comment points to building a host-capable Primus-Turbo per the repro guide (§6.3). Co-authored-by: Cursor --- tests/trainer/test_odc_megatron_trainer.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/tests/trainer/test_odc_megatron_trainer.py b/tests/trainer/test_odc_megatron_trainer.py index f589b1b22..d675ee719 100644 --- a/tests/trainer/test_odc_megatron_trainer.py +++ b/tests/trainer/test_odc_megatron_trainer.py @@ -309,8 +309,23 @@ def _require_odc_hardware(self): if not torch.cuda.is_available() or torch.cuda.device_count() < 1: self.skipTest("ODC needs a ROCm GPU; none available") - if importlib.util.find_spec("primus_turbo") is None: - self.skipTest("ODC rocSHMEM ops need a Primus-Turbo build (primus_turbo not importable)") + # ODC's rocSHMEM ops must be COMPILED INTO Primus-Turbo -- a stock/pip + # primus_turbo is built with -DDISABLE_ROCSHMEM and carries NEITHER op, so a + # bare find_spec("primus_turbo") is not enough (a "GPU present but turbo lacks + # ODC ops" box would crash mid-run instead of skipping). Verify the op this + # fixture actually needs is present. This single-node (TP/PP/EP=1) rocshmem + # run uses the host/XGMI-IPC op (odc_rocshmem_host); a multi-node GDA run + # would additionally require odc_rocshmem_gda. Build a host-capable + # Primus-Turbo from source against rocSHMEM per the repro guide (§6.3) to run. + try: + import primus_turbo.pytorch._C as _turbo_c + except Exception as e: # noqa: BLE001 -- any import failure -> not runnable here + self.skipTest(f"Primus-Turbo not importable ({type(e).__name__}: {e})") + if not hasattr(_turbo_c, "odc_rocshmem_host"): + self.skipTest( + "Primus-Turbo lacks the ODC rocSHMEM host op (odc_rocshmem_host); build " + "Primus-Turbo from source against rocSHMEM (repro guide §6.3)" + ) props = torch.cuda.get_device_properties(0) arch = (getattr(props, "gcnArchName", "") or "").split(":")[0].strip() platform = _GFX_TO_PLATFORM.get(arch)