Skip to content

feat(miles): allow tms preload on Blackwell with cu13+ torch wheels#24

Closed
zhenyulincs wants to merge 2 commits into
zhenyu/m11-mvp-testfrom
zhenyu/m11-tms-preload-cu13
Closed

feat(miles): allow tms preload on Blackwell with cu13+ torch wheels#24
zhenyulincs wants to merge 2 commits into
zhenyu/m11-mvp-testfrom
zhenyu/m11-tms-preload-cu13

Conversation

@zhenyulincs

Copy link
Copy Markdown

What

Make assert_tms_hook_mode_matches_arch CUDA-version-aware: preload hook mode on Blackwell (cc major >= 10) is now allowed when torch.version.cuda parses to major >= 13. Pre-cu13 / unknown CUDA stays blocked (RuntimeError with updated remedies); the MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL=1 escape hatch is unchanged.

Also fixes env plumbing so the knobs actually reach the actor process:

  • both rlix example drivers now forward MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL + MILES_MAX_RESIDUAL_GPU_MEM_GB
  • both actor_group.py spawn paths now forward MILES_TMS_HOOK_MODE + the escape hatch (the legacy colocate path forwarded neither; the guard reads them inside the actor)
  • stale "CUDA 12.9" comments refreshed

Why

2026-07-05 offload audit (4x RTX 5090, torch 2.11.0+cu130): the historical tms 0.0.9 preload segfault is gone on cu13 wheels — tms mock pause/resume PASSES and the full dual-pipeline E2E smoke completes with EXIT_CODE=0 in preload mode. Preload is the only hook mode whose catchment includes Megatron optimizer state (torch-hook leaves ~11.07 GiB whole-GPU residual at cycle 2 vs 5.02 GiB stable under preload), which unblocks lowering the rlix residual gate from 13.0 to 7.0 (paired rlix PR).

The guard's dev-note anticipated exactly this follow-up: relax once a fixed tms/CUDA build is confirmed.

Tests

PYTHONPATH=. python -m pytest tests/fast/backends/test_tms_utils.py tests/fast/utils/test_rlix_env_forwarding.py -q
32 passed  (on 4x RTX 5090 cu130; torch.version.cuda == '13.0' confirmed)

New coverage: cu13.0/13.1 allowed (env preload AND env-unset default), cu12.9/None/unparseable raise, escape hatch on cu12.9, pre-Blackwell unchanged; source-level tests pin the four env-forwarding requirements.

E2E

Paired trio (this + #5 + the rlix threshold-7.0 PR) dual smoke on 4x RTX 5090: EXIT_CODE=0, whole-GPU gate readings 3.45-3.47 GB (cycle 1) / 5.02-5.04 GB (cycle 2) vs threshold 7.00 GB.

Codex review: APPROVE_WITH_NOTES, 0 blockers.

Refs: rlops/rlix#17 (merged; audit evidence in its description), #5 (paired sender-side diagnostics).

assert_tms_hook_mode_matches_arch becomes CUDA-version-aware: preload on
cc>=10 is allowed when torch.version.cuda parses to major >= 13. The tms
0.0.9 preload segfault is specific to pre-cu13 wheels — verified gone on
4x RTX 5090 + torch 2.11.0+cu130 (2026-07-05 audit: tms mock pause/resume
+ dual-pipeline E2E, EXIT_CODE=0). None/unparseable CUDA stays
conservative (raise); MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL escape hatch
unchanged for pre-cu13 stacks.

Also:
- forward MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL + MILES_MAX_RESIDUAL_GPU_MEM_GB
  in both rlix example drivers (operator rollback knobs must reach actors)
- forward hook mode + escape hatch on BOTH actor_group spawn paths (the
  guard reads them inside the actor process; legacy path forwarded nothing)
- refresh stale CUDA-12.9 comments in tms_utils/actor/actor_group
- extend guard test matrix (cu13.x / None / unparseable / env-unset preload
  default) + new source-level env-forwarding tests

Tests: 32 passed on 4x RTX 5090 cu130 (tests/fast/backends/test_tms_utils.py,
tests/fast/utils/test_rlix_env_forwarding.py). Codex APPROVE_WITH_NOTES, 0 blockers.
Conflict resolution: union of env-forwarding tuple additions in both rlix
drivers (this branch's MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL + #5's
MILES_MAX_RESIDUAL_GPU_MEM_GB, each once). Codex APPROVE.
@zhenyulincs

Copy link
Copy Markdown
Author

Closing per maintainer decision: no miles-side code change needed. The tms arch guard already ships an escape hatch (MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL=1) that authorizes preload on stacks where the cu12.9 segfault is confirmed gone; the rlix smoke (rlops/rlix#31) now exports it directly. The audit evidence (preload works on cu130 Blackwell, residual 5.02 GiB) stands and is recorded in rlops/rlix#17 / #31.

@zhenyulincs zhenyulincs closed this Jul 5, 2026
@zhenyulincs
zhenyulincs deleted the zhenyu/m11-tms-preload-cu13 branch July 5, 2026 17:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant