From 4d27a1fc4e631942740e47fdd9398eaeba279099 Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Fri, 24 Jul 2026 17:29:16 -0700 Subject: [PATCH 1/2] feat(optimizer): NVMe optimizer-state streaming as a miles plugin Move the store out of Megatron into miles_plugins/optimizers/nvme_stream.py, and turn the five DistributedOptimizer entry points it overrides into instance bindings here instead of `if self._nvme_state_store is not None` branches there. Pairs with radixark/Megatron-LM#63, whose optimizer diff is now empty: the only thing left on that side is the two checkpoint hooks, which have to sit inside save_checkpoint/load_checkpoint to cover every call path. The per-bucket copy of mains into the param buffer is adapted here from _copy_main_params_to_model_params rather than extracted there, with the source pinned by commit and line range. Extracting it would have handed out a partial operation across a repo boundary: the fp8 branch depends on quantize_param_shard() having run over the whole fp8 param set, which is a DP collective and cannot be split per bucket, so a subset-taking method silently skips fp8 updates unless the caller knows to compensate. Keeping the copy next to its only caller makes that precondition local, and __init__ now rejects fp8 params (and the MXFP8 param all-gather) outright. --stream-optimizer-state-to-disk is the switch, replacing Megatron's --optimizer-state-nvme-dir where "enable" and "where" were the same knob so there could be no default location. It requires --offload-train-target=disk: both mechanisms answer the same pressure and are only deployed as a pair, so that is the only combination covered. --offload-train-disk-dir and --offload-train-disk-chunk-mb are therefore reused as-is, with each mechanism taking its own subdirectory of the shared root, and their validation is untouched. --stream-optimizer-state-moment-dtype carries the on-disk dtype, which is a constructor argument now, so the silent getattr fallback on the config field is gone. It is a serialization format, not a compute precision, which is what distinguishes it from --exp-avg-dtype: the step still hands fp32 to FusedAdam, and those flags require the precision-aware optimizer that streaming excludes. Streaming also asserts against the paths that read the main params or the master optimizer's state directly (--reset-optimizer-states, --save-local-weight-checksum, --enable-witness) instead of letting them return emptiness. --- docs/advanced/disk-offload.md | 75 +++- miles/backends/megatron_utils/model.py | 6 + miles/utils/arguments.py | 78 +++- miles_plugins/optimizers/__init__.py | 0 miles_plugins/optimizers/nvme_stream.py | 490 ++++++++++++++++++++++++ 5 files changed, 639 insertions(+), 10 deletions(-) create mode 100644 miles_plugins/optimizers/__init__.py create mode 100644 miles_plugins/optimizers/nvme_stream.py diff --git a/docs/advanced/disk-offload.md b/docs/advanced/disk-offload.md index 8cccf08ec2..996aedb2d8 100644 --- a/docs/advanced/disk-offload.md +++ b/docs/advanced/disk-offload.md @@ -38,12 +38,75 @@ matching `TMS_DISK_BACKUP_*` environment, and reclaiming the files at startup an Because pause and resume happen at phase boundaries, everything is resident again by the time the optimizer step runs. If the binding constraint is instead that the optimizer state does not fit the GPU *during* the step, actor offload cannot help; that case is what -Megatron's `--optimizer-state-nvme-dir` addresses (see -[radixark/Megatron-LM#63](https://github.com/radixark/Megatron-LM/pull/63)), and the two -compose. +`--stream-optimizer-state-to-disk` addresses, and the two compose. + +## Streaming the optimizer state + +```bash +--offload-train --offload-train-target disk \ +--stream-optimizer-state-to-disk \ +--offload-train-disk-dir /scratch/miles_offload +``` + +The two together are what this is for. Actor offload gets the paused actor out of HBM for +the whole rollout window; streaming keeps the largest part of it — 12 bytes per parameter +of fp32 main params and Adam moments — from being in HBM in the first place, including +while the step runs, which offload cannot do. Data parallelism divides the optimizer state, +so a run with GPUs to spare shards it small enough to stay resident and needs neither; the +runs that need both are the ones tight enough to sit at DP=1. GLM-5.2 744B on 8 GB300 nodes +is the shape of it: 279 GB of state per rank against 276.6 GB of HBM, on half the nodes the +run would otherwise need. + +The fp32 main params and Adam moments live in per-bucket files instead of on the GPU. +Each step brings in one bucket, updates it, and writes it back, so peak residency is one +bucket rather than the whole state. Buckets are capped at 200M elements independently of +DDP's bucket sizes, which reach tens of GB at DP=1. Native-fp32 model params (a router's +`expert_bias`, a GDN/Mamba `A_log`) stay GPU-resident under a small separate Adam: they +are tiny, and unlike the bf16 path their optimizer shards alias the model params directly. + +`fp32` storage is bit-identical to keeping the state on GPU, so turning this on does not +change results — it trades step time for memory. The step is I/O bound and the moments +tolerate less precision than the master copy, so they can be stored narrower: + +```bash +--stream-optimizer-state-moment-dtype bf16 +``` + +That cuts streaming volume by a third (12 bytes per param to 8). A checkpoint records the +dtypes it was written with and a resume verifies them, so bytes written as bf16 can never +be read back as fp32. The fp8 options work but are not recommended: `exp_avg_sq` needs +per-block scaling to survive 8-bit storage, which this does not implement. + +Two limits to know about. Resume is same-topology only — the on-disk layout follows this +rank's DP shard, so changing TP/PP/DP/EP fails the layout assert rather than resharding. +And the optimizer state is copied to the checkpoint directory synchronously, outside +`--async-save`, so expect checkpoint saves to take noticeably longer. + +The two also help each other. With the optimizer state already on disk there is that much +less to move when the actor is paused: on Qwen3-30B-A3B, sleep/wake went from 24s/8.9s to +5.2s/1.3s once the paused actor no longer carried the state. ## Choosing -- Host RAM holds the paused actor: keep the default `--offload-train-target=cpu`. It is - faster and needs no scratch space. -- Host RAM does not hold it: switch to `--offload-train-target=disk`. +`--offload-train` is the base mechanism, and it is not tied to colocation: the actor is +resident only during `train()` and sleeps for the whole rollout window either way. What +changes is who takes the freed HBM. Colocated, it is the engine on the same GPUs, so +offload is mandatory and defaults on. Disaggregated, it is whatever else you fit on the +training GPUs, which is the point when you are sizing engines against a fixed cluster. + +- Paused actor fits in host RAM: keep `--offload-train-target=cpu` (default). Fastest, no + scratch space. +- It does not fit: `--offload-train-target=disk`. +- On top of that, the optimizer state does not fit the GPU *during the step*, which offload + cannot help with: add `--stream-optimizer-state-to-disk`, and consider + `--stream-optimizer-state-moment-dtype bf16` to claw back some of the I/O cost. + +Streaming requires `--offload-train-target=disk`; miles asserts it. The two answer the same +pressure and are only deployed as a pair, so that is the only combination covered. +Streaming on its own does run — it is the one case where offload has nothing to do, because +nothing else wants the training GPUs during rollout — but nothing validates it. + +Both mechanisms share `--offload-train-disk-dir` and `--offload-train-disk-chunk-mb`. Point the +directory at real node-local NVMe: a tmpfs mount, which `/tmp` is on many systems, keeps +the data in RAM and defeats both. The chunk is a pinned host staging buffer and each +mechanism allocates its own, so enabling both costs 2x that per rank. diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 30e70661b9..466fe473ca 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -181,6 +181,12 @@ def setup_model_and_optimizer( model_chunks=model, use_gloo_process_groups=args.enable_gloo_process_groups, ) + + if args.stream_optimizer_state_to_disk: + from miles_plugins.optimizers.nvme_stream import setup_optimizer_state_streaming + + setup_optimizer_state_streaming(args, optimizer) + opt_param_scheduler = get_optimizer_param_scheduler(args, optimizer) return model, optimizer, opt_param_scheduler diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index a975a3f04d..3cd926926a 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -138,13 +138,48 @@ def add_cluster_arguments(parser): "NVMe (--offload-train-disk-dir) for the case where even CPU RAM cannot hold it." ), ) + parser.add_argument( + "--stream-optimizer-state-to-disk", + action="store_true", + help=( + "Stream the fp32 main params and Adam moments through per-bucket files on " + "node-local NVMe during optimizer.step(), bounding GPU residency to one bucket. " + "For when the optimizer state does not fit the GPU *while the step runs*: " + "--offload-train-target=disk cannot help there, because pause/resume happen at " + "phase boundaries and everything is resident again by the time Adam launches. " + "Bit-identical to keeping the state on GPU, at the cost of disk traffic every " + "step. Distinct from --offload-optimizer-states and --optimizer-cpu-offload, " + "and mutually exclusive with both." + ), + ) + parser.add_argument( + "--stream-optimizer-state-moment-dtype", + type=str, + default="fp32", + choices=["fp32", "bf16", "fp16", "fp8e4m3", "fp8e5m2"], + help=( + "On-disk dtype for the streamed Adam moments; the fp32 master copy is always " + "fp32. This is a serialization format, not a compute precision: the step still " + "hands fp32 tensors to FusedAdam, and the cast happens on the way to and from " + "disk. That makes it distinct from --exp-avg-dtype / --exp-avg-sq-dtype, which " + "change what the optimizer holds and require --use-precision-aware-optimizer, " + "so they cannot be combined with streaming at all. " + "The step is I/O bound and the moments tolerate less precision than the master " + "copy, so bf16 cuts streaming volume by a third (12 bytes per param to 8). " + "fp32 is bit-identical to keeping the moments on GPU. The fp8 options need " + "per-block scaling for exp_avg_sq to be sound, which this does not implement, " + "and are not recommended." + ), + ) parser.add_argument( "--offload-train-disk-dir", type=str, default=None, help=( - "Node-local directory for train disk-offload files when " - "--offload-train-target=disk. Should be fast local NVMe (e.g. /scratch). " + "Node-local directory for the disk-offload files, used by both " + "--offload-train-target=disk and --stream-optimizer-state-to-disk (each under " + "its own subdirectory). Should be fast local NVMe (e.g. /scratch); a tmpfs " + "mount, which /tmp is on many systems, keeps the data in RAM and defeats both. " "Files are per-process and overwritten in place every step (bounded size); " "defaults to $SCRATCH/miles_train_offload_." ), @@ -154,8 +189,10 @@ def add_cluster_arguments(parser): type=int, default=256, help=( - "Chunk size (MiB) for streaming the training actor GPU<->disk in disk-offload " - "mode. Bounds the pinned host staging buffer regardless of the total offloaded size." + "Chunk size (MiB) for the GPU<->disk transfers, i.e. the pinned host staging " + "buffer, which bounds host memory regardless of how much is moved. Used by both " + "--offload-train-target=disk and --stream-optimizer-state-to-disk, and each " + "allocates its own, so enabling both costs 2x this per rank." ), ) @@ -2853,6 +2890,39 @@ def miles_validate_args(args): f"chunk={args.offload_train_disk_chunk_mb}MB" ) + if args.stream_optimizer_state_to_disk: + assert args.use_distributed_optimizer, "--stream-optimizer-state-to-disk requires the distributed optimizer" + assert ( + args.optimizer == "adam" + ), f"--stream-optimizer-state-to-disk requires --optimizer adam, got {args.optimizer}" + assert not args.multi_lora, "--stream-optimizer-state-to-disk does not support multi-LoRA" + assert not args.optimizer_cpu_offload, "--stream-optimizer-state-to-disk excludes --optimizer-cpu-offload" + assert ( + not args.offload_optimizer_states + ), "--stream-optimizer-state-to-disk excludes --offload-optimizer-states" + assert ( + not args.use_precision_aware_optimizer + ), "--stream-optimizer-state-to-disk requires mcore to hold the fp32 main params" + assert not args.reset_optimizer_states, ( + "--reset-optimizer-states walks the master optimizer's state, which the NVMe store " + "leaves empty, so the reset would silently do nothing" + ) + assert not args.save_local_weight_checksum, ( + "--save-local-weight-checksum reads param.main_param, whose storage the NVMe store " + "resizes to 0 between steps" + ) + assert not args.enable_witness, ( + "--enable-witness reads the master optimizer's per-param state, which the NVMe store owns" + ) + assert args.offload_train_target == "disk", ( + "--stream-optimizer-state-to-disk is only supported alongside " + "--offload-train-target=disk; pass both." + ) + logger.info( + f"Streaming optimizer state to disk, dir={args.offload_train_disk_dir}, " + f"chunk={args.offload_train_disk_chunk_mb}MB, moments={args.stream_optimizer_state_moment_dtype}" + ) + if args.async_max_concurrent_samples is not None: assert args.async_max_concurrent_samples >= args.n_samples_per_prompt, ( f"--async-max-concurrent-samples ({args.async_max_concurrent_samples}) must be at least " diff --git a/miles_plugins/optimizers/__init__.py b/miles_plugins/optimizers/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/miles_plugins/optimizers/nvme_stream.py b/miles_plugins/optimizers/nvme_stream.py new file mode 100644 index 0000000000..5b1415ebb2 --- /dev/null +++ b/miles_plugins/optimizers/nvme_stream.py @@ -0,0 +1,490 @@ +"""NVMe streaming of Megatron's DistributedOptimizer state. + +The fp32 main params and Adam moments live in per-bucket files on node-local NVMe +instead of on the GPU. ``step()`` walks the buckets one at a time -- materialize, +load, per-bucket Adam step, copy the updated mains into the param buffer, write +back, release -- so GPU residency is bounded by one bucket rather than the whole +state. For runs where the state does not fit the GPU *while the step runs*, which +sleep-window offload cannot help with. + +``setup_optimizer_state_streaming`` gives each ``DistributedOptimizer`` in the chain +a store and routes the five entry points that touch optimizer state at it, so +Megatron carries no NVMe-specific behaviour. The one piece that stays on Megatron's +side is the pair of checkpoint hooks in ``megatron/training/checkpointing.py``, which +have to be inside ``save_checkpoint`` / ``load_checkpoint`` to cover every call path; +they reach this class through four methods: + + step() + refresh_main_from_model_params(copy_fn) + save_to(base_dir) + load_from(base_dir) -> bool + +Both directory arguments are checkpoint *bases*; the per-rank layout underneath is +this file's business, matching the layout of the live scratch directory. +""" + +import atexit +import errno +import json +import logging +import os +import shutil +import time +from types import MethodType +from typing import TYPE_CHECKING, NamedTuple + +import torch +from megatron.core.fp8_utils import is_float8tensor + +if TYPE_CHECKING: + from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer + +logger = logging.getLogger(__name__) + +SEGMENTS = ("main", "exp_avg", "exp_avg_sq") +DTYPES = { + "fp32": torch.float32, + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp8e4m3": torch.float8_e4m3fn, + "fp8e5m2": torch.float8_e5m2, +} +BUCKET_NUMEL_LIMIT = 200_000_000 +FP32_RESIDENT_WARN_MB = 256 +IO_ALIGN = 4096 + + +class _Entry(NamedTuple): + model_param: torch.nn.Parameter + main_param: torch.Tensor + group_index: int + + +def _align(nbytes: int) -> int: + return (nbytes + IO_ALIGN - 1) // IO_ALIGN * IO_ALIGN + + +def _resize(tensor: torch.Tensor, numel: int) -> None: + tensor.untyped_storage().resize_(numel * tensor.element_size()) + + +def _allocate_file(path: str, nbytes: int) -> int: + fd = os.open(path, os.O_RDWR | os.O_CREAT | os.O_CLOEXEC, 0o600) + try: + os.posix_fallocate(fd, 0, nbytes) + except OSError as e: + if e.errno not in (errno.EOPNOTSUPP, errno.ENOTSUP, errno.EINVAL): + raise + os.ftruncate(fd, nbytes) + return fd + + +def _rw_full(op, fd: int, offset: int, buf) -> None: + mv = memoryview(buf).cast("B") + done = 0 + while done < len(mv): + n = op(fd, [mv[done:]], offset + done) + if n <= 0: + raise OSError(f"short {op.__name__} ({n}) on optimizer state file at offset {offset + done}") + done += n + + +def plan_buckets(entries_by_ddp_bucket: dict, limit: int = BUCKET_NUMEL_LIMIT) -> list[list[_Entry]]: + planned, current, numel = [], [], 0 + for _, entries in sorted(entries_by_ddp_bucket.items(), key=lambda kv: kv[0]): + for entry in entries: + current.append(entry) + numel += entry.main_param.numel() + if numel >= limit: + planned.append(current) + current, numel = [], 0 + if current: + planned.append(current) + current, numel = [], 0 + return planned + + +class _Stager: + def __init__(self, nbytes: int): + self._buf = torch.empty(nbytes, dtype=torch.uint8, pin_memory=True) + self._bytes = self._buf.numpy() + self._device_buf = None + + def _device_staging(self, dtype: torch.dtype, numel: int, like: torch.Tensor) -> torch.Tensor: + # A cross-dtype copy between GPU and pinned host memory does not take the DMA + # path: casting on the device first and moving same-dtype bytes is ~40x faster + # for bf16 and ~100x for fp8. + size = self._buf.numel() + if self._device_buf is None or self._device_buf.device != like.device: + self._device_buf = torch.empty(size, dtype=torch.uint8, device=like.device) + return self._device_buf[: numel * dtype.itemsize].view(dtype) + + def transfer(self, fd: int, offset: int, tensor: torch.Tensor, dtype: torch.dtype, *, to_disk: bool) -> int: + flat = tensor.view(-1) + cast = dtype != flat.dtype + chunk = self._buf.numel() // dtype.itemsize + pos = 0 + while pos < flat.numel(): + numel = min(chunk, flat.numel() - pos) + host = self._buf[: numel * dtype.itemsize].view(dtype) + at = offset + pos * dtype.itemsize + nbytes = numel * dtype.itemsize + if to_disk: + if cast: + staged = self._device_staging(dtype, numel, flat) + staged.copy_(flat[pos : pos + numel]) + host.copy_(staged) + else: + host.copy_(flat[pos : pos + numel]) + _rw_full(os.pwritev, fd, at, self._bytes[:nbytes]) + else: + _rw_full(os.preadv, fd, at, self._bytes[:nbytes]) + if cast: + staged = self._device_staging(dtype, numel, flat) + staged.copy_(host) + flat[pos : pos + numel].copy_(staged) + else: + flat[pos : pos + numel].copy_(host) + pos += numel + return flat.numel() * dtype.itemsize + + +class _Bucket: + def __init__(self, path: str, entries: list[_Entry], adam, stager: _Stager, dtypes: dict): + self.path, self.entries, self.adam, self.dtypes = path, entries, adam, dtypes + self._stager = stager + self.group_indices = sorted({e.group_index for e in entries}) + self.numel = sum(e.main_param.numel() for e in entries) + + self.offsets: dict[str, list[int]] = {} + at = 0 + for segment in SEGMENTS: + self.offsets[segment] = [] + for entry in entries: + self.offsets[segment].append(at) + at += _align(entry.main_param.numel() * dtypes[segment].itemsize) + self.nbytes = at + self.fd = _allocate_file(path, at) + self.moments_ready = False + + def _tensors(self, segment: str): + for index, entry in enumerate(self.entries): + tensor = entry.main_param if segment == "main" else self.adam.state[entry.main_param][segment] + yield tensor, self.offsets[segment][index] + + def _move(self, segments, *, to_disk: bool) -> int: + moved = 0 + for segment in segments: + for tensor, offset in self._tensors(segment): + if not to_disk: + _resize(tensor, tensor.numel()) + moved += self._stager.transfer(self.fd, offset, tensor, self.dtypes[segment], to_disk=to_disk) + if to_disk: + _resize(tensor, 0) + return moved + + def fetch(self) -> int: + return self._move(SEGMENTS if self.moments_ready else SEGMENTS[:1], to_disk=False) + + def flush(self, segments=SEGMENTS) -> int: + moved = self._move(segments, to_disk=True) + self.moments_ready = self.moments_ready or tuple(segments) == SEGMENTS + return moved + + def materialize_main(self) -> None: + for tensor, _ in self._tensors("main"): + _resize(tensor, tensor.numel()) + + def allocate_moments(self) -> None: + for entry in self.entries: + state = self.adam.state.setdefault(entry.main_param, {}) + for segment in SEGMENTS[1:]: + if segment not in state: + state[segment] = torch.empty_like(entry.main_param) + _resize(state[segment], 0) + self.moments_ready = True + + +class NVMeOptimizerStateStore: + _next_uid = 0 + + def __init__( + self, distrib_optimizer: "DistributedOptimizer", dir_root: str, chunk_mb: int, moment_dtype: str = "fp32" + ): + self.dist_opt = distrib_optimizer + self.uid = NVMeOptimizerStateStore._next_uid + NVMeOptimizerStateStore._next_uid += 1 + config = distrib_optimizer.config + + assert not config.use_precision_aware_optimizer, ( + "NVMe state store requires the non-precision-aware optimizer " "(fp32 main params held by mcore)." + ) + assert not config.optimizer_cpu_offload, "NVMe state store is mutually exclusive with CPU offload." + assert ( + not config.offload_optimizer_states + ), "NVMe state store is mutually exclusive with --offload-optimizer-states." + assert not distrib_optimizer.ddp_config.use_megatron_fsdp + # See _copy_main_to_model_params for why fp8 params cannot be streamed. + assert not any( + is_float8tensor(p) for group in distrib_optimizer.model_float16_groups for p in group + ), "NVMe state store does not support fp8 params." + assert not config.reuse_grad_buf_for_mxfp8_param_ag, ( + "NVMe state store does not support the MXFP8 param all-gather, which copies mains " + "into the param buffer through a different path." + ) + + moments = moment_dtype + assert moments in DTYPES, f"unknown moment dtype {moments!r}, expected one of {sorted(DTYPES)}" + if DTYPES[moments].itemsize == 1: + logger.warning( + f"Storing Adam moments as {moments} without per-block scaling is numerically risky " + "for exp_avg_sq; bf16 is the safe way to halve moment I/O." + ) + self.dtypes = {"main": torch.float32, "exp_avg": DTYPES[moments], "exp_avg_sq": DTYPES[moments]} + + self._rank = torch.distributed.get_rank() + self._instance = distrib_optimizer.distributed_optimizer_instance_id + self.dir = os.path.join(dir_root, self.relative_dir) + shutil.rmtree(self.dir, ignore_errors=True) + os.makedirs(self.dir, exist_ok=True) + atexit.register(shutil.rmtree, self.dir, ignore_errors=True) + + self._stager = _Stager(chunk_mb * 1024 * 1024) + self.buckets = self._build_buckets() + self._fp32_group_indices, self._fp32_adam = self._build_fp32_optimizer() + + for bucket in self.buckets: + bucket.flush(segments=("main",)) + + total_gb = sum(b.nbytes for b in self.buckets) / 1024**3 + logger.info( + f"NVMe optimizer state store: {len(self.buckets)} buckets, {total_gb:.1f} GB at " + f"{self.dir} (moments stored as {moments})" + ) + + @property + def relative_dir(self) -> str: + """This store's location under any base directory -- scratch or checkpoint. + + ``uid`` is what disambiguates the chained dense and expert optimizers, which can + share an instance id. + """ + return os.path.join(f"rank{self._rank:05d}", f"opt{self._instance}_{self.uid}") + + def _build_buckets(self) -> list[_Bucket]: + by_ddp_bucket: dict[tuple, list[_Entry]] = {} + groups = zip(self.dist_opt.model_float16_groups, self.dist_opt.shard_fp32_from_float16_groups, strict=True) + for group_index, (model_group, main_group) in enumerate(groups): + for model_param, main_param in zip(model_group, main_group, strict=True): + assert main_param is not None and main_param.dtype == torch.float32 + key = self.dist_opt.model_param_gbuf_map[model_param] + by_ddp_bucket.setdefault(key, []).append(_Entry(model_param, main_param, group_index)) + + buckets = [] + for index, entries in enumerate(plan_buckets(by_ddp_bucket)): + params: dict[int, list[torch.Tensor]] = {} + for entry in entries: + params.setdefault(entry.group_index, []).append(entry.main_param) + path = os.path.join(self.dir, f"bucket{index:05d}.bin") + buckets.append(_Bucket(path, entries, self._adam_for(params), self._stager, self.dtypes)) + return buckets + + def _build_fp32_optimizer(self): + params: dict[int, list[torch.Tensor]] = {} + total_bytes = 0 + for group_index, (model_group, shard_group) in enumerate( + zip(self.dist_opt.model_fp32_groups, self.dist_opt.shard_fp32_groups, strict=True) + ): + if model_group: + params[group_index] = list(shard_group) + total_bytes += sum(p.numel() * p.element_size() for p in shard_group) + if not params: + return [], None + + total_mb = total_bytes / 1024**2 + log = logger.warning if total_mb > FP32_RESIDENT_WARN_MB else logger.info + log(f"NVMe optimizer state store: {total_mb:.1f} MB of native-fp32 params stay GPU-resident") + return sorted(params), self._adam_for(params) + + def _adam_for(self, params_by_group: dict[int, list[torch.Tensor]]): + from megatron.core.optimizer import Adam + + master_groups = self.dist_opt.optimizer.param_groups + groups = [] + for group_index in sorted(params_by_group): + group = {k: v for k, v in master_groups[group_index].items() if k != "params"} + group["params"] = params_by_group[group_index] + groups.append(group) + return Adam(groups, adam_w_mode=self.dist_opt.config.decoupled_weight_decay) + + def _sync_lr_wd(self, adam, group_indices) -> None: + master_groups = self.dist_opt.optimizer.param_groups + for group, group_index in zip(adam.param_groups, group_indices, strict=True): + group["lr"] = master_groups[group_index]["lr"] + group["weight_decay"] = master_groups[group_index]["weight_decay"] + + # Adapted from DistributedOptimizer._copy_main_params_to_model_params, which walks every + # group at once, at + # https://github.com/radixark/Megatron-LM/blob/4716f75475c78e2fc2c6f0d3af095f1681b770b4/megatron/core/optimizer/distrib_optimizer.py#L2469-L2519 + # Recheck against that revision when bumping Megatron. Its fp8 branch is absent because + # quantize_param_shard() is a DP collective over the whole fp8 param set and cannot be + # split per bucket; __init__ rejects fp8 params instead. + def _copy_main_to_model_params(self, entries: list[_Entry]) -> None: + dist_opt = self.dist_opt + for entry in entries: + param_range_map = dist_opt._get_model_param_range_map(entry.model_param) + world_range = param_range_map["gbuf_world_in_bucket"] + assert world_range.size == entry.main_param.nelement() + gbuf_index, _, bucket_id = dist_opt.model_param_gbuf_map[entry.model_param] + param_data = dist_opt.buffers[gbuf_index].buckets[bucket_id].param_data + param_data.view(-1)[world_range.start : world_range.end].copy_(entry.main_param) + + @torch.no_grad() + def step(self) -> bool: + started = time.monotonic() + read = written = 0 + for bucket in self.buckets: + read += bucket.fetch() + self._sync_lr_wd(bucket.adam, bucket.group_indices) + bucket.adam.step() + self._copy_main_to_model_params(bucket.entries) + written += bucket.flush() + if self._fp32_adam is not None: + self._sync_lr_wd(self._fp32_adam, self._fp32_group_indices) + self._fp32_adam.step() + logger.info( + f"NVMe streaming step: {len(self.buckets)} buckets, read {read / 1024**3:.1f} GB, " + f"wrote {written / 1024**3:.1f} GB in {time.monotonic() - started:.1f}s" + ) + return True + + @torch.no_grad() + def refresh_main_from_model_params(self, copy_fn) -> None: + for bucket in self.buckets: + bucket.materialize_main() + copy_fn() + for bucket in self.buckets: + bucket.flush(segments=("main",)) + + @torch.no_grad() + def save_to(self, base_dir: str) -> None: + dirpath = os.path.join(base_dir, self.relative_dir) + os.makedirs(dirpath, exist_ok=True) + manifest = { + "dtypes": {segment: str(dtype) for segment, dtype in self.dtypes.items()}, + "buckets": [ + { + "numel": bucket.numel, + "entry_numels": [e.main_param.numel() for e in bucket.entries], + "steps": [g.get("step", 0) for g in bucket.adam.param_groups], + "file": os.path.basename(bucket.path), + } + for bucket in self.buckets + ], + } + for bucket in self.buckets: + shutil.copyfile(bucket.path, os.path.join(dirpath, os.path.basename(bucket.path))) + if self._fp32_adam is not None: + torch.save(self._fp32_adam.state_dict(), os.path.join(dirpath, "fp32_resident_optimizer.pt")) + with open(os.path.join(dirpath, "manifest.json"), "w") as f: + json.dump(manifest, f) + logger.info(f"NVMe optimizer state saved: {len(self.buckets)} buckets -> {dirpath}") + + @torch.no_grad() + def load_from(self, base_dir: str) -> bool: + """Restore this store from a checkpoint base, or report that there is nothing there. + + Returns False for a checkpoint written without streaming, which keeps pre-streaming + checkpoints loadable; every other mismatch fails on the asserts below. + """ + dirpath = os.path.join(base_dir, self.relative_dir) + if not os.path.isdir(dirpath): + logger.warning(f"no NVMe optimizer state at {dirpath}; starting from a fresh optimizer state") + return False + + with open(os.path.join(dirpath, "manifest.json")) as f: + manifest = json.load(f) + + saved = manifest.get("dtypes", {segment: str(torch.float32) for segment in SEGMENTS}) + current = {segment: str(dtype) for segment, dtype in self.dtypes.items()} + assert saved == current, ( + f"NVMe state dtype mismatch: checkpoint stores {saved}, this run stores {current} " + "-- the bytes would be misread" + ) + assert len(manifest["buckets"]) == len(self.buckets), ( + f"NVMe state layout mismatch: checkpoint has {len(manifest['buckets'])} buckets, " + f"current topology builds {len(self.buckets)} (same-topology resume only)" + ) + + for bucket, meta in zip(self.buckets, manifest["buckets"], strict=True): + assert meta["numel"] == bucket.numel + assert meta["entry_numels"] == [e.main_param.numel() for e in bucket.entries] + shutil.copyfile(os.path.join(dirpath, meta["file"]), bucket.path) + for group, step in zip(bucket.adam.param_groups, meta["steps"], strict=True): + if step: + group["step"] = step + bucket.allocate_moments() + fp32_state = os.path.join(dirpath, "fp32_resident_optimizer.pt") + if self._fp32_adam is not None and os.path.isfile(fp32_state): + self._fp32_adam.load_state_dict(torch.load(fp32_state)) + logger.info(f"NVMe optimizer state loaded: {len(self.buckets)} buckets <- {dirpath}") + return True + + +def setup_optimizer_state_streaming(args, optimizer) -> None: + """Give every DistributedOptimizer in ``optimizer``'s chain a store and route it there. + + Must run before load_checkpoint: the constructor evicts the fp32 main params, and the + bindings are what keep the load path from writing into the evicted storage. + """ + from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer + + dir_root = os.path.join(args.offload_train_disk_dir, "optimizer_state") + for dist_opt in optimizer.chained_optimizers: + assert isinstance( + dist_opt, DistributedOptimizer + ), f"--stream-optimizer-state-to-disk requires the distributed optimizer, got {type(dist_opt).__name__}" + if dist_opt.is_stub_optimizer: + continue + store = NVMeOptimizerStateStore( + dist_opt, dir_root, args.offload_train_disk_chunk_mb, args.stream_optimizer_state_moment_dtype + ) + _bind(dist_opt, store) + + +def _bind(dist_opt: "DistributedOptimizer", store: NVMeOptimizerStateStore) -> None: + """Point the five DistributedOptimizer entry points that touch optimizer state at ``store``.""" + from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer + + # The all-gather tail is copied from DistributedOptimizer.step_with_ready_grads and has + # to stay in sync with it; the timer calls around it are dropped, hence the assert. + def step_with_ready_grads(self) -> bool: + assert self.config.timers is None + update_successful = store.step() + if not self.ddp_config.overlap_param_gather: + for model_chunk in self.model_chunks: + model_chunk.start_param_sync() + return update_successful + + def reload_model_params(self, state_dict=None) -> None: + store.refresh_main_from_model_params( + lambda: DistributedOptimizer.reload_model_params(self, state_dict=state_dict) + ) + + # save_to()/load_from() carry the real state; returning empties here rather than forcing + # --no-save-optim keeps opt_param_scheduler, saved under the same guard, working. + def state_dict(self): + return {"nvme_state_store": True} + + def load_state_dict(self, state_dict) -> None: + return + + def sharded_state_dict(self, model_sharded_state_dict, is_loading=False, sharding_type=None, metadata=None): + return {} + + dist_opt.step_with_ready_grads = MethodType(step_with_ready_grads, dist_opt) + dist_opt.reload_model_params = MethodType(reload_model_params, dist_opt) + dist_opt.state_dict = MethodType(state_dict, dist_opt) + dist_opt.load_state_dict = MethodType(load_state_dict, dist_opt) + dist_opt.sharded_state_dict = MethodType(sharded_state_dict, dist_opt) + dist_opt._nvme_state_store = store # how checkpointing.py finds it From c0b03e5d5f6b092899673a7893b5e20162281b7c Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Sat, 25 Jul 2026 02:50:25 -0700 Subject: [PATCH 2/2] fix(offload): reject a torch_memory_saver without the disk backend --offload-train-target=disk sets TMS_INIT_ENABLE_CPU_BACKUP=0 alongside TMS_INIT_ENABLE_DISK_BACKUP=1. A torch_memory_saver built before fzyzcjy/torch_memory_saver#80 ignores the disk variable but still honours the cpu one, so pause() frees the actor with no backup anywhere and resume() hands back uninitialized memory. Nothing fails at that point: the weights are garbage, and it surfaces a rollout later as nan ref_log_probs/log_probs and then "found NaN in local grad norm for bucket #0", which points at the model, not at the offload path. Hit on radixark/miles:dev, which ships 0.0.9.post1 (no disk backend). Diagnosing it cost a 57G checkpoint reconversion chasing the wrong suspect. Checked against the binary that actually gets LD_PRELOADed, since that is what reads the env vars -- a Python-level probe would not cover a stale .so. Verified both ways on 8xH200: the pinned build has the symbol, all three preload .so in the 0.0.9.post1 wheel do not. --- miles/ray/train/actor_factory.py | 5 +++++ miles/utils/arguments.py | 9 ++++----- .../backends/megatron_utils/test_lora_model_branches.py | 1 + 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/miles/ray/train/actor_factory.py b/miles/ray/train/actor_factory.py index c3c0116528..3c295af2cf 100644 --- a/miles/ray/train/actor_factory.py +++ b/miles/ray/train/actor_factory.py @@ -1,4 +1,5 @@ import os +from pathlib import Path import ray from ray.util.placement_group import PlacementGroup @@ -48,6 +49,10 @@ def allocate_gpus_for_actor( env_vars["LD_PRELOAD"] = dynlib_path env_vars["TMS_INIT_ENABLE"] = "1" if args.offload_train_target == "disk": + assert b"TMS_INIT_ENABLE_DISK_BACKUP" in Path(dynlib_path).read_bytes(), ( + f"{dynlib_path} has no disk backend; reinstall torch_memory_saver at the commit " + f"docker/Dockerfile pins." + ) env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "0" env_vars["TMS_INIT_ENABLE_DISK_BACKUP"] = "1" env_vars["TMS_DISK_BACKUP_CHUNK_MB"] = str(args.offload_train_disk_chunk_mb) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 3cd926926a..1536a62f23 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -2911,12 +2911,11 @@ def miles_validate_args(args): "--save-local-weight-checksum reads param.main_param, whose storage the NVMe store " "resizes to 0 between steps" ) - assert not args.enable_witness, ( - "--enable-witness reads the master optimizer's per-param state, which the NVMe store owns" - ) + assert ( + not args.enable_witness + ), "--enable-witness reads the master optimizer's per-param state, which the NVMe store owns" assert args.offload_train_target == "disk", ( - "--stream-optimizer-state-to-disk is only supported alongside " - "--offload-train-target=disk; pass both." + "--stream-optimizer-state-to-disk is only supported alongside " "--offload-train-target=disk; pass both." ) logger.info( f"Streaming optimizer state to disk, dir={args.offload_train_disk_dir}, " diff --git a/tests/fast/backends/megatron_utils/test_lora_model_branches.py b/tests/fast/backends/megatron_utils/test_lora_model_branches.py index 3fb081f6a9..3a201a0158 100644 --- a/tests/fast/backends/megatron_utils/test_lora_model_branches.py +++ b/tests/fast/backends/megatron_utils/test_lora_model_branches.py @@ -70,6 +70,7 @@ def _make_args(self, lora_rank=32, role="actor", mode="bridge"): megatron_to_hf_mode=mode, moe_use_upcycling=False, debug_disable_optimizer=False, + stream_optimizer_state_to_disk=False, load="/some/path", pretrained_checkpoint=None, # optimizer fields