feat(miles): allow tms preload on Blackwell with cu13+ torch wheels#24
Closed
zhenyulincs wants to merge 2 commits into
Closed
feat(miles): allow tms preload on Blackwell with cu13+ torch wheels#24zhenyulincs wants to merge 2 commits into
zhenyulincs wants to merge 2 commits into
Conversation
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.
Author
|
Closing per maintainer decision: no miles-side code change needed. The tms arch guard already ships an escape hatch ( |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Make
assert_tms_hook_mode_matches_archCUDA-version-aware:preloadhook mode on Blackwell (cc major >= 10) is now allowed whentorch.version.cudaparses to major >= 13. Pre-cu13 / unknown CUDA stays blocked (RuntimeError with updated remedies); theMILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL=1escape hatch is unchanged.Also fixes env plumbing so the knobs actually reach the actor process:
MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL+MILES_MAX_RESIDUAL_GPU_MEM_GBactor_group.pyspawn paths now forwardMILES_TMS_HOOK_MODE+ the escape hatch (the legacy colocate path forwarded neither; the guard reads them inside the actor)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=0in 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
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).