From 25eb11e29efab26a0ddec02ea52028da64c5e4a9 Mon Sep 17 00:00:00 2001 From: Jason Mancuso <7891333+jvmncs@users.noreply.github.com> Date: Mon, 22 Jun 2026 23:51:35 -0700 Subject: [PATCH] fix(disagg): NVFP4 + publish-only rollout fixes for disaggregated training Found bringing up disaggregated NVFP4 rollout (cookbook/miles_disagg, Moonlight/Kimi-K2.6) on Modal: - megatron_to_hf/processors: route quant_algo=="NVFP4" to quantize_params_nvfp4 (modelopt NVFP4 checkpoints advertise quant_method="modelopt", so dispatch on quant_algo); NVFP4 export was never reached before. - rollout/sglang_rollout: GenerateState.semaphore was Semaphore(0) in publish-only mode (rollout_num_gpus==0) -> every rollout deadlocked. Bound generation concurrency by sglang_server_concurrency when rollout_endpoint_url is set. - utils/http_utils: init_http_client early-returned when rollout_num_gpus==0, leaving _http_client=None -> "'NoneType' has no attribute 'post'". Initialize it and bound _client_concurrency by sglang_server_concurrency in publish-only. - update_weight/update_weight_from_disk_delta: flatten before .view(torch.uint8) (.contiguous().reshape(-1).view) so 0-dim NVFP4 scalar tensors (weight_scale_2, input_scale) don't crash the disk-delta encode/snapshot. - chat_template_utils/deepseek_v4: guard the `encoding_dsv4` import so non-V4 models load on sglang-miles builds without that symbol. --- .../megatron_to_hf/processors/__init__.py | 6 ++++++ .../update_weight/update_weight_from_disk_delta.py | 9 +++++++-- miles/rollout/sglang_rollout.py | 14 +++++++++++--- miles/utils/chat_template_utils/deepseek_v4.py | 5 ++++- miles/utils/http_utils.py | 12 ++++++++++-- 5 files changed, 38 insertions(+), 8 deletions(-) diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py b/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py index fba2c4e165..ec61504d6d 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py @@ -2,6 +2,7 @@ from .quantizer_compressed_tensors import quantize_params_compressed_tensors from .quantizer_fp8 import quantize_params_fp8 from .quantizer_mxfp8 import quantize_params_mxfp8 +from .quantizer_nvfp4 import quantize_params_nvfp4 __all__ = [ "remove_padding", @@ -9,6 +10,7 @@ "quantize_params_fp8", "quantize_params_mxfp8", "quantize_params_compressed_tensors", + "quantize_params_nvfp4", ] @@ -22,3 +24,7 @@ def quantize_params(args, megatron_name, converted_named_params, quantization_co elif quantization_config["quant_method"] == "compressed-tensors": # only int4 at the moment. return quantize_params_compressed_tensors(converted_named_params, quantization_config) + elif quantization_config.get("quant_algo") == "NVFP4" or quantization_config.get("quant_method") == "nvfp4": + # modelopt NVFP4 checkpoints (e.g. nvidia/Kimi-K2.6-NVFP4) advertise + # quant_method="modelopt", so route on quant_algo, not quant_method. + return quantize_params_nvfp4(args, megatron_name, converted_named_params, quantization_config) diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py b/miles/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py index c3f584af6d..f21a64bedf 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py @@ -143,7 +143,8 @@ def _snapshot_bucket(self, converted_named_tensors: list[tuple[str, torch.Tensor try: self._snapshot[name] = self._baseline_reader(name) except KeyError: - self._snapshot[name] = tensor.detach().cpu().contiguous().view(torch.uint8).numpy().reshape(-1) + # flatten before the byte-view (0-dim scalars like NVFP4 scales) + self._snapshot[name] = tensor.detach().cpu().contiguous().reshape(-1).view(torch.uint8).numpy() logger.warning("seed: %s absent from hf_checkpoint; seeding from current weights", name) converted_named_tensors.clear() @@ -196,7 +197,11 @@ def _encode_bucket(self, converted_named_tensors: list[tuple[str, torch.Tensor]] """Copy each gathered HF tensor to host and submit it to the diff/compress pool, draining once enough work is in flight to backpressure the gather (source ranks only).""" for name, tensor in converted_named_tensors: - flat = tensor.detach().contiguous().view(torch.uint8).reshape(-1) + # Flatten to 1-D *before* the byte-view: NVFP4 export emits 0-dim + # scalar tensors (weight_scale_2, input_scale), and .view(torch.uint8) + # cannot reinterpret a 0-dim tensor ("self.dim() cannot be 0 to view + # Float as Byte"). reshape(-1) on a 0-dim tensor yields [1] first. + flat = tensor.detach().contiguous().reshape(-1).view(torch.uint8) nbytes = int(flat.numel()) if self._use_pinned and nbytes <= self._max_bytes: buf = self._free_q.get() # blocks when all buffers are in flight -> backpressures the gather diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py index f49d9b1f61..eb7e573dc2 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -107,9 +107,17 @@ def __init__(self, args: Namespace) -> None: ) self.processor = load_processor(args.hf_checkpoint, trust_remote_code=True) - self.semaphore = asyncio.Semaphore( - args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine - ) + if getattr(args, "rollout_endpoint_url", None): + # Publish-only / disaggregated: no miles-launched engines, so + # rollout_num_gpus == 0 and the GPU-based formula below would be 0, + # making asyncio.Semaphore(0) deadlock every rollout. Bound generation + # concurrency by the external pool's per-engine concurrency instead. + generation_concurrency = args.sglang_server_concurrency + else: + generation_concurrency = ( + args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + ) + self.semaphore = asyncio.Semaphore(generation_concurrency) self.sampling_params: dict[str, Any] = dict( temperature=args.rollout_temperature, top_p=args.rollout_top_p, diff --git a/miles/utils/chat_template_utils/deepseek_v4.py b/miles/utils/chat_template_utils/deepseek_v4.py index b3a70513fe..49fff42119 100644 --- a/miles/utils/chat_template_utils/deepseek_v4.py +++ b/miles/utils/chat_template_utils/deepseek_v4.py @@ -7,7 +7,10 @@ import os from typing import Any -from sglang.srt.entrypoints.openai import encoding_dsv4 +try: + from sglang.srt.entrypoints.openai import encoding_dsv4 +except ImportError: # older sglang-miles without the DeepSeek-V4 encoder; fine for non-V4 models + encoding_dsv4 = None from sglang.srt.entrypoints.openai.protocol import Tool logger = logging.getLogger(__name__) diff --git a/miles/utils/http_utils.py b/miles/utils/http_utils.py index 4e7cb61efd..2556b159a7 100644 --- a/miles/utils/http_utils.py +++ b/miles/utils/http_utils.py @@ -227,10 +227,18 @@ async def _post(client, url, payload, max_retries=60, action="post", headers=Non def init_http_client(args): """Initialize HTTP client and optionally enable distributed POST via Ray.""" global _http_client, _client_concurrency, _distributed_post_enabled - if not args.rollout_num_gpus: + publish_only = bool(getattr(args, "rollout_endpoint_url", None)) + if not args.rollout_num_gpus and not publish_only: return - _client_concurrency = args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + if publish_only: + # Disaggregated / publish-only: no miles-launched engines, so + # rollout_num_gpus == 0 and the GPU-based formula below would be 0 — the + # client would never be created (NoneType.post). The external pool bounds + # concurrency instead. + _client_concurrency = args.sglang_server_concurrency + else: + _client_concurrency = args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine if _http_client is None: _http_client = httpx.AsyncClient( limits=httpx.Limits(max_connections=_client_concurrency),