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),