Repository navigation
swift_megatron: map resumed anchors and add a host-memory census - #343
Merged
Merged
Conversation
michaellchung
force-pushed
the
pr/swift-megatron-resume-rejoin
branch
from
November 11, 2026 00:51
01e4573 to
f10af40
Compare
michaellchung
force-pushed
the
pr/swift-megatron-host-memory
branch
2 times, most recently
from
November 11, 2026 00:52
74fcc33 to
dc98e5b
Compare
michaellchung
force-pushed
the
pr/swift-megatron-resume-rejoin
branch
from
November 11, 2026 00:52
f10af40 to
3b5b378
Compare
A resume with saved anchors needs the old anchor of every shard of the fragments the syncer moved past, beside the new anchors the catch-up installs. _load_resume read the whole per-rank file into a second resident set: for the 120B fast island's pp0 stage that is 1869 shards and 5.377e9 FP32 elements, +30.0 GiB per rank (+240 GiB per 8-card host) measured, held from the restart until the last fragment lands -- the tightest point of a restart, and the island only ever reads the ~1/8 of it that belongs to a shifted fragment. Map the file instead (torch.load(mmap=True), the new default): the shards the restart never touches stay out of RSS, and the ones it does touch are reclaimable file-backed pages rather than AnonPages. The values are identical, so this is a residency change and not a numerical one; YETO_DILOCO_RESUME_ANCHOR_MMAP=0 restores the eager read, and a checkpoint torch cannot map (an old non-zipfile one) falls back to it with a warning instead of losing its anchors. Also add anchor_dtype / YETO_DILOCO_ANCHOR_DTYPE, which halves the resident anchors (20.03 -> 10.02 GiB per rank for that stage). It rounds both `local - anchor` and the resume's `global_new + (local - anchor_old)`, so the default stays the broadcast dtype -- FP32 for an int4 or int8 push session, as production runs -- and bf16 is for an island that otherwise does not fit. The record notes which dtype wrote the anchors; a resume reads whatever the file holds, since either way it is the anchor the island really pushed from. Measured on one node with the production shard shapes, CPU only: torch.save of the whole dict costs +0.357 GiB, so the save side was never the problem; the eager torch.load is what costs the 30 GiB. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The island's host memory is what ends these runs, and the numbers we have are not good enough to tell a leak from a sawtooth. Production samples only /proc/meminfo MemAvailable, which mixes page cache with anonymous memory, and the one per-rank signal a callback did log came from the accelerator's host allocator, where only some counters are gauges. Measured on one card: allocated_bytes.current is a gauge (the bytes held from the OS -- freeing a tensor leaves it alone, only host_empty_cache() brings it down), while active_bytes.freed never moves, so active_bytes.current is every byte ever handed out and rises for the life of the process by construction. On the fast island that running total rose +0.155 GiB per rank per step for 90 steps while the gauge sat at 28 GiB and MemAvailable on those nodes did not fall at all. Read as a gauge it is a leak that is not there. Measure the thing itself instead. host_memory_census() returns this rank's VmRSS / VmHWM, its smaps_rollup Anonymous / Private_Dirty, and the node's MemAvailable / AnonPages, with the allocator's byte counters beside them rather than in place of them -- and each one as .current, .freed and .peak, so a gauge can be told from a running total by whether .freed ever moves instead of being assumed. Four small /proc reads and one allocator query, cheap enough per step. Deriving a quantity from those two keys is also why the ops-side callback never did anything: it released the host cache only when allocated_bytes.current - active_bytes.current passed a threshold, and subtracting a running total from a gauge goes negative and stays there, so the release never ran. YETO_DILOCO_HOST_MEM_CENSUS=1 adds the census to the existing step log (sync rank 0). Default off, and 0 leaves the record's keys exactly as they are today. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…rding The step-log census test asserted /proc numbers on every platform; stand in for the census where /proc does not exist. List STEP_LOG and HOST_MEM_CENSUS with the other YETO_DILOCO_* variables, and keep the allocator notes to what a reader of the code needs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Small eager reads can be served from the brk heap, which /proc/self/maps names [heap], so the CI runner saw a non-empty pathname for memory that is still anonymous. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
michaellchung
force-pushed
the
pr/swift-megatron-host-memory
branch
from
November 11, 2026 00:53
dc98e5b to
75d255c
Compare
michaellchung
changed the base branch from
pr/swift-megatron-resume-rejoin
to
main
November 11, 2026 00:53
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.
Depends on #342 (stacked on
pr/swift-megatron-resume-rejoin).Summary
Two host-memory changes for resumed 120B islands: map the saved anchors instead of reading them into anonymous memory, and an optional per-step host-memory census.
Changes
torch.load(mmap=True)by default (YETO_DILOCO_RESUME_ANCHOR_MMAP=0restores the eager read; a non-zipfile checkpoint falls back with a warning). Same values; the shards a restart never touches stay out of RSS (measured +30 GiB per rank, +240 GiB per 8-card host, at the tightest point of a restart, down to ~0).YETO_DILOCO_ANCHOR_DTYPE(default: the broadcast dtype) can halve resident anchors with bf16 for an island that otherwise does not fit.YETO_DILOCO_HOST_MEM_CENSUS=1adds this rank's VmRSS / VmHWM, smaps_rollup Anonymous / Private_Dirty, the node's MemAvailable / AnonPages and the accelerator host allocator's.current / .freed / .peakcounters to the step log. Off by default; with0the record is unchanged.Testing
cargo testinsyncer/: 325 passed.pytest tests/test_swift_megatron_sync.py tests/test_swift_megatron_plugin.py tests/test_rl_full_parameter_restart_rejoin.py tests/test_dispense_data_mix_samplers.py: 82 passed.pr/swift-megatron-host-memory, this branch): 8997 passed, 181 skipped, 2 failed while three suites ran in parallel;test_dispense_data_mix_samplers::test_engine_samplers_claim_from_the_dispenser...passes on rerun (load), andtest_rl_full_parameter_restart_rejoin::test_heartbeats_keep_the_lease_while_the_pool_waits_for_initexposed the graceful-stop wait fixed in the last commit of the resume PR (passes after it). For reference,mainunder the same load: 8920 passed, 1 failed (test_integration::test_learner_budget_restart_freezes_exactly_and_marks_complete_checkpoint, timing).scripts/precommit_check.pyover the whole tree: clean.Run locally on macOS arm64, Python 3.11 with the CI install (
torch==2.8.0CPU,-r pyproject.toml --extra test),PYTHONPATH=. python -m pytest ... --timeout=600. Optional stacks that are not installed (Miles, harbor, lark-oapi, matplotlib) skip as on CI. The/procparts of the census are exercised on Linux only.Notes
/procoff Linux, plugin env table, docstring wording.🤖 Generated with Claude Code