From 2e937598518eeb4b975e0366683877e9317d29e8 Mon Sep 17 00:00:00 2001 From: "Gavin.Zhu" Date: Mon, 13 Jul 2026 03:40:13 +0000 Subject: [PATCH 1/7] miles: align optim to fork Ray interface + overlay fork miles in profile - backend.apply_optimizer_step: drop the non-existent set_learning_rate call; pass learning_rate into apply_optimizer_step[_and_sync](learning_rate=) to match the miles fork's RayTrainGroup interface (github.com/GavinZhu-GMI/miles main). - setup_container.sh (miles profile): overlay the fork's miles onto the base image (baked image predates the Tinker orchestration interface); default MILES_REF=main. Validated end-to-end on ns.config: create_model + 3x forward_backward + apply_optimizer_step(lr) all success=True; empty-data now correctly ValueErrors. --- scripts/lib/setup_container.sh | 12 +++++++++++- training/backends/miles/backend.py | 11 ++++++----- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/scripts/lib/setup_container.sh b/scripts/lib/setup_container.sh index 933f56d..05e163c 100755 --- a/scripts/lib/setup_container.sh +++ b/scripts/lib/setup_container.sh @@ -61,7 +61,17 @@ if [ "$PROFILE" = nemo_rl ]; then PKG=$(python3 -c 'import nemo_rl, os; print(os.path.dirname(nemo_rl.__file__))') cp -r /tmp/RL/nemo_rl/. "$PKG/" elif [ "$PROFILE" = miles ]; then - python3 -c 'import miles; print("miles pre-installed OK")' + # Overlay the GMI miles fork onto the image's editable install: the base image's + # baked miles predates the Tinker Ray-orchestration interface (forward_backward_only, + # apply_optimizer_step(learning_rate=), apply_optimizer_step_and_sync). Fork main is + # the maintained branch (PR merges fix_hackathon_rebased -> main). + MILES_REPO="${MILES_REPO:-https://github.com/GavinZhu-GMI/miles.git}" + MILES_REF="${MILES_REF:-main}" + MPKG=$(python3 -c 'import miles, os; print(os.path.dirname(miles.__file__))') + rm -rf /tmp/miles_src + git clone -q --depth 1 --branch "$MILES_REF" "$MILES_REPO" /tmp/miles_src + cp -r /tmp/miles_src/miles/. "$MPKG/" + python3 -c 'import miles; print("miles fork overlaid OK ('"$MILES_REF"')")' fi # SDK + cookbook editable installs diff --git a/training/backends/miles/backend.py b/training/backends/miles/backend.py index 23c9ec2..fe2db06 100644 --- a/training/backends/miles/backend.py +++ b/training/backends/miles/backend.py @@ -291,16 +291,17 @@ async def apply_optimizer_step( ) -> Dict[str, Any]: h: MilesHandle = handle # type: ignore[assignment] try: - if learning_rate is not None: - await asyncio.to_thread(h.train_group.set_learning_rate, learning_rate) - + # Miles' RayTrainGroup takes the LR as a parameter (no set_learning_rate); + # apply_optimizer_step_and_sync = apply_optimizer_step + update_weights. offload_train = h.args.offload_train if h.args else True offload_rollout = h.args.offload_rollout if h.args else True if not offload_train and not offload_rollout: - results = await asyncio.to_thread(h.train_group.apply_optimizer_step_and_sync) + results = await asyncio.to_thread( + h.train_group.apply_optimizer_step_and_sync, learning_rate) else: - results = await asyncio.to_thread(h.train_group.apply_optimizer_step) + results = await asyncio.to_thread( + h.train_group.apply_optimizer_step, learning_rate) if h.rollout_manager is not None: from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS From 0901a2358f009b126a2516a759615df8bf58c547 Mon Sep 17 00:00:00 2001 From: "Gavin.Zhu" Date: Mon, 13 Jul 2026 07:10:32 +0000 Subject: [PATCH 2/7] miles: LoRA alpha defaults to rank, not 0 (zero-gradient bug) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Omitted alpha became 0 via lora_config.get("alpha", 0), making LoRA scaling alpha/rank = 0: gradients exactly zero, constant loss — the grad_norm=0 seen in every LoRA round-trip. Schema already documents alpha-defaults-to-rank. Confirmed on ns.config (specs/005 G1 spike): with nonzero alpha, grad_norm=207, loss 11.7->3.2 over 5 steps. Co-Authored-By: Claude Fable 5 --- training/core/slime_builder.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/training/core/slime_builder.py b/training/core/slime_builder.py index 65d6571..8799c17 100644 --- a/training/core/slime_builder.py +++ b/training/core/slime_builder.py @@ -328,9 +328,11 @@ def _configure_model_args( if checkpoint_path: args.load = parse_checkpoint_uri(checkpoint_path, args.save) - # LoRA configuration + # LoRA configuration. alpha defaults to rank per the API schema + # (requests.py LoraConfig); alpha=0 zeroes LoRA scaling and gradients + # entirely — see specs/005-miles-ray-interface/design.md (grad_norm=0). args.lora_rank = lora_config.get("rank", 0) if lora_config else 0 - args.lora_alpha = lora_config.get("alpha", 0) if lora_config else 0 + args.lora_alpha = (lora_config.get("alpha") or args.lora_rank) if lora_config else 0 args.lora_dropout = lora_config.get("dropout", 0.0) if lora_config else 0.0 # Parallelism settings - use values from parallel_config (already auto-detected in build_args) From df0eca7eeba64baecc67f466c27fb2de28d034cc Mon Sep 17 00:00:00 2001 From: "Gavin.Zhu" Date: Mon, 13 Jul 2026 07:18:02 +0000 Subject: [PATCH 3/7] miles profile: overlay via git checkout, not cp cp -r over the image's stale checkout left orphan files and a 73-file dirty git status that masqueraded as unsaved work (see specs/005 HANDOFF 2026-07-13). When the install is a git checkout, fetch+checkout the pinned ref and clean the package dir; else delete-then-copy. Co-Authored-By: Claude Fable 5 --- scripts/lib/setup_container.sh | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/scripts/lib/setup_container.sh b/scripts/lib/setup_container.sh index 05e163c..e2a90e3 100755 --- a/scripts/lib/setup_container.sh +++ b/scripts/lib/setup_container.sh @@ -68,9 +68,19 @@ elif [ "$PROFILE" = miles ]; then MILES_REPO="${MILES_REPO:-https://github.com/GavinZhu-GMI/miles.git}" MILES_REF="${MILES_REF:-main}" MPKG=$(python3 -c 'import miles, os; print(os.path.dirname(miles.__file__))') - rm -rf /tmp/miles_src - git clone -q --depth 1 --branch "$MILES_REF" "$MILES_REPO" /tmp/miles_src - cp -r /tmp/miles_src/miles/. "$MPKG/" + MROOT=$(dirname "$MPKG") + if git -C "$MROOT" rev-parse --git-dir >/dev/null 2>&1; then + # miles is an editable checkout (miles image): update it via git so the + # worktree IS the pinned ref — a bare cp leaves stale files and a dirty + # git status that masquerades as unsaved work. + git -C "$MROOT" fetch -q --depth 1 "$MILES_REPO" "$MILES_REF" + git -C "$MROOT" checkout -qf FETCH_HEAD + git -C "$MROOT" clean -qfd -- miles/ + else + rm -rf /tmp/miles_src "$MPKG" + git clone -q --depth 1 --branch "$MILES_REF" "$MILES_REPO" /tmp/miles_src + cp -r /tmp/miles_src/miles "$MPKG" + fi python3 -c 'import miles; print("miles fork overlaid OK ('"$MILES_REF"')")' fi From 2898626537d848213f4c281083c20e2a33ef1cdc Mon Sep 17 00:00:00 2001 From: "Gavin.Zhu" Date: Mon, 13 Jul 2026 07:43:05 +0000 Subject: [PATCH 4/7] miles: adapt backend to the tinker-seam interface (upstream-based miles) Targets GavinZhu-GMI/miles branch tinker-seam (6e79b54), which re-ports the Tinker orchestration onto current upstream radixark/miles (specs/005). - create_model: reuse upstream factories (create_placement_groups, create_rollout_manager) and the new TinkerTrainGroup async fanout; await group.init(); router address from args.sglang_router_{ip,port}; startup weight sync mirrors upstream train.py's offload dance. - forward_backward: await forward_backward_only; average pp-last metrics across DP ranks. fb returns no per-sample logprobs on the seam (open item). - apply_optimizer_step: await apply_optimizer_step[_and_sync](lr); upstream RolloutManager onload_weights/onload_kv replace tag-based onload. - forward: forward_logprobs returns client-ordered tensors; result assembled inline (rollout_to_forward_result's per-shard juggling is obsolete here). - save/load_checkpoint via group broadcast (load restores optimizer state). - delete_model: _actor_handles rename; placement-group dict cleanup. - converter: send dynamic_global_batch_size=len(batch), _loss_norm_total=1 (pure-sum gradients, G1 contract) and _loss_type_override per request. - builder: use_dynamic_global_batch_size=True, delay_split_train_data_by_dp=True. Not runtime-tested: needs the tinker-seam miles image (upstream stack). Co-Authored-By: Claude Fable 5 --- training/backends/miles/backend.py | 243 ++++++++++++--------------- training/backends/miles/converter.py | 27 ++- training/core/slime_builder.py | 8 +- 3 files changed, 144 insertions(+), 134 deletions(-) diff --git a/training/backends/miles/backend.py b/training/backends/miles/backend.py index fe2db06..d567110 100644 --- a/training/backends/miles/backend.py +++ b/training/backends/miles/backend.py @@ -1,8 +1,13 @@ -"""Miles backend — wraps RayTrainGroup/RolloutManager/SlimeArgumentBuilder -behind the TrainingBackend interface (refactor only, no behavior change).""" +"""Miles backend — wraps TinkerTrainGroup/RolloutManager/SlimeArgumentBuilder +behind the TrainingBackend interface. + +Targets the miles `tinker-seam` branch (upstream-based; specs/005 in +tinker-nemorl): async TinkerTrainGroup fanout, decoupled +forward_backward_only / apply_optimizer_step, pure-sum loss via rollout keys +set in the converter.""" import asyncio import logging -from dataclasses import dataclass, field +from dataclasses import dataclass from datetime import datetime from typing import Any, Dict, List, Optional @@ -97,93 +102,64 @@ async def create_model( ) logger.info("[%s] Miles args built, hf_path=%s", request_id, hf_path) - from miles.ray.actor_group import RayTrainGroup - - num_nodes = 1 - num_gpus_per_node = num_gpus - bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_nodes * num_gpus_per_node)] - pg = ray.util.placement_group(bundles, strategy="PACK") + # Reuse upstream's own wiring (miles tinker-seam branch): placement + # groups + RolloutManager from the factories train.py uses, and the + # TinkerTrainGroup fanout for the decoupled train-step seam. + from miles.ray.placement_group import create_placement_groups, create_rollout_manager + from miles.ray.tinker_group import TinkerTrainGroup - await asyncio.wait_for( - asyncio.wrap_future(pg.ready().future()), - timeout=120.0, - ) + pgs = create_placement_groups(args) - reordered_indices = list(range(len(bundles))) + rollout_manager = None + router_ip = None + router_port = None + if not debug_train_only: + rollout_manager, _ = await asyncio.to_thread( + create_rollout_manager, args, pgs["rollout"] + ) - train_group = RayTrainGroup( + train_group = TinkerTrainGroup( args=args, - num_nodes=num_nodes, - num_gpus_per_node=num_gpus_per_node, - pg=(pg, reordered_indices), - num_gpus_per_actor=0.8, + num_nodes=args.actor_num_nodes, + num_gpus_per_node=args.actor_num_gpus_per_node, + pg=pgs["actor"], + num_gpus_per_actor=0.4, role="actor", + with_ref=False, + rollout_manager=rollout_manager, ) - init_refs = train_group.async_init(args, role="actor", with_ref=False) try: - await asyncio.wait_for( - asyncio.gather(*[asyncio.wrap_future(ref.future()) for ref in init_refs]), - timeout=1800.0, - ) + await asyncio.wait_for(train_group.init(), timeout=1800.0) except asyncio.TimeoutError: - ray.util.remove_placement_group(pg) raise BackendError( "Actor initialization timeout after 1800s", backend="miles", operation="create_model", ) - # Create RolloutManager for SGLang (RL mode only) - rollout_manager = None - router_ip = None - router_port = None - - if not debug_train_only: - from miles.ray.rollout import RolloutManager - - rollout_manager = RolloutManager.options( - num_cpus=1, num_gpus=0, - ).remote(args, (pg, reordered_indices)) - - await asyncio.to_thread(train_group.set_rollout_manager, rollout_manager) - - # Initialize SGLang memory state - from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS, GPU_MEMORY_TYPE_KV_CACHE - try: - from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH - except ImportError: - GPU_MEMORY_TYPE_CUDA_GRAPH = None + if rollout_manager is not None: + await train_group.set_rollout_manager() + # Mirror upstream train.py startup: load weights into SGLang + # before anything samples, honoring rollout offload state. if args.offload_rollout: - await asyncio.to_thread(lambda: ray.get(rollout_manager.offload.remote())) - await asyncio.to_thread(lambda: ray.get( - rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]) - )) - - await asyncio.to_thread(train_group.update_weights) - + await rollout_manager.onload_weights.remote() + await train_group.update_weights() if args.offload_rollout: - if GPU_MEMORY_TYPE_CUDA_GRAPH is not None: - await asyncio.to_thread(lambda: ray.get( - rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH]) - )) - await asyncio.to_thread(lambda: ray.get( - rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE]) - )) - - try: - router_address_ref = rollout_manager.get_router_address.remote() - router_ip, router_port = await asyncio.wrap_future(router_address_ref.future()) - except Exception as e: - logger.error("[%s] Failed to get router address: %s", request_id, e) + await rollout_manager.onload_kv.remote() + + router_ip = getattr(args, "sglang_router_ip", None) + router_port = getattr(args, "sglang_router_port", None) + if not router_ip: + logger.error("[%s] SGLang router address missing from args", request_id) handle = MilesHandle( model_id=model_id, backend_type="miles", train_group=train_group, rollout_manager=rollout_manager, - placement_group=pg, + placement_group=pgs, args=args, hf_path=hf_path, router_ip=router_ip, @@ -215,15 +191,23 @@ async def forward( from miles.utils.ray_utils import Box if h.rollout_manager is not None and h.args.offload_rollout: - await asyncio.to_thread(lambda: ray.get(h.rollout_manager.offload.remote())) + await h.rollout_manager.offload.remote() rollout_data = self.converter.forward_to_backend(data, h.args) - results = await asyncio.to_thread( - h.train_group.forward_only, - rollout_id=0, - rollout_data_ref=Box(ray.put(rollout_data)), - ) - return self.converter.backend_to_forward_result(results, data) + # TinkerTrainGroup returns per-sample logprob tensors already + # merged into the client's submission order. + logprobs = await h.train_group.forward_logprobs(0, Box(ray.put(rollout_data))) + + loss_fn_outputs = [ + {"logprobs": {"data": lp.tolist(), "shape": [len(lp)], "dtype": "float32"}} + for lp in logprobs + ] + return { + "type": "forward", + "loss_fn_output_type": loss_fn, + "loss_fn_outputs": loss_fn_outputs, + "metrics": {}, + } except BackendError: raise @@ -243,7 +227,7 @@ async def forward_backward( from miles.utils.ray_utils import Box if h.rollout_manager is not None and h.args.offload_rollout: - await asyncio.to_thread(lambda: ray.get(h.rollout_manager.offload.remote())) + await h.rollout_manager.offload.remote() is_rl = not h.args.debug_train_only @@ -266,14 +250,28 @@ async def forward_backward( data, loss_fn, h.args, ) - results = await asyncio.to_thread( - h.train_group.forward_backward_only, - rollout_id=0, - rollout_data_ref=Box(ray.put(rollout_data)), - ) - result = self.converter.backend_to_forward_backward_result(results, data) - result["deferred"] = False - return result + results = await h.train_group.forward_backward_only(0, Box(ray.put(rollout_data))) + + # Only pipeline-last-stage actors return metrics; average across + # the DP ranks that did. Per-sample logprobs are not emitted by the + # seam's fb pass (specs/005 HANDOFF, open item). + summed: Dict[str, float] = {} + reporting = 0 + for r in results or []: + loss_dict = (r or {}).get("loss") or {} + if loss_dict: + reporting += 1 + for k, v in loss_dict.items(): + summed[k] = summed.get(k, 0.0) + float(v) + metrics = {k: v / reporting for k, v in summed.items()} if reporting else {} + + return { + "loss_fn_output_type": loss_fn, + "loss": metrics.get("loss"), + "metrics": metrics, + "loss_fn_outputs": [], + "deferred": False, + } except BackendError: raise @@ -291,43 +289,26 @@ async def apply_optimizer_step( ) -> Dict[str, Any]: h: MilesHandle = handle # type: ignore[assignment] try: - # Miles' RayTrainGroup takes the LR as a parameter (no set_learning_rate); - # apply_optimizer_step_and_sync = apply_optimizer_step + update_weights. + # TinkerTrainGroup: apply_optimizer_step(learning_rate) fans out to + # the actors; _and_sync additionally pushes weights to SGLang. offload_train = h.args.offload_train if h.args else True offload_rollout = h.args.offload_rollout if h.args else True - if not offload_train and not offload_rollout: - results = await asyncio.to_thread( - h.train_group.apply_optimizer_step_and_sync, learning_rate) + if h.rollout_manager is None: + results = await h.train_group.apply_optimizer_step(learning_rate) + elif not offload_train and not offload_rollout: + results = await h.train_group.apply_optimizer_step_and_sync(learning_rate) else: - results = await asyncio.to_thread( - h.train_group.apply_optimizer_step, learning_rate) - - if h.rollout_manager is not None: - from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS - try: - from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH - except ImportError: - GPU_MEMORY_TYPE_CUDA_GRAPH = None - - if offload_train: - await asyncio.to_thread(h.train_group.offload) - - if offload_rollout: - await asyncio.to_thread(lambda: ray.get( - h.rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]) - )) - - await asyncio.to_thread(h.train_group.update_weights) - - if offload_rollout: - if GPU_MEMORY_TYPE_CUDA_GRAPH is not None: - await asyncio.to_thread(lambda: ray.get( - h.rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH]) - )) - await asyncio.to_thread(lambda: ray.get( - h.rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE]) - )) + results = await h.train_group.apply_optimizer_step(learning_rate) + + # Mirror upstream train.py's offload dance around weight sync. + if offload_train: + await h.train_group.offload() + if offload_rollout: + await h.rollout_manager.onload_weights.remote() + await h.train_group.update_weights() + if offload_rollout: + await h.rollout_manager.onload_kv.remote() return { "success": results[0]["success"], @@ -346,7 +327,7 @@ async def apply_optimizer_step( async def update_inference_weights(self, handle: BackendHandle) -> None: h: MilesHandle = handle # type: ignore[assignment] try: - await asyncio.to_thread(h.train_group.update_weights) + await h.train_group.update_weights() except Exception as e: raise BackendError( str(e), backend="miles", operation="update_inference_weights", original_error=e, @@ -368,11 +349,7 @@ async def save_checkpoint( if step_id is None: step_id = 0 - object_refs = [ - actor.save_model.remote(step_id) - for actor in h.train_group._actor_handlers - ] - await asyncio.gather(*[asyncio.wrap_future(ref.future()) for ref in object_refs]) + await h.train_group.save_model(step_id) return checkpoint_path except Exception as e: @@ -387,15 +364,11 @@ async def load_checkpoint( ) -> None: h: MilesHandle = handle # type: ignore[assignment] try: - object_refs = [ - actor.load_checkpoint.remote(checkpoint_path) - for actor in h.train_group._actor_handlers - ] - await asyncio.gather(*[asyncio.wrap_future(ref.future()) for ref in object_refs]) + await h.train_group.load_checkpoint(checkpoint_path) # Sync loaded weights to inference engine if h.rollout_manager is not None: - await asyncio.to_thread(h.train_group.update_weights) + await h.train_group.update_weights() logger.info("Miles checkpoint loaded from %s", checkpoint_path) @@ -408,7 +381,7 @@ async def delete_model(self, handle: BackendHandle) -> None: h: MilesHandle = handle # type: ignore[assignment] try: resources_freed = [] - for actor in h.train_group._actor_handlers: + for actor in h.train_group._actor_handles: ray.kill(actor, no_restart=True) resources_freed.append("actor") @@ -416,9 +389,17 @@ async def delete_model(self, handle: BackendHandle) -> None: ray.kill(h.rollout_manager, no_restart=True) resources_freed.append("rollout_manager") - if h.placement_group is not None: - ray.util.remove_placement_group(h.placement_group) - resources_freed.append("placement_group") + # placement_group holds the create_placement_groups() dict of + # (pg, bundle_indices, gpu_ids) tuples; pgs may be shared between + # roles (colocate), so dedupe before removal. + if h.placement_group: + seen = set() + for pg_tuple in h.placement_group.values(): + pg_obj = pg_tuple[0] if isinstance(pg_tuple, tuple) else pg_tuple + if id(pg_obj) not in seen: + seen.add(id(pg_obj)) + ray.util.remove_placement_group(pg_obj) + resources_freed.append("placement_group") logger.info("Miles model %s deleted, freed %d resources", h.model_id, len(resources_freed)) diff --git a/training/backends/miles/converter.py b/training/backends/miles/converter.py index cc6bd78..015f789 100644 --- a/training/backends/miles/converter.py +++ b/training/backends/miles/converter.py @@ -23,7 +23,9 @@ def forward_to_backend( args: Any, ) -> Any: """Convert Tinker data to Miles rollout_data for forward pass.""" - return self._inner.forward_to_rollout(data) + rollout_data = self._inner.forward_to_rollout(data) + self._add_tinker_seam_keys(rollout_data, len(data)) + return rollout_data def forward_backward_to_backend( self, @@ -33,7 +35,28 @@ def forward_backward_to_backend( ) -> Any: """Convert Tinker data to Miles rollout_data for training.""" is_rl = not getattr(args, "debug_train_only", False) - return self._inner.forward_backward_to_rollout(data, is_rl=is_rl) + rollout_data = self._inner.forward_backward_to_rollout(data, is_rl=is_rl) + self._add_tinker_seam_keys(rollout_data, len(data)) + # Per-request loss selection (upstream dispatches on args.loss_type at + # startup; the seam overrides per batch). The inner converter already + # sets sft_loss when it detects SFT-shaped data — don't override that. + rollout_data.setdefault( + "_loss_type_override", + "sft_loss" if loss_fn == "cross_entropy" else "policy_loss", + ) + return rollout_data + + @staticmethod + def _add_tinker_seam_keys(rollout_data: Any, num_samples: int) -> None: + """Keys the tinker-seam miles branch consumes (specs/005 design.md). + + - dynamic_global_batch_size: actual request size, so upstream + get_data_iterator schedules correctly for variable batches. + - _loss_norm_total=1: pure-sum gradients — invariant to how a logical + batch is split across forward_backward calls (G1 contract). + """ + rollout_data["dynamic_global_batch_size"] = num_samples + rollout_data["_loss_norm_total"] = 1 def backend_to_forward_result( self, diff --git a/training/core/slime_builder.py b/training/core/slime_builder.py index 8799c17..fcb556c 100644 --- a/training/core/slime_builder.py +++ b/training/core/slime_builder.py @@ -12,7 +12,6 @@ from ..utils.model_config import ( load_model_config, - get_parallelism_config, auto_detect_all_parallelism, detect_torch_dist_path, parse_checkpoint_uri, @@ -479,4 +478,11 @@ def _configure_model_args( args.offload_rollout = False args.train_memory_margin_bytes = 0 + # Tinker seam (miles tinker-seam branch, specs/005 design.md): + # per-request batch sizes ride the dynamic_global_batch_size rollout key + # (its assert couples this arg to the key's presence), and each actor + # splits the fanned-out batch by DP rank locally. + args.use_dynamic_global_batch_size = True + args.delay_split_train_data_by_dp = True + return args From 86ca0332f9f95f64727e860f0730b05ae9d0f2b5 Mon Sep 17 00:00:00 2001 From: "Gavin.Zhu" Date: Mon, 13 Jul 2026 09:17:27 +0000 Subject: [PATCH 5/7] miles seam: first-boot fixes from gate validation on ns.config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five runtime-drift fixes found running G1/G2 gates on the tinker-seam stack (miles_tinker_seam:20260713 image): - model_config: transformers 5.x moved rope_theta into rope_parameters; the 10000 fallback tripped upstream's hf_validate_args. - slime_builder: reset disable_{grad,param}_buffers_cpu_backup after forcing offload_train=False — parse-time --offload derived them True, and actors crashed allocating grad buffers in a TMS region without LD_PRELOAD. - data_converter (forward path): response_length is the mask LENGTH (not its nonzero count) + causal N-1 trim when response==total; fixes shape assert and empty logprob slices. - backend: run create_placement_groups / group construction off the event loop (retrieve_future polls stalled); skip None rollout pg on delete. - deploy script: forward MILES_REPO/MILES_REF into the container setup (the overlay silently checked out main and reverted the baked seam). Gate results (Qwen2.5-0.5B LoRA r32/a64, DP=4): G1 ratio EXACTLY 1.0 (grad_norm 1485.871... bit-identical for 1x8 vs 4+4 fb split — pure-sum contract holds); G2 client-order exact; weights move; loss 11.78->5.11 over 5 steps; delete_model clean. Co-Authored-By: Claude Fable 5 --- scripts/deploy_tinkercloud.sh | 2 +- training/backends/miles/backend.py | 11 +++++++---- training/core/data_converter.py | 12 ++++++++---- training/core/slime_builder.py | 7 ++++++- training/utils/model_config.py | 7 ++++++- 5 files changed, 28 insertions(+), 11 deletions(-) diff --git a/scripts/deploy_tinkercloud.sh b/scripts/deploy_tinkercloud.sh index d762b6b..65c5039 100755 --- a/scripts/deploy_tinkercloud.sh +++ b/scripts/deploy_tinkercloud.sh @@ -276,7 +276,7 @@ fi # --- 4. in-container setup --------------------------------------------------- echo "==> [4/6] running in-container setup (profile=$PROFILE)" -EX "PROFILE=$PROFILE bash /tmp/setup_container.sh" +EX "PROFILE=$PROFILE MILES_REPO='${MILES_REPO:-}' MILES_REF='${MILES_REF:-}' bash /tmp/setup_container.sh" if [ "$RUN_SERVER" = 0 ]; then # bionemo profile: env only — no Ray, no server (no working evo2 server yet). diff --git a/training/backends/miles/backend.py b/training/backends/miles/backend.py index d567110..0c27739 100644 --- a/training/backends/miles/backend.py +++ b/training/backends/miles/backend.py @@ -108,7 +108,9 @@ async def create_model( from miles.ray.placement_group import create_placement_groups, create_rollout_manager from miles.ray.tinker_group import TinkerTrainGroup - pgs = create_placement_groups(args) + # Sync ray calls (pg.ready waits, actor allocation) — keep them off + # the event loop or /retrieve_future polls stall and clients time out. + pgs = await asyncio.to_thread(create_placement_groups, args) rollout_manager = None router_ip = None @@ -118,7 +120,7 @@ async def create_model( create_rollout_manager, args, pgs["rollout"] ) - train_group = TinkerTrainGroup( + train_group = await asyncio.to_thread(lambda: TinkerTrainGroup( args=args, num_nodes=args.actor_num_nodes, num_gpus_per_node=args.actor_num_gpus_per_node, @@ -127,7 +129,7 @@ async def create_model( role="actor", with_ref=False, rollout_manager=rollout_manager, - ) + )) try: await asyncio.wait_for(train_group.init(), timeout=1800.0) @@ -396,7 +398,8 @@ async def delete_model(self, handle: BackendHandle) -> None: seen = set() for pg_tuple in h.placement_group.values(): pg_obj = pg_tuple[0] if isinstance(pg_tuple, tuple) else pg_tuple - if id(pg_obj) not in seen: + # debug_train_only leaves the rollout entry as None + if pg_obj is not None and id(pg_obj) not in seen: seen.add(id(pg_obj)) ray.util.remove_placement_group(pg_obj) resources_freed.append("placement_group") diff --git a/training/core/data_converter.py b/training/core/data_converter.py index 228879e..9bba2f9 100644 --- a/training/core/data_converter.py +++ b/training/core/data_converter.py @@ -112,11 +112,15 @@ def forward_to_rollout(cls, data: List[Any]) -> Dict[str, Any]: # Default: all ones (no masking) loss_mask = torch.ones(len(tokens), dtype=torch.float32) + # Miles convention: the mask covers the response REGION, so + # response_length is the mask's length (zeros inside are allowed); + # counting nonzeros breaks prompt/response alignment in get_batch. + # Causal N-1 trim: response == total yields an empty logit slice + # (same handling as the forward_backward RL path). + if len(loss_mask) == len(tokens) and len(tokens) > 1: + loss_mask = loss_mask[1:] loss_masks_list.append(loss_mask) - - # Response length is number of non-zero mask elements - response_length = int(loss_mask.sum().item()) - response_lengths_list.append(response_length) + response_lengths_list.append(len(loss_mask)) # print(f"[CONVERTER DEBUG SFT] Sample {len(loss_masks_list)-1}: loss_mask sum={response_length}, len={len(loss_mask)}", flush=True) # Build rollout_data with dummy RL fields (not used for forward-only) diff --git a/training/core/slime_builder.py b/training/core/slime_builder.py index fcb556c..d394790 100644 --- a/training/core/slime_builder.py +++ b/training/core/slime_builder.py @@ -473,10 +473,15 @@ def _configure_model_args( } # Disable offload_train for simpler GPU memory management - # When offload_train=False, must also set train_memory_margin_bytes=0 to avoid assert + # When offload_train=False, must also set train_memory_margin_bytes=0 to avoid assert. + # Also reset flags parse_args DERIVED from the minimal-args --offload flag, + # or actors allocate grad buffers in a torch_memory_saver region without + # LD_PRELOAD (only set for offload actors) and crash at init. args.offload_train = False args.offload_rollout = False args.train_memory_margin_bytes = 0 + args.disable_grad_buffers_cpu_backup = False + args.disable_param_buffers_cpu_backup = False # Tinker seam (miles tinker-seam branch, specs/005 design.md): # per-request batch sizes ride the dynamic_global_batch_size rollout key diff --git a/training/utils/model_config.py b/training/utils/model_config.py index b8751e3..97db5e6 100644 --- a/training/utils/model_config.py +++ b/training/utils/model_config.py @@ -107,7 +107,12 @@ def load_model_config(base_model: str) -> Dict[str, Any]: config, 'rms_norm_eps', getattr(config, 'layer_norm_eps', 1e-6) ), - 'rotary_base': getattr(config, 'rope_theta', 10000), + # transformers 5.x moved rope_theta into the rope_parameters dict + 'rotary_base': ( + getattr(config, 'rope_theta', None) + or (getattr(config, 'rope_parameters', None) or {}).get('rope_theta') + or 10000 + ), 'tie_word_embeddings': getattr(config, 'tie_word_embeddings', False), 'max_position_embeddings': getattr(config, 'max_position_embeddings', 2048), } From 8438606e075a0daa1d386acab8002fb490ea58ec Mon Sep 17 00:00:00 2001 From: "Gavin.Zhu" Date: Mon, 13 Jul 2026 12:59:03 +0000 Subject: [PATCH 6/7] =?UTF-8?q?miles=20seam:=20G3=20fixes=20=E2=80=94=20fu?= =?UTF-8?q?ll=20RL-mode=20stack=20+=20cookbook=20SFT=20run=20green?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes found running sl_basic (NoRobots, 30 steps) through the full stack (RolloutManager + SGLang engines + bridge-LoRA), each one layer deeper: - converter: concatenate ALL model_input chunks (chunks[0]-only silently truncated multi-chunk datums: 4-token tokens vs 438-token weights, which also masqueraded as a dynamo compile failure in fused cross-entropy) - builder: rollout_global_dataset=False (Tinker clients drive data; upstream RolloutManager otherwise requires a server-side prompt dataset file) - builder: target_modules from Tinker train_attn/train_mlp flags (upstream requires them for LoRA), megatron_to_hf_mode=bridge (raw mode silently builds full-finetune: no lora_A/B params, weight sync fails) - backend: fb metrics keys as name:mean (SDK reduction convention); fb returns per-datum logprobs in client order (SDK weights reductions by len(loss_fn_outputs); cookbook computes NLL from them; miles tinker-seam b1c2c0f computes them actor-side on pre-step weights) G3 result: 30/30 steps, per-request linear LR schedule honored, NLL stable 2.7->2.4 on an already-instruction-tuned base, checkpoint saved, clean exit. KNOWN ISSUE: multi-fb grad accumulation on bridge-LoRA shows ratio ~1/sqrt(2) (second fb appears to overwrite adapter grads) — single-fb-per-step flows (the cookbook default) are unaffected; tracked in specs/005 HANDOFF. Co-Authored-By: Claude Fable 5 --- training/backends/miles/backend.py | 19 ++++++++++++++++--- training/core/data_converter.py | 19 ++++++++++++++----- training/core/slime_builder.py | 21 +++++++++++++++++++-- 3 files changed, 49 insertions(+), 10 deletions(-) diff --git a/training/backends/miles/backend.py b/training/backends/miles/backend.py index 0c27739..a950e56 100644 --- a/training/backends/miles/backend.py +++ b/training/backends/miles/backend.py @@ -265,13 +265,26 @@ async def forward_backward( reporting += 1 for k, v in loss_dict.items(): summed[k] = summed.get(k, 0.0) + float(v) - metrics = {k: v / reporting for k, v in summed.items()} if reporting else {} + averaged = {k: v / reporting for k, v in summed.items()} if reporting else {} + # SDK metric keys carry their cross-chunk reduction as ":" + # (chunked_fwdbwd_helpers._metrics_reduction splits on ":"). + metrics = {f"{k}:mean": v for k, v in averaged.items()} + + # Per-datum response logprobs in client order (the SDK weights its + # metric reduction by len(loss_fn_outputs), and the cookbook + # computes NLL from these). + from miles.ray.tinker_group import merge_dp_sample_outputs + logprobs_list = merge_dp_sample_outputs(results or [], key="log_probs") + loss_fn_outputs = [ + {"logprobs": {"data": lp.tolist(), "shape": [len(lp)], "dtype": "float32"}} + for lp in logprobs_list + ] return { "loss_fn_output_type": loss_fn, - "loss": metrics.get("loss"), + "loss": averaged.get("loss"), "metrics": metrics, - "loss_fn_outputs": [], + "loss_fn_outputs": loss_fn_outputs, "deferred": False, } diff --git a/training/core/data_converter.py b/training/core/data_converter.py index 9bba2f9..a3e0c0a 100644 --- a/training/core/data_converter.py +++ b/training/core/data_converter.py @@ -40,13 +40,22 @@ def extract_tokens_from_model_input(model_input: Any) -> List[int]: Works with both dict and Pydantic model inputs. """ - # Try chunks first + # Try chunks first. A ModelInput may carry MULTIPLE chunks (the SDK + # splits inputs); concatenate them all — taking only chunks[0] + # silently truncates the sample (specs/005: 4-token tokens vs + # 438-token weights crash in the miles loss). chunks = TinkerDataConverter._get_field(model_input, "chunks") if chunks: - if not chunks: - raise ValueError("Empty chunks in model_input") - first_chunk = chunks[0] - return TinkerDataConverter._get_field(first_chunk, "tokens") + tokens: List[int] = [] + for chunk in chunks: + chunk_tokens = TinkerDataConverter._get_field(chunk, "tokens") + if chunk_tokens is None: + raise ValueError( + "model_input chunk without tokens (non-text chunks are " + "not supported by the miles backend)" + ) + tokens.extend(chunk_tokens) + return tokens # Try tokens tokens = TinkerDataConverter._get_field(model_input, "tokens") diff --git a/training/core/slime_builder.py b/training/core/slime_builder.py index d394790..bd97041 100644 --- a/training/core/slime_builder.py +++ b/training/core/slime_builder.py @@ -333,6 +333,21 @@ def _configure_model_args( args.lora_rank = lora_config.get("rank", 0) if lora_config else 0 args.lora_alpha = (lora_config.get("alpha") or args.lora_rank) if lora_config else 0 args.lora_dropout = lora_config.get("dropout", 0.0) if lora_config else 0.0 + # Megatron module targets from Tinker's train_attn/train_mlp flags + # (upstream requires target_modules when LoRA is enabled; unembed LoRA + # is not supported by miles' injector). + target_modules = [] + if not lora_config or lora_config.get("train_attn", True): + target_modules += ["linear_qkv", "linear_proj"] + if not lora_config or lora_config.get("train_mlp", True): + target_modules += ["linear_fc1", "linear_fc2"] + args.target_modules = target_modules + # Upstream only injects megatron-side LoRA adapters on the bridge path + # (model.py: is_lora_enabled and megatron_to_hf_mode == "bridge"); + # without this the model silently builds as full-finetune and LoRA + # weight sync to SGLang fails (no lora_A/lora_B params). + if args.lora_rank > 0: + args.megatron_to_hf_mode = "bridge" # Parallelism settings - use values from parallel_config (already auto-detected in build_args) tp_size = parallel_config.get('tensor_parallel_size', 2) @@ -420,8 +435,10 @@ def _configure_model_args( args.rollout_function_path = "miles.rollout.sglang_rollout.generate_rollout" args.eval_function_path = "miles.rollout.sglang_rollout.generate_rollout" - # Dataset configuration (fallback for testing) - args.rollout_global_dataset = True + # No server-side prompt dataset: Tinker clients drive all training and + # sampling data, so RolloutManager must not load one (its data source + # raises if the file is missing). + args.rollout_global_dataset = False args.prompt_data = "/data/datasets/gsm8k_rl.jsonl" args.rollout_shuffle = False args.rollout_max_prompt_len = 2048 From 562bad4e69472132b15da916057ff225e4931685 Mon Sep 17 00:00:00 2001 From: "Gavin.Zhu" Date: Tue, 14 Jul 2026 00:27:34 +0000 Subject: [PATCH 7/7] miles: self-contained comments (drop monorepo spec references) Co-Authored-By: Claude Fable 5 --- training/backends/miles/backend.py | 6 +++--- training/backends/miles/converter.py | 2 +- training/core/data_converter.py | 4 ++-- training/core/slime_builder.py | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/training/backends/miles/backend.py b/training/backends/miles/backend.py index a950e56..df07e3b 100644 --- a/training/backends/miles/backend.py +++ b/training/backends/miles/backend.py @@ -1,8 +1,8 @@ """Miles backend — wraps TinkerTrainGroup/RolloutManager/SlimeArgumentBuilder behind the TrainingBackend interface. -Targets the miles `tinker-seam` branch (upstream-based; specs/005 in -tinker-nemorl): async TinkerTrainGroup fanout, decoupled +Targets the miles `tinker-seam` branch (upstream-based): async +TinkerTrainGroup fanout, decoupled forward_backward_only / apply_optimizer_step, pure-sum loss via rollout keys set in the converter.""" import asyncio @@ -256,7 +256,7 @@ async def forward_backward( # Only pipeline-last-stage actors return metrics; average across # the DP ranks that did. Per-sample logprobs are not emitted by the - # seam's fb pass (specs/005 HANDOFF, open item). + # seam's fb pass itself (they ride a separate forward). summed: Dict[str, float] = {} reporting = 0 for r in results or []: diff --git a/training/backends/miles/converter.py b/training/backends/miles/converter.py index 015f789..76269e6 100644 --- a/training/backends/miles/converter.py +++ b/training/backends/miles/converter.py @@ -48,7 +48,7 @@ def forward_backward_to_backend( @staticmethod def _add_tinker_seam_keys(rollout_data: Any, num_samples: int) -> None: - """Keys the tinker-seam miles branch consumes (specs/005 design.md). + """Keys the tinker-seam miles branch consumes. - dynamic_global_batch_size: actual request size, so upstream get_data_iterator schedules correctly for variable batches. diff --git a/training/core/data_converter.py b/training/core/data_converter.py index a3e0c0a..273572c 100644 --- a/training/core/data_converter.py +++ b/training/core/data_converter.py @@ -42,8 +42,8 @@ def extract_tokens_from_model_input(model_input: Any) -> List[int]: """ # Try chunks first. A ModelInput may carry MULTIPLE chunks (the SDK # splits inputs); concatenate them all — taking only chunks[0] - # silently truncates the sample (specs/005: 4-token tokens vs - # 438-token weights crash in the miles loss). + # silently truncates the sample (e.g. 4-token tokens vs 438-token + # weights, crashing the miles loss on shape mismatch). chunks = TinkerDataConverter._get_field(model_input, "chunks") if chunks: tokens: List[int] = [] diff --git a/training/core/slime_builder.py b/training/core/slime_builder.py index bd97041..fab06ef 100644 --- a/training/core/slime_builder.py +++ b/training/core/slime_builder.py @@ -329,7 +329,7 @@ def _configure_model_args( # LoRA configuration. alpha defaults to rank per the API schema # (requests.py LoraConfig); alpha=0 zeroes LoRA scaling and gradients - # entirely — see specs/005-miles-ray-interface/design.md (grad_norm=0). + # entirely (the historical grad_norm=0 constant-loss bug). args.lora_rank = lora_config.get("rank", 0) if lora_config else 0 args.lora_alpha = (lora_config.get("alpha") or args.lora_rank) if lora_config else 0 args.lora_dropout = lora_config.get("dropout", 0.0) if lora_config else 0.0 @@ -500,7 +500,7 @@ def _configure_model_args( args.disable_grad_buffers_cpu_backup = False args.disable_param_buffers_cpu_backup = False - # Tinker seam (miles tinker-seam branch, specs/005 design.md): + # Tinker seam (miles tinker-seam branch): # per-request batch sizes ride the dynamic_global_batch_size rollout key # (its assert couples this arg to the key's presence), and each actor # splits the fanned-out batch by DP rank locally.