diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 9d93db1199..08b48c6296 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -615,12 +615,6 @@ def _ensure_cpu_bucket_cache(self): max_bytes = int(getattr(self.args, "miles_model_update_bucket_size_mb", 512)) * 1024 * 1024 self._cpu_bucket_cache = CPUBucketCache(max_bucket_size_bytes=max_bytes) - # F20: bucket build / sync session each acquire this lock for - # the whole critical section (single-method-single-critical- - # section). Cross-RPC locking is forbidden (Layer 1 / F04). - import threading as _threading - - self._cache_lock = _threading.Lock() return self._cpu_bucket_cache @staticmethod @@ -697,65 +691,56 @@ def build_cpu_bucket_cache(self, step: int) -> int: # weights. megatron_local_weights = self.weights_backuper.get("actor") - with self._cache_lock: - if not is_owner: - # Non-owner ranks must still drive the collective gather - # (each chunk is implicitly cross-rank inside - # _get_megatron_full_params + all_gather_params_async), - # but they discard the resulting tensors and advance - # their pointer in lockstep. - for _ in iterator.get_hf_weight_chunks(megatron_local_weights): - pass - cache.put_empty_step(int(step)) - return int(step) - - # F4 bucket layout (cache_owner only): pack (name, tensor) - # pairs into buckets up to max_bucket_size_bytes each. The - # iterator yields chunks of (name, hf_tensor); tensors come - # back on the GPU device (cuda.current_device()) for the - # standalone broadcast path, so we materialize to CPU here - # before storing — BucketEntry rejects CUDA tensors per the - # cpu_serialize transport contract. - max_bytes = cache.max_bucket_size_bytes - buckets: list[BucketEntry] = [] - current: dict[str, torch.Tensor] = {} - current_bytes = 0 - current_elements = 0 - current_idx = 0 - for chunk in iterator.get_hf_weight_chunks(megatron_local_weights): - for name, tensor in chunk: - if not isinstance(tensor, torch.Tensor): - continue - if tensor.is_cuda: - tensor = tensor.detach().to("cpu") - tensor_bytes = tensor.element_size() * tensor.numel() - if current_bytes + tensor_bytes > max_bytes and current: - buckets.append( - BucketEntry( - bucket_index=current_idx, - params=current, - size_bytes=current_bytes, - element_count=current_elements, - ) + if not is_owner: + # Non-owner ranks must still drive the collective gather, but + # discard the resulting tensors and advance their pointer. + for _ in iterator.get_hf_weight_chunks(megatron_local_weights): + pass + cache.put_empty_step(int(step)) + return int(step) + + # F4 bucket layout (cache_owner only): pack (name, tensor) pairs + # into CPU buckets up to max_bucket_size_bytes each. + max_bytes = cache.max_bucket_size_bytes + buckets: list[BucketEntry] = [] + current: dict[str, torch.Tensor] = {} + current_bytes = 0 + current_elements = 0 + current_idx = 0 + for chunk in iterator.get_hf_weight_chunks(megatron_local_weights): + for name, tensor in chunk: + if not isinstance(tensor, torch.Tensor): + continue + if tensor.is_cuda: + tensor = tensor.detach().to("cpu") + tensor_bytes = tensor.element_size() * tensor.numel() + if current_bytes + tensor_bytes > max_bytes and current: + buckets.append( + BucketEntry( + bucket_index=current_idx, + params=current, + size_bytes=current_bytes, + element_count=current_elements, ) - current_idx += 1 - current = {} - current_bytes = 0 - current_elements = 0 - current[name] = tensor - current_bytes += tensor_bytes - current_elements += tensor.numel() - if current: - buckets.append( - BucketEntry( - bucket_index=current_idx, - params=current, - size_bytes=current_bytes, - element_count=current_elements, ) + current_idx += 1 + current = {} + current_bytes = 0 + current_elements = 0 + current[name] = tensor + current_bytes += tensor_bytes + current_elements += tensor.numel() + if current: + buckets.append( + BucketEntry( + bucket_index=current_idx, + params=current, + size_bytes=current_bytes, + element_count=current_elements, ) + ) - cache.put_step(int(step), buckets) + cache.put_step(int(step), buckets) return int(step) def run_sync_session(self, plan) -> int: @@ -764,10 +749,7 @@ def run_sync_session(self, plan) -> int: Per scope F04 (Layer 1 forbidden) the cache_owner exposes ONE top-level Ray method for transporting a sync session; helpers (per-engine cpu_serialize, NCCL group setup/broadcast/destroy) - are in-method private helpers, NOT separate Ray RPCs. The - ``_cache_lock`` is held for the whole transport phase so the - bucket list snapshot + payload generation cannot be torn by a - concurrent build_cpu_bucket_cache. + are in-method private helpers, NOT separate Ray RPCs. ``plan`` is a plain mapping (per cozy-plan §Shared protocol contract — RLix may use a frozen dataclass internally but @@ -845,56 +827,47 @@ def run_sync_session(self, plan) -> int: world_size: int = int(plan["world_size"]) cache = self._ensure_cpu_bucket_cache() - with self._cache_lock: - buckets = cache.get_step(version) - if not buckets: - logger.info( - "run_sync_session sync_id=%s version=%s found 0 buckets — empty publish", - sync_id, - version, - ) - return version - - # Path A: cpu_serialize per-engine RPC (tmpfs payload). The - # wrapper owns the tmpfs file lifecycle (try/finally - # os.unlink) per scope F28; payload is materialized once - # per (bucket, engine) pair so peak /dev/shm = 1× bucket - # size (serial per-bucket receiver invocation). - for bucket in buckets: - if not cpu_serialize_local_ranks: - break - self._dispatch_cpu_serialize_bucket( - sync_id=sync_id, - bucket=bucket, - target_handles={ - idx: target_handles[idx] - for idx in cpu_serialize_local_ranks - if idx in target_handles - }, - ) + buckets = cache.get_step(version) + if not buckets: + logger.info( + "run_sync_session sync_id=%s version=%s found 0 buckets — empty publish", + sync_id, + version, + ) + return version + + # Path A: cpu_serialize per-engine RPC (tmpfs payload). The + # wrapper owns the tmpfs file lifecycle (try/finally os.unlink). + for bucket in buckets: + if not cpu_serialize_local_ranks: + break + self._dispatch_cpu_serialize_bucket( + sync_id=sync_id, + bucket=bucket, + target_handles={ + idx: target_handles[idx] + for idx in cpu_serialize_local_ranks + if idx in target_handles + }, + ) - # Path B: NCCL broadcast non-colocate path. Set up a dynamic - # group with TCP rendezvous, broadcast each bucket from - # cache_owner (rank 0), and tear the group down after the - # last bucket. F25: warmup allreduce on every CREATE; F26 - # already enforced master_port != 0 above; F03/Anti-regression - # invariant #3: is_group_exist no-op guard on destroy is - # provided by SGLangEngine.destroy_collective_group. - if broadcast_local_ranks: - self._dispatch_nccl_broadcast( - sync_id=sync_id, - buckets=buckets, - target_handles={ - idx: target_handles[idx] - for idx in broadcast_local_ranks - if idx in target_handles - }, - group_name=group_name, - master_addr=master_addr, - master_port=master_port, - comm_ranks=comm_ranks, - world_size=world_size, - ) + # Path B: NCCL broadcast non-colocate path. The sender path is + # currently guarded in _dispatch_nccl_broadcast. + if broadcast_local_ranks: + self._dispatch_nccl_broadcast( + sync_id=sync_id, + buckets=buckets, + target_handles={ + idx: target_handles[idx] + for idx in broadcast_local_ranks + if idx in target_handles + }, + group_name=group_name, + master_addr=master_addr, + master_port=master_port, + comm_ranks=comm_ranks, + world_size=world_size, + ) return version diff --git a/miles/backends/megatron_utils/update_weight/cpu_bucket_cache.py b/miles/backends/megatron_utils/update_weight/cpu_bucket_cache.py index 48ea45ec2e..8bb3dd8288 100644 --- a/miles/backends/megatron_utils/update_weight/cpu_bucket_cache.py +++ b/miles/backends/megatron_utils/update_weight/cpu_bucket_cache.py @@ -1,41 +1,11 @@ -"""F4 CPU bucket cache — HF-format weight buckets keyed by training step. - -The cache_owner rank (pp0 + dp0 + tp0 + cp0 — see F18 / scope F18 cache_owner -uniqueness) builds and stores buckets after every training step the -RLix scheduler asks for. Receivers (SGLang engines) load the buckets -during the F4-F6 selective sync atomic unit driven by -:class:`MilesModelUpdateService`. - -Design invariants ------------------ - -- **Single ready slot** (``_cache_ready_step``): only the most recently - built step is exposed for sync. Pipelines do not pin historical - versions — the M11.1 SLA is base v=-1 plus the latest training step. - Lookup by any other step raises. -- **Per-bucket payload contains NO ``weight_version``** (scope F21). The - weight_version is published exactly once per sync via - ``manager.set_weight_version`` at the end of the atomic unit, never - per-bucket. -- **HF-format gather** (scope F18): names + shapes + dtypes are HF / - HuggingFace conventions, not Megatron-internal. The Megatron→HF - conversion runs upstream in - :class:`MegatronTrainRayActor.build_cpu_bucket_cache` (iter 11) and - passes already-converted tensors into :meth:`CPUBucketCache.put`. -- **Cache owner uniqueness**: only the cache_owner rank actually stores - bucket data. Non-cache_owner ranks instantiate their own - :class:`CPUBucketCache` for per-rank state, but they MUST NOT call - ``put`` — the receive-side ``run_sync_session`` body (iter 12) drives - data transport through the cache_owner only. -- **No per-rank version inversion** (scope F20): publishing a new - ``_cache_ready_step`` MUST happen inside the same critical section - that wrote the bucket list. The class itself is single-method-single- - critical-section friendly: ``put_step`` builds and publishes - atomically; ``get_step`` is a pure read. -- **Tmpfs naming convention** (scope F66): callers that materialize - bucket payloads to ``/dev/shm`` for the cpu_serialize transport - use the format ``miles_cpu_bucket_{uuid}.pt`` so leak detection is - grep-friendly (``ls /dev/shm | grep miles_cpu_bucket_``). +"""CPU bucket cache for HF-format training weights. + +Contract: +- only the cache_owner train rank stores non-empty buckets; +- only the latest ready step is retained and readable; +- bucket payloads do not carry weight versions; +- tmpfs payload files use the ``miles_cpu_bucket_`` prefix for receiver + handoff and leak detection. """ from __future__ import annotations @@ -50,7 +20,7 @@ logger = logging.getLogger(__name__) -# F66 — leak-detection-friendly file naming for the cpu_serialize transport. +# Leak-detection-friendly file naming for the cpu_serialize transport. TMPFS_FILE_PREFIX = "miles_cpu_bucket_" @@ -62,9 +32,8 @@ class BucketEntry: (or pinned RAM). The tensors are HF-format (post Megatron→HF conversion); receivers load them by name. - ``size_bytes`` is the post-conversion total payload size used for the - F10 startup S2 / S3a-2 capacity checks; ``element_count`` is the - count of scalar elements (debug-only). + ``size_bytes`` is the post-conversion total payload size; ``element_count`` + is the count of scalar elements (debug-only). """ bucket_index: int @@ -94,15 +63,8 @@ class CPUBucketCache: Only the cache_owner rank holds non-empty buckets. Other ranks instantiate this class but call ``put_empty_step`` so their ``_cache_ready_step`` advances in lockstep without retaining bucket - payloads. - - Thread safety: a single ``threading.Lock`` guards bucket + ready-step - mutation. The cache_owner builds buckets serially during a training - step (Megatron->HF gather is collective; the lock serializes - publishing within the actor process). Read-side lookup - (:meth:`get_step` / :meth:`is_ready_for`) is also lock-guarded so - publish/lookup races against ``MilesModelUpdateService`` see a - consistent snapshot. + payloads. The internal lock protects the bucket list and ready-step + marker as one piece of state. """ def __init__(self, *, max_bucket_size_bytes: int): @@ -143,13 +105,7 @@ def put_step(self, step: int, buckets: Iterable[BucketEntry]) -> None: Discards any prior step's data (single-ready-slot invariant). Validates that bucket sizes do not exceed - ``max_bucket_size_bytes`` so the F10 startup capacity check - (S2 / S3a-2) holds at runtime too. - - Atomic: the bucket-list write and the ``_cache_ready_step`` - advance happen in the same critical section so concurrent - readers see either the old (step, buckets) tuple or the new - one — never a torn state. + ``max_bucket_size_bytes``. """ bucket_list = list(buckets) for entry in bucket_list: @@ -176,10 +132,7 @@ def put_empty_step(self, step: int) -> None: Used by non-cache_owner ranks: they participate in the collective gather (so the cache_owner can produce HF-format weights) but - discard the locally-constructed tensors. Their - ``_cache_ready_step`` still advances in lockstep so the rest of - the F4-F6 atomic unit can verify cross-rank readiness if it - wants to. + discard the locally-constructed tensors. """ step_int = int(step) with self._lock: @@ -223,12 +176,11 @@ def get_total_bytes(self, step: int) -> int: @staticmethod def make_tmpfs_filename(bucket_index: int, *, suffix: str = ".pt") -> str: - """F66: produce a leak-detection-friendly tmpfs file name. + """Produce a leak-detection-friendly tmpfs file name. - Used by :class:`MilesModelUpdateService` (iter 19) when - materializing a bucket onto ``/dev/shm`` for the cpu_serialize - transport. The wrapper code is responsible for ``try/finally - os.unlink`` of the returned path; the SGLang server-side route - only reads. + Used when materializing a bucket onto ``/dev/shm`` for the + cpu_serialize transport. The wrapper code is responsible for + ``try/finally os.unlink`` of the returned path; the SGLang + server-side route only reads. """ return f"{TMPFS_FILE_PREFIX}{bucket_index:04d}_{uuid.uuid4().hex}{suffix}"