Skip to content

Simplify CPU bucket cache locking #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
TianyeGGBond wants to merge 2 commits into
base: zhenyu/m11-mvp-test
Choose a base branch
from
Open

Simplify CPU bucket cache locking #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f318518
Simplify CPU bucket cache locking
TianyeGGBond Jun 4, 2026
f47642b
Clarify CPU bucket cache contracts
TianyeGGBond Jun 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
201 changes: 87 additions & 114 deletions miles/backends/megatron_utils/actor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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

Expand Down
88 changes: 20 additions & 68 deletions miles/backends/megatron_utils/update_weight/cpu_bucket_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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_"


Expand All @@ -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
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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}"