diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index b09c2c3a03b..0475fb5cf59 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -28,6 +28,11 @@ from ...utils.profile_utils import TrainProfiler from ...utils.tensor_backper import TensorBackuper from ..training_utils.data import DataIterator, get_data_iterator, get_rollout_data, sync_actor_critic_data +from ..training_utils.higgs_policy import ( + is_higgs_policy_enabled, + validate_higgs_logprob_parity, + validate_higgs_weight_versions, +) from ..training_utils.log_utils import log_cpu_memory, log_perf_data, log_rollout_data from ..training_utils.loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values from ..training_utils.parallel import get_parallel_state @@ -39,9 +44,6 @@ from .parallel import verify_megatron_parallel_state from .replay_utils import register_replay_list_moe from .update_weight.common import named_params_and_buffers -from .update_weight.update_weight_from_distributed.broadcast import UpdateWeightFromDistributed -from .update_weight.update_weight_from_distributed.p2p import UpdateWeightP2P -from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor if TYPE_CHECKING: from miles.ray.rollout.rollout_manager import EnginesAndLock @@ -173,9 +175,13 @@ def init( self.args.vocab_size = self.tokenizer.vocab_size if self.args.colocate: + from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor + update_weight_cls = UpdateWeightFromTensor else: if self.args.update_weight_transfer_mode == "broadcast": + from .update_weight.update_weight_from_distributed.broadcast import UpdateWeightFromDistributed + update_weight_cls = UpdateWeightFromDistributed elif self.args.update_weight_transfer_mode == "disk-delta": # Lazy import: keeps the delta deps (numpy/zstandard/xxhash) off the other paths. @@ -183,6 +189,8 @@ def init( update_weight_cls = UpdateWeightFromDiskDelta else: + from .update_weight.update_weight_from_distributed.p2p import UpdateWeightP2P + update_weight_cls = UpdateWeightP2P self.weight_updater = update_weight_cls( self.args, @@ -322,6 +330,12 @@ def _use_rollout_replay(self, m) -> bool: return getattr(self.args, f"use_rollout_{m.name}_replay", False) def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: + higgs_policy = is_higgs_policy_enabled(self.args) + if higgs_policy: + validate_higgs_weight_versions( + rollout_data.get("weight_versions"), + trainer_weight_version=self.weight_updater.weight_version, + ) # Create data iterator for log_probs and train. data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) @@ -364,7 +378,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: ) ) self._switch_model("old_actor" if self.args.keep_old_actor else "actor") - if not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics: + if higgs_policy or not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics: for m in all_replay_managers: if m.enabled: if self._use_rollout_replay(m): @@ -381,6 +395,25 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: for m in all_replay_managers: if self._use_rollout_replay(m): m.clear_all_forward() + if higgs_policy: + parity = validate_higgs_logprob_parity( + rollout_data["action_traces"], + rollout_data["log_probs"], + atol=self.args.higgs_logprob_parity_atol, + ) + log = logger.info if parity["within_tolerance"] else logger.warning + log( + "Higgs pre-optimizer joint-logprob parity %s: " + "max_abs_diff=%.6g mean_abs_diff=%.6g warning_atol=%.6g", + ( + "within measured tolerance" + if parity["within_tolerance"] + else "exceeded measured tolerance" + ), + parity["max_abs_diff"], + parity["mean_abs_diff"], + self.args.higgs_logprob_parity_atol, + ) if self.args.use_critic: sync_actor_critic_data( @@ -466,6 +499,12 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: if self.args.offload_train: destroy_process_groups() + @timer + def disconnect_rollout_engines(self) -> None: + disconnect = getattr(self.weight_updater, "disconnect_rollout_engines", None) + if disconnect is not None: + disconnect() + @timer def update_weights(self, info: "EnginesAndLock") -> None: if self.args.debug_train_only or self.args.debug_rollout_only: diff --git a/miles/backends/megatron_utils/checkpoint.py b/miles/backends/megatron_utils/checkpoint.py index 7d6ba7681ec..42bacebfedd 100644 --- a/miles/backends/megatron_utils/checkpoint.py +++ b/miles/backends/megatron_utils/checkpoint.py @@ -1,8 +1,10 @@ import logging import os import re +from collections import Counter from pathlib import Path +import torch import torch.distributed as dist # TODO: may need to copy those 2 functions and do refactoring. @@ -97,23 +99,94 @@ def _init_from_local_shards_and_global_metadata( # type: ignore[override] __all__ = ["save_checkpoint", "save_checkpoint_with_lora", "load_checkpoint"] +def _normalize_torch_optimizer_steps_for_checkpoint_load(optimizer) -> None: + """Make disposable native-Adam load templates internally consistent.""" + + states = [] + for wrapped_optimizer in getattr(optimizer, "chained_optimizers", (optimizer,)): + torch_optimizer = getattr(wrapped_optimizer, "optimizer", None) + if torch_optimizer is None: + continue + serialized_state = torch_optimizer.state_dict().get("state", {}) + if not serialized_state: + initialize_states = getattr(wrapped_optimizer, "_init_optimizer_states_with_dummy_values", None) + if initialize_states is not None: + logger.info("Initializing temporary native optimizer state for Higgs checkpoint load") + initialize_states() + serialized_state = torch_optimizer.state_dict().get("state", {}) + states.extend(state for state in serialized_state.values() if "step" in state) + + step_counts = Counter(float(state["step"].item()) for state in states) + if len(step_counts) <= 1: + return + + logger.warning( + "Normalizing divergent temporary Torch optimizer steps before Higgs checkpoint load: %s", + dict(sorted(step_counts.items())), + ) + for state in states: + step = state["step"] + if isinstance(step, torch.Tensor): + step.zero_() + else: + state["step"] = 0 + + +def _load_higgs_megatron_checkpoint_with_consistent_optimizer_steps(checkpoint_optimizer, **load_kwargs): + """Normalize native-Adam template steps at the point Megatron materializes them.""" + + patched_state_dicts = [] + for wrapped_optimizer in getattr(checkpoint_optimizer, "chained_optimizers", (checkpoint_optimizer,)): + original_state_dict = wrapped_optimizer.state_dict + + def state_dict(_optimizer=wrapped_optimizer, _original=original_state_dict): + _normalize_torch_optimizer_steps_for_checkpoint_load(_optimizer) + return _original() + + patched_state_dicts.append((wrapped_optimizer, original_state_dict)) + wrapped_optimizer.state_dict = state_dict + try: + return _load_checkpoint_megatron(**load_kwargs) + finally: + for wrapped_optimizer, original_state_dict in patched_state_dicts: + wrapped_optimizer.state_dict = original_state_dict + + def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, checkpointing_context, skip_load_to_model_and_opt): # ref: how megatron `load_checkpoint` gets directory args = get_args() load_path = args.load + from miles.backends.training_utils.higgs_policy import is_higgs_policy_enabled + + if is_higgs_policy_enabled(args): + from .higgs_checkpoint import resolve_higgs_checkpoint_path + + load_path = str(resolve_higgs_checkpoint_path(load_path)) + args.load = load_path + assert Path(load_path).exists() and _is_dir_nonempty( load_path ), f"{args.load=} does not exist or is an empty directory. Did you specify the wrong folder?" if _is_megatron_checkpoint(load_path): - result = _load_checkpoint_megatron( - ddp_model=ddp_model, - optimizer=optimizer, - opt_param_scheduler=opt_param_scheduler, - checkpointing_context=checkpointing_context, - skip_load_to_model_and_opt=skip_load_to_model_and_opt, - ) + if is_higgs_policy_enabled(args) and optimizer is not None: + result = _load_higgs_megatron_checkpoint_with_consistent_optimizer_steps( + optimizer, + ddp_model=ddp_model, + optimizer=optimizer, + opt_param_scheduler=opt_param_scheduler, + checkpointing_context=checkpointing_context, + skip_load_to_model_and_opt=skip_load_to_model_and_opt, + ) + else: + result = _load_checkpoint_megatron( + ddp_model=ddp_model, + optimizer=optimizer, + opt_param_scheduler=opt_param_scheduler, + checkpointing_context=checkpointing_context, + skip_load_to_model_and_opt=skip_load_to_model_and_opt, + ) else: result = _load_checkpoint_hf( ddp_model=ddp_model, @@ -172,14 +245,36 @@ def _is_megatron_checkpoint(path: str | Path) -> bool: def _load_checkpoint_hf(ddp_model, optimizer, args, load_path: str): - assert args.megatron_to_hf_mode == "bridge", "Only bridge mode is supported for loading HF checkpoint" - from megatron.bridge import AutoBridge + from miles.backends.training_utils.higgs_policy import is_higgs_policy_enabled, validate_higgs_single_device_config logger.info(f"Load checkpoint from HuggingFace model into Megatron (path={load_path})") - with megatron_bridge_utils.patch_megatron_model(ddp_model): - bridge = AutoBridge.from_hf_pretrained(load_path, trust_remote_code=True) - bridge.load_hf_weights(ddp_model) + if is_higgs_policy_enabled(args): + from megatron.core.utils import unwrap_model + + from .higgs_checkpoint import HIGGS_TEXT_VOCAB_SIZE, load_higgs_policy_checkpoint + + validate_higgs_single_device_config(args) + if args.megatron_to_hf_mode != "raw": + raise ValueError("Higgs HF loading requires megatron_to_hf_mode='raw'") + if args.vocab_size != HIGGS_TEXT_VOCAB_SIZE or args.padded_vocab_size != HIGGS_TEXT_VOCAB_SIZE: + raise ValueError( + "Higgs HF loading requires vocab_size=padded_vocab_size=" + f"{HIGGS_TEXT_VOCAB_SIZE}, got vocab_size={args.vocab_size!r} " + f"and padded_vocab_size={args.padded_vocab_size!r}" + ) + unwrapped_model = unwrap_model(ddp_model) + if len(unwrapped_model) != 1: + raise ValueError("the initial Higgs raw loader requires exactly one Megatron model chunk") + load_higgs_policy_checkpoint(unwrapped_model[0], load_path) + else: + if args.megatron_to_hf_mode != "bridge": + raise ValueError("only bridge mode is supported for loading a generic HF checkpoint") + from megatron.bridge import AutoBridge + + with megatron_bridge_utils.patch_megatron_model(ddp_model): + bridge = AutoBridge.from_hf_pretrained(load_path, trust_remote_code=True) + bridge.load_hf_weights(ddp_model) # Copied from Megatron-core :: load_checkpoint (with simplifications) if (args.fp16 or args.bf16) and optimizer is not None: diff --git a/miles/backends/megatron_utils/higgs_checkpoint.py b/miles/backends/megatron_utils/higgs_checkpoint.py new file mode 100644 index 00000000000..3b4e7fdf483 --- /dev/null +++ b/miles/backends/megatron_utils/higgs_checkpoint.py @@ -0,0 +1,380 @@ +"""Strict raw-checkpoint loading for the initial single-device Higgs policy.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from contextlib import ExitStack +from pathlib import Path +from types import MappingProxyType + +import torch +from safetensors import safe_open + +HIGGS_NUM_LAYERS = 36 +HIGGS_HIDDEN_SIZE = 2560 +HIGGS_TEXT_VOCAB_SIZE = 151936 +HIGGS_NUM_ATTENTION_HEADS = 32 +HIGGS_NUM_QUERY_GROUPS = 8 +HIGGS_HEAD_DIM = 128 +HIGGS_FFN_HIDDEN_SIZE = 9728 +HIGGS_NUM_CODEBOOKS = 8 +HIGGS_CODEBOOK_VOCAB_SIZE = 1026 +HIGGS_CODEC_ROWS = HIGGS_NUM_CODEBOOKS * HIGGS_CODEBOOK_VOCAB_SIZE + +_FROZEN_CODEC_PREFIX = "tied.embedding.modality_embeddings.0.model." +_INDEX_NAME = "model.safetensors.index.json" +_EXPECTED_SAFETENSORS_DTYPE = "BF16" + + +def _build_policy_shapes() -> dict[str, tuple[int, ...]]: + shapes: dict[str, tuple[int, ...]] = { + "body.norm.weight": (HIGGS_HIDDEN_SIZE,), + "tied.embedding.text_embedding.weight": ( + HIGGS_TEXT_VOCAB_SIZE, + HIGGS_HIDDEN_SIZE, + ), + "tied.embedding.modality_embeddings.0.embedding.weight": ( + HIGGS_CODEC_ROWS, + HIGGS_HIDDEN_SIZE, + ), + } + q_rows = HIGGS_NUM_ATTENTION_HEADS * HIGGS_HEAD_DIM + kv_rows = HIGGS_NUM_QUERY_GROUPS * HIGGS_HEAD_DIM + for layer in range(HIGGS_NUM_LAYERS): + prefix = f"body.layers.{layer}" + shapes.update( + { + f"{prefix}.input_layernorm.weight": (HIGGS_HIDDEN_SIZE,), + f"{prefix}.post_attention_layernorm.weight": (HIGGS_HIDDEN_SIZE,), + f"{prefix}.self_attn.q_norm.weight": (HIGGS_HEAD_DIM,), + f"{prefix}.self_attn.k_norm.weight": (HIGGS_HEAD_DIM,), + f"{prefix}.self_attn.q_proj.weight": (q_rows, HIGGS_HIDDEN_SIZE), + f"{prefix}.self_attn.k_proj.weight": (kv_rows, HIGGS_HIDDEN_SIZE), + f"{prefix}.self_attn.v_proj.weight": (kv_rows, HIGGS_HIDDEN_SIZE), + f"{prefix}.self_attn.o_proj.weight": (HIGGS_HIDDEN_SIZE, q_rows), + f"{prefix}.mlp.gate_proj.weight": ( + HIGGS_FFN_HIDDEN_SIZE, + HIGGS_HIDDEN_SIZE, + ), + f"{prefix}.mlp.up_proj.weight": ( + HIGGS_FFN_HIDDEN_SIZE, + HIGGS_HIDDEN_SIZE, + ), + f"{prefix}.mlp.down_proj.weight": ( + HIGGS_HIDDEN_SIZE, + HIGGS_FFN_HIDDEN_SIZE, + ), + } + ) + return shapes + + +_CANONICAL_HIGGS_POLICY_SHAPES: Mapping[str, tuple[int, ...]] = MappingProxyType(_build_policy_shapes()) + + +def canonical_higgs_policy_shapes() -> dict[str, tuple[int, ...]]: + """Return the 399 indexed trainable-policy names and exact v3 shapes.""" + + return dict(_CANONICAL_HIGGS_POLICY_SHAPES) + + +def resolve_higgs_checkpoint_path(load_path: str | Path) -> Path: + """Resolve a local checkpoint directory or download a Hugging Face repo ID.""" + + if not isinstance(load_path, (str, Path)): + raise FileNotFoundError(f"Higgs checkpoint path must be a string or Path, got {load_path!r}") + path = Path(load_path) + if path.is_dir(): + return path + repo_id = str(load_path) + if path.is_absolute() or len(path.parts) != 2: + raise FileNotFoundError( + f"Higgs checkpoint {repo_id!r} is not a local directory or a Hugging Face owner/repo ID" + ) + try: + from huggingface_hub import snapshot_download + + resolved = Path(snapshot_download(repo_id=repo_id)) + except Exception as error: + raise FileNotFoundError(f"failed to resolve Higgs checkpoint {repo_id!r}: {error}") from error + if not resolved.is_dir(): + raise FileNotFoundError(f"resolved Higgs checkpoint is not a directory: {resolved}") + return resolved + + +def _summarize_names(names: set[str]) -> str: + ordered = sorted(names) + shown = ordered[:8] + suffix = f" ... ({len(ordered)} total)" if len(ordered) > len(shown) else "" + return f"{shown}{suffix}" + + +def _load_weight_map(checkpoint_dir: Path) -> dict[str, str]: + index_path = checkpoint_dir / _INDEX_NAME + if not index_path.is_file(): + raise ValueError(f"Higgs raw loading requires {_INDEX_NAME} in {checkpoint_dir}") + try: + payload = json.loads(index_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"failed to read Higgs safetensors index {index_path}: {error}") from error + if not isinstance(payload, dict) or not isinstance(payload.get("weight_map"), dict): + raise ValueError(f"Higgs safetensors index {index_path} has no object weight_map") + weight_map = payload["weight_map"] + if any(not isinstance(name, str) or not isinstance(filename, str) for name, filename in weight_map.items()): + raise ValueError("Higgs safetensors weight_map must contain string names and filenames") + return dict(weight_map) + + +def _resolve_weight_files(checkpoint_dir: Path, weight_map: Mapping[str, str]) -> dict[str, Path]: + files: dict[str, Path] = {} + for filename in sorted(set(weight_map.values())): + relative_path = Path(filename) + if relative_path.parts != (filename,): + raise ValueError(f"invalid Higgs safetensors filename {filename!r}") + path = checkpoint_dir / relative_path + if not path.is_file(): + raise ValueError(f"invalid or missing Higgs safetensors file {filename!r}") + files[filename] = path + return files + + +def validate_higgs_checkpoint_manifest(checkpoint_dir: str | Path) -> dict[str, str]: + """Validate the indexed canonical policy surface without loading tensor data.""" + + checkpoint_dir = Path(checkpoint_dir) + weight_map = _load_weight_map(checkpoint_dir) + expected = set(_CANONICAL_HIGGS_POLICY_SHAPES) + indexed_policy = {name for name in weight_map if not name.startswith(_FROZEN_CODEC_PREFIX)} + missing = expected - indexed_policy + unexpected = indexed_policy - expected + if missing or unexpected: + details = [] + if missing: + details.append(f"missing={_summarize_names(missing)}") + if unexpected: + details.append(f"unexpected={_summarize_names(unexpected)}") + raise ValueError("Higgs checkpoint policy manifest mismatch: " + "; ".join(details)) + + files = _resolve_weight_files(checkpoint_dir, {name: weight_map[name] for name in expected}) + names_by_file: dict[str, list[str]] = {} + for name in expected: + names_by_file.setdefault(weight_map[name], []).append(name) + + for filename, names in names_by_file.items(): + with safe_open(files[filename], framework="pt", device="cpu") as handle: + available = set(handle.keys()) + missing_from_file = set(names) - available + if missing_from_file: + raise ValueError( + f"Higgs safetensors file {filename!r} is missing indexed tensors " + f"{_summarize_names(missing_from_file)}" + ) + for name in names: + tensor_slice = handle.get_slice(name) + actual_shape = tuple(tensor_slice.get_shape()) + expected_shape = _CANONICAL_HIGGS_POLICY_SHAPES[name] + if actual_shape != expected_shape: + raise ValueError( + f"Higgs checkpoint shape mismatch for {name}: " + f"expected {expected_shape}, got {actual_shape}" + ) + actual_dtype = tensor_slice.get_dtype() + if actual_dtype != _EXPECTED_SAFETENSORS_DTYPE: + raise ValueError( + f"Higgs checkpoint dtype mismatch for {name}: " + f"expected {_EXPECTED_SAFETENSORS_DTYPE}, got {actual_dtype}" + ) + + return {name: weight_map[name] for name in expected} + + +def _target_parameter_shapes(norm_layout: str) -> dict[str, tuple[int, ...]]: + if norm_layout not in {"transformer_engine", "local"}: + raise ValueError(f"unknown Higgs norm layout {norm_layout!r}") + shapes: dict[str, tuple[int, ...]] = { + "embedding.word_embeddings.weight": ( + HIGGS_TEXT_VOCAB_SIZE, + HIGGS_HIDDEN_SIZE, + ), + "decoder.final_layernorm.weight": (HIGGS_HIDDEN_SIZE,), + "codec_embeddings.weight": (HIGGS_CODEC_ROWS, HIGGS_HIDDEN_SIZE), + } + qkv_rows = (HIGGS_NUM_ATTENTION_HEADS + 2 * HIGGS_NUM_QUERY_GROUPS) * HIGGS_HEAD_DIM + q_rows = HIGGS_NUM_ATTENTION_HEADS * HIGGS_HEAD_DIM + for layer in range(HIGGS_NUM_LAYERS): + prefix = f"decoder.layers.{layer}" + if norm_layout == "transformer_engine": + input_norm = f"{prefix}.self_attention.linear_qkv.layer_norm_weight" + post_norm = f"{prefix}.mlp.linear_fc1.layer_norm_weight" + else: + input_norm = f"{prefix}.input_layernorm.weight" + post_norm = f"{prefix}.pre_mlp_layernorm.weight" + shapes.update( + { + input_norm: (HIGGS_HIDDEN_SIZE,), + post_norm: (HIGGS_HIDDEN_SIZE,), + f"{prefix}.self_attention.q_layernorm.weight": (HIGGS_HEAD_DIM,), + f"{prefix}.self_attention.k_layernorm.weight": (HIGGS_HEAD_DIM,), + f"{prefix}.self_attention.linear_qkv.weight": ( + qkv_rows, + HIGGS_HIDDEN_SIZE, + ), + f"{prefix}.self_attention.linear_proj.weight": ( + HIGGS_HIDDEN_SIZE, + q_rows, + ), + f"{prefix}.mlp.linear_fc1.weight": ( + 2 * HIGGS_FFN_HIDDEN_SIZE, + HIGGS_HIDDEN_SIZE, + ), + f"{prefix}.mlp.linear_fc2.weight": ( + HIGGS_HIDDEN_SIZE, + HIGGS_FFN_HIDDEN_SIZE, + ), + } + ) + return shapes + + +def validate_higgs_target_parameters(named_parameters: Mapping[str, torch.Tensor]) -> str: + """Validate the wrapper parameter manifest and return its norm layout.""" + + actual_names = set(named_parameters) + matching_layouts = [] + for layout in ("transformer_engine", "local"): + if actual_names == set(_target_parameter_shapes(layout)): + matching_layouts.append(layout) + if len(matching_layouts) != 1: + te_names = set(_target_parameter_shapes("transformer_engine")) + local_names = set(_target_parameter_shapes("local")) + missing_te, unexpected_te = te_names - actual_names, actual_names - te_names + missing_local, unexpected_local = local_names - actual_names, actual_names - local_names + raise ValueError( + "Higgs Megatron parameter manifest does not match a supported TP1 model; " + f"transformer_engine missing={_summarize_names(missing_te)} " + f"unexpected={_summarize_names(unexpected_te)}; " + f"local missing={_summarize_names(missing_local)} " + f"unexpected={_summarize_names(unexpected_local)}" + ) + + layout = matching_layouts[0] + for name, expected_shape in _target_parameter_shapes(layout).items(): + parameter = named_parameters[name] + if tuple(parameter.shape) != expected_shape: + raise ValueError( + f"Higgs Megatron shape mismatch for {name}: " + f"expected {expected_shape}, got {tuple(parameter.shape)}" + ) + if parameter.dtype != torch.bfloat16: + raise ValueError( + f"Higgs Megatron dtype mismatch for {name}: expected torch.bfloat16, got {parameter.dtype}" + ) + return layout + + +def _copy_parameter(parameter: torch.Tensor, source: torch.Tensor, name: str) -> None: + if tuple(source.shape) != tuple(parameter.shape): + raise ValueError( + f"Higgs load shape mismatch for {name}: expected {tuple(parameter.shape)}, got {tuple(source.shape)}" + ) + if source.dtype != torch.bfloat16: + raise ValueError(f"Higgs load dtype mismatch for {name}: expected torch.bfloat16, got {source.dtype}") + parameter.copy_(source) + + +def load_higgs_policy_checkpoint(model: torch.nn.Module, checkpoint_dir: str | Path) -> set[str]: + """Load the canonical v3 policy weights into an unwrapped TP1 Higgs model.""" + + if model.num_codebooks != HIGGS_NUM_CODEBOOKS: + raise ValueError(f"Higgs model must have {HIGGS_NUM_CODEBOOKS} codebooks") + if model.codebook_vocab_size != HIGGS_CODEBOOK_VOCAB_SIZE: + raise ValueError(f"Higgs model codebook vocabulary must be {HIGGS_CODEBOOK_VOCAB_SIZE}") + + named_parameters = dict(model.named_parameters()) + norm_layout = validate_higgs_target_parameters(named_parameters) + weight_map = validate_higgs_checkpoint_manifest(checkpoint_dir) + files = _resolve_weight_files(Path(checkpoint_dir), weight_map) + loaded: set[str] = set() + + with ExitStack() as stack, torch.no_grad(): + handles = { + filename: stack.enter_context(safe_open(path, framework="pt", device="cpu")) + for filename, path in files.items() + } + + def source(name: str) -> torch.Tensor: + tensor = handles[weight_map[name]].get_tensor(name) + loaded.add(name) + return tensor + + _copy_parameter( + named_parameters["embedding.word_embeddings.weight"], + source("tied.embedding.text_embedding.weight"), + "embedding.word_embeddings.weight", + ) + _copy_parameter( + named_parameters["codec_embeddings.weight"], + source("tied.embedding.modality_embeddings.0.embedding.weight"), + "codec_embeddings.weight", + ) + _copy_parameter( + named_parameters["decoder.final_layernorm.weight"], + source("body.norm.weight"), + "decoder.final_layernorm.weight", + ) + + for layer in range(HIGGS_NUM_LAYERS): + source_prefix = f"body.layers.{layer}" + target_prefix = f"decoder.layers.{layer}" + if norm_layout == "transformer_engine": + input_norm = f"{target_prefix}.self_attention.linear_qkv.layer_norm_weight" + post_norm = f"{target_prefix}.mlp.linear_fc1.layer_norm_weight" + else: + input_norm = f"{target_prefix}.input_layernorm.weight" + post_norm = f"{target_prefix}.pre_mlp_layernorm.weight" + + direct_mappings = { + input_norm: f"{source_prefix}.input_layernorm.weight", + post_norm: f"{source_prefix}.post_attention_layernorm.weight", + f"{target_prefix}.self_attention.q_layernorm.weight": f"{source_prefix}.self_attn.q_norm.weight", + f"{target_prefix}.self_attention.k_layernorm.weight": f"{source_prefix}.self_attn.k_norm.weight", + f"{target_prefix}.self_attention.linear_proj.weight": f"{source_prefix}.self_attn.o_proj.weight", + f"{target_prefix}.mlp.linear_fc2.weight": f"{source_prefix}.mlp.down_proj.weight", + } + for target_name, source_name in direct_mappings.items(): + _copy_parameter(named_parameters[target_name], source(source_name), target_name) + + q = source(f"{source_prefix}.self_attn.q_proj.weight") + k = source(f"{source_prefix}.self_attn.k_proj.weight") + v = source(f"{source_prefix}.self_attn.v_proj.weight") + q = q.view(HIGGS_NUM_QUERY_GROUPS, -1, HIGGS_HEAD_DIM, HIGGS_HIDDEN_SIZE) + k = k.view(HIGGS_NUM_QUERY_GROUPS, 1, HIGGS_HEAD_DIM, HIGGS_HIDDEN_SIZE) + v = v.view(HIGGS_NUM_QUERY_GROUPS, 1, HIGGS_HEAD_DIM, HIGGS_HIDDEN_SIZE) + qkv = torch.cat((q, k, v), dim=1).reshape(-1, HIGGS_HIDDEN_SIZE) + qkv_name = f"{target_prefix}.self_attention.linear_qkv.weight" + _copy_parameter(named_parameters[qkv_name], qkv, qkv_name) + + gate = source(f"{source_prefix}.mlp.gate_proj.weight") + up = source(f"{source_prefix}.mlp.up_proj.weight") + fc1 = torch.cat((gate, up), dim=0) + fc1_name = f"{target_prefix}.mlp.linear_fc1.weight" + _copy_parameter(named_parameters[fc1_name], fc1, fc1_name) + + expected = set(_CANONICAL_HIGGS_POLICY_SHAPES) + if loaded != expected: + raise RuntimeError( + "Higgs loader did not consume the canonical policy surface: " + f"missing={_summarize_names(expected - loaded)} " + f"unexpected={_summarize_names(loaded - expected)}" + ) + return loaded + + +__all__ = [ + "canonical_higgs_policy_shapes", + "load_higgs_policy_checkpoint", + "resolve_higgs_checkpoint_path", + "validate_higgs_checkpoint_manifest", + "validate_higgs_target_parameters", +] diff --git a/miles/backends/megatron_utils/higgs_model.py b/miles/backends/megatron_utils/higgs_model.py new file mode 100644 index 00000000000..7e8af81da61 --- /dev/null +++ b/miles/backends/megatron_utils/higgs_model.py @@ -0,0 +1,168 @@ +"""Megatron-native Higgs codec policy wrapper. + +Megatron imports are intentionally local so the structured tensor contract can +be unit-tested in environments that do not install Megatron-LM. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from miles.backends.training_utils.higgs_policy import build_higgs_teacher_embeddings + + +def build_higgs_megatron_model( + *, + gpt_model_kwargs: dict[str, Any], + num_codebooks: int, + codebook_vocab_size: int, +): + """Build a Qwen3 GPTModel with one tied fused codec embedding/head. + + The initial backend is deliberately TP=PP=CP=DP=1. It still uses + Megatron modules so DDP wrapping, optimizer construction, scheduling, and + checkpoint lifecycle remain owned by Megatron. + """ + + from megatron.core import tensor_parallel + from megatron.core.models.gpt import GPTModel + + if num_codebooks <= 0 or codebook_vocab_size <= 0: + raise ValueError("Higgs codebook dimensions must be positive") + + class HiggsMegatronModel(GPTModel): + def setup_embeddings_and_output_layer(self) -> None: + """Install parameter attributes without tying the temporary text head. + + ``GPTModel.__init__`` calls this hook before the codec modules exist. + The initial call must therefore avoid the virtual shared-weight + lookup; the second call below completes setup after construction. + """ + + if self.pre_process: + self.embedding.word_embeddings.weight.is_embedding_or_output_parameter = True + if not hasattr(self, "codec_embeddings"): + return + self.codec_embeddings.weight.is_embedding_or_output_parameter = True + self.codec_embeddings.weight.zero_out_wgrad = True + + def __init__(self) -> None: + super().__init__(**gpt_model_kwargs) + if not self.pre_process or not self.post_process: + raise ValueError("the initial Higgs Megatron model requires pipeline parallel size 1") + if self.config.sequence_parallel: + raise ValueError("the initial Higgs Megatron model does not support sequence parallelism") + + self.num_codebooks = int(num_codebooks) + self.codebook_vocab_size = int(codebook_vocab_size) + codec_rows = self.num_codebooks * self.codebook_vocab_size + self.codec_embeddings = tensor_parallel.VocabParallelEmbedding( + num_embeddings=codec_rows, + embedding_dim=self.config.hidden_size, + init_method=self.config.embedding_init_method, + config=self.config, + tp_group=self.pg_collection.tp, + ) + # The checkpoint ties modality input and output weights. Allocate + # the parameter once on codec_embeddings and pass it to this head. + self.output_layer = tensor_parallel.ColumnParallelLinear( + self.config.hidden_size, + codec_rows, + config=self.config, + init_method=self.config.init_method, + bias=False, + gather_output=False, + skip_bias_add=False, + skip_weight_param_allocation=True, + tp_group=self.pg_collection.tp, + ) + self.share_embeddings_and_output_weights = True + self.vocab_size = codec_rows + self.setup_embeddings_and_output_layer() + + def shared_embedding_or_output_weight(self) -> torch.Tensor: + return self.codec_embeddings.weight + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + *, + higgs_prior_codes: torch.Tensor | None = None, + higgs_codec_position_mask: torch.Tensor | None = None, + higgs_sequence_mask: torch.Tensor | None = None, + higgs_prediction_positions: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + required = ( + higgs_prior_codes, + higgs_codec_position_mask, + higgs_sequence_mask, + higgs_prediction_positions, + ) + if any(value is None for value in required): + raise ValueError("HiggsMegatronModel requires the complete structured teacher-forcing batch") + if input_ids.ndim != 2: + raise ValueError("Higgs input_ids must have shape [batch, sequence]") + if higgs_sequence_mask.shape != input_ids.shape: + raise ValueError("Higgs sequence mask must match input_ids") + if higgs_prediction_positions.ndim != 2 or higgs_prediction_positions.shape[0] != input_ids.shape[0]: + raise ValueError("Higgs prediction positions must have shape [batch, action_rows]") + + text_embeddings = self.embedding.word_embeddings(input_ids) + embeddings = build_higgs_teacher_embeddings( + text_embeddings, + self.codec_embeddings.weight, + higgs_prior_codes, + higgs_codec_position_mask, + ) + decoder_input = embeddings.transpose(0, 1).contiguous() + if self.config.fp32_residual_connection: + decoder_input = decoder_input.float() + decoder_input = self.embedding.embedding_dropout(decoder_input) + + logits = super().forward( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + decoder_input=decoder_input, + labels=None, + packed_seq_params=None, + padding_mask=~higgs_sequence_mask, + **kwargs, + ) + if logits.ndim != 3 or logits.shape[:2] != input_ids.shape: + raise RuntimeError(f"unexpected Higgs Megatron logits shape {tuple(logits.shape)}") + if logits.shape[-1] != self.num_codebooks * self.codebook_vocab_size: + raise RuntimeError("Higgs codec head returned the wrong flattened vocabulary size") + + gather_index = higgs_prediction_positions.unsqueeze(-1).expand(-1, -1, logits.shape[-1]) + action_logits = logits.gather(dim=1, index=gather_index) + return action_logits.reshape( + input_ids.shape[0], + higgs_prediction_positions.shape[1], + self.num_codebooks, + self.codebook_vocab_size, + ) + + return HiggsMegatronModel() + + +def higgs_model_forward_kwargs(batch: Any) -> dict[str, torch.Tensor]: + """Translate a ``HiggsPolicyBatch`` into the model's explicit API.""" + + return { + "input_ids": batch.input_ids, + "position_ids": None, + "attention_mask": None, + "higgs_prior_codes": batch.prior_codes, + "higgs_codec_position_mask": batch.codec_position_mask, + "higgs_sequence_mask": batch.sequence_mask, + "higgs_prediction_positions": batch.prediction_positions, + } + + +__all__ = ["build_higgs_megatron_model", "higgs_model_forward_kwargs"] diff --git a/miles/backends/megatron_utils/megatron_to_hf/__init__.py b/miles/backends/megatron_utils/megatron_to_hf/__init__.py index 0873e6f8e32..74bd82d742f 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/__init__.py +++ b/miles/backends/megatron_utils/megatron_to_hf/__init__.py @@ -2,6 +2,7 @@ from .deepseekv4 import convert_deepseekv4_to_hf from .glm4 import convert_glm4_to_hf from .glm4moe import convert_glm4moe_to_hf +from .higgs_tts import convert_higgs_to_hf from .kimi_vl import convert_kimi_k25_to_hf, convert_kimivl_to_hf from .llama import convert_llama_to_hf from .mimo import convert_mimo_to_hf @@ -31,7 +32,14 @@ def convert_to_hf(args, model_name, name, param, quantization_config=None): # TODO optimize code details def _convert_to_hf_core(args, model_name, name, param): model_name = model_name.lower() + normalized_model_name = model_name.replace("-", "").replace("_", "") if ( + getattr(args, "structured_policy_model_family", None) == "higgs_tts" + or "higgsmultimodalqwen3" in normalized_model_name + or normalized_model_name == "higgstts" + ): + converted_named_tensors = convert_higgs_to_hf(args, name, param) + elif ( "glm4moelite" in model_name or "deepseekv3" in model_name or "glmmoedsa" in model_name diff --git a/miles/backends/megatron_utils/megatron_to_hf/higgs_tts.py b/miles/backends/megatron_utils/megatron_to_hf/higgs_tts.py new file mode 100644 index 00000000000..252826ee8fb --- /dev/null +++ b/miles/backends/megatron_utils/megatron_to_hf/higgs_tts.py @@ -0,0 +1,39 @@ +from argparse import Namespace + +import torch + +from .qwen2 import convert_qwen2_to_hf + + +def convert_higgs_to_hf(args: Namespace, name: str, param: torch.Tensor) -> list[tuple[str, torch.Tensor]]: + """Convert the TP-gathered Higgs policy weights to canonical checkpoint names.""" + + if name == "module.module.codec_embeddings.weight": + expected_shape = ( + args.higgs_num_codebooks * args.higgs_codebook_vocab_size, + args.hidden_size, + ) + if tuple(param.shape) != expected_shape: + raise ValueError( + "Higgs codec embedding shape mismatch: " f"expected {expected_shape}, got {tuple(param.shape)}" + ) + return [("tied.embedding.modality_embeddings.0.embedding.weight", param)] + + converted = convert_qwen2_to_hf(args, name, param) + renamed: list[tuple[str, torch.Tensor]] = [] + for hf_name, tensor in converted: + if hf_name == "model.embed_tokens.weight": + canonical_name = "tied.embedding.text_embedding.weight" + elif hf_name == "model.norm.weight": + canonical_name = "body.norm.weight" + elif hf_name.startswith("model.layers."): + canonical_name = "body.layers." + hf_name.removeprefix("model.layers.") + elif hf_name == "lm_head.weight": + raise ValueError("the tied Higgs text head must not be a separate Megatron parameter") + else: + raise ValueError(f"unsupported Higgs policy parameter mapping: {name!r} -> {hf_name!r}") + renamed.append((canonical_name, tensor)) + return renamed + + +__all__ = ["convert_higgs_to_hf"] diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 44fdaeedcd8..38ef6bb50d8 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -26,7 +26,8 @@ from miles.utils.memory_utils import clear_memory from ..training_utils.ci_utils import check_grad_norm, check_kl -from ..training_utils.data import DataIterator, get_batch +from ..training_utils.data import DataIterator, get_batch, get_higgs_batch +from ..training_utils.higgs_policy import get_higgs_joint_log_probs, is_higgs_policy_enabled from ..training_utils.log_utils import aggregate_forward_results, aggregate_train_losses, log_train_step from ..training_utils.loss import loss_function from ..training_utils.parallel import get_parallel_state @@ -37,6 +38,7 @@ compute_model_hashes_by_layer, save_model_hashes, ) +from .higgs_model import higgs_model_forward_kwargs from .initialize import is_megatron_main_rank from .lora_utils import is_lora_enabled, is_lora_model from .model_provider import get_model_provider_func @@ -255,6 +257,11 @@ def forward_step( assert not return_schedule_plan, "forward_only step should never return schedule plan" + if is_higgs_policy_enabled(args): + higgs_batch = get_higgs_batch(data_iterator, args, require_advantages=False) + output_tensor = model(**higgs_model_forward_kwargs(higgs_batch)) + return output_tensor, partial(get_higgs_joint_log_probs, higgs_batch=higgs_batch) + # Get the batch. batch = get_batch( data_iterator, @@ -399,29 +406,35 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p (loss, num_elems, {"keys": list[str], "values": torch.Tensor}). """ - # Get the batch. - batch = get_batch( - data_iterator, - [ - "tokens", - "multimodal_train_inputs", - "packed_seq_params", - "total_lengths", - "response_lengths", - "loss_masks", - "log_probs", - "ref_log_probs", - "values", - "advantages", - "returns", - "rollout_log_probs", - "max_seq_lens", - "opd_reverse_kl", - ], - args.data_pad_size_multiplier, - args.qkv_format, - allgather_cp=args.allgather_cp, - ) + if is_higgs_policy_enabled(args): + if return_schedule_plan: + raise ValueError("combined 1f1b schedule plans are not implemented for Higgs structured policy") + higgs_batch = get_higgs_batch(data_iterator, args, require_advantages=True) + batch = {"higgs_policy_batch": higgs_batch} + else: + # Get the legacy causal-text batch. + batch = get_batch( + data_iterator, + [ + "tokens", + "multimodal_train_inputs", + "packed_seq_params", + "total_lengths", + "response_lengths", + "loss_masks", + "log_probs", + "ref_log_probs", + "values", + "advantages", + "returns", + "rollout_log_probs", + "max_seq_lens", + "opd_reverse_kl", + ], + args.data_pad_size_multiplier, + args.qkv_format, + allgather_cp=args.allgather_cp, + ) from miles.utils.replay_base import all_replay_managers @@ -429,7 +442,9 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p for m in all_replay_managers: m.stage = "replay_forward" - if return_schedule_plan: + if is_higgs_policy_enabled(args): + output_tensor = model(**higgs_model_forward_kwargs(higgs_batch)) + elif return_schedule_plan: assert not args.enable_mtp_training, "MTP training should not be enabled when using combined 1f1b" output_tensor = model.build_schedule_plan( input_ids=batch["tokens"], diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py index 6aa694e8761..c54a6cec1be 100644 --- a/miles/backends/megatron_utils/model_provider.py +++ b/miles/backends/megatron_utils/model_provider.py @@ -20,6 +20,9 @@ from miles.utils.misc import load_function from miles.utils.replay_base import routing_replay_manager +from ..training_utils.higgs_policy import is_higgs_policy_enabled, validate_higgs_single_device_config +from .higgs_model import build_higgs_megatron_model + logger = logging.getLogger(__name__) @@ -131,6 +134,14 @@ def get_model_provider_func( args: argparse.Namespace, role: Literal["actor", "critic"] = "actor", ): + higgs_policy = is_higgs_policy_enabled(args) + if higgs_policy: + validate_higgs_single_device_config(args) + if role != "actor": + raise ValueError("the initial Higgs structured policy path does not implement a critic") + if getattr(args, "custom_model_provider_path", None): + raise ValueError("Higgs structured policy owns its Megatron model provider") + # Support custom model provider path (similar to --custom-rm-path for reward models) if getattr(args, "custom_model_provider_path", None): @@ -159,6 +170,10 @@ def wrapped_model_provider( return wrapped_model_provider if args.megatron_to_hf_mode == "bridge": + if higgs_policy: + raise ValueError( + "Higgs Megatron-Bridge loading is disabled until a checkpoint mapping passes server/trainer parity" + ) from megatron.bridge import AutoBridge bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True) @@ -313,7 +328,14 @@ def model_provider( routing_replay_manager.enabled = True with build_model_context(**build_model_context_args): - model = GPTModel(**kwargs) + if higgs_policy: + model = build_higgs_megatron_model( + gpt_model_kwargs=kwargs, + num_codebooks=args.higgs_num_codebooks, + codebook_vocab_size=args.higgs_codebook_vocab_size, + ) + else: + model = GPTModel(**kwargs) if post_process and role == "critic": model.output_layer = LinearForLastLayer(input_size=config.hidden_size, output_size=1, config=config) diff --git a/miles/backends/megatron_utils/update_weight/common.py b/miles/backends/megatron_utils/update_weight/common.py index dc483e5e543..5bf59a76f3f 100644 --- a/miles/backends/megatron_utils/update_weight/common.py +++ b/miles/backends/megatron_utils/update_weight/common.py @@ -175,6 +175,9 @@ def all_gather_param(args: Namespace, name: str, param: torch.nn.Parameter) -> t tp_size = get_parallel_state().tp.size tp_group = get_parallel_state().tp.group + if tp_size == 1: + return param.data + param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)] dist.all_gather(param_partitions, param.data, group=tp_group) partition_dim = param.partition_dim diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py index 679d31d4c92..0566eb1a7b5 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py @@ -1,3 +1,4 @@ +import os import socket import time from argparse import Namespace @@ -18,6 +19,42 @@ from .mixin import DistBucketedWeightUpdateMixin +def _nccl_version_string() -> str: + version = torch.cuda.nccl.version() + if isinstance(version, tuple): + return ".".join(str(part) for part in version) + return str(version) + + +def _validate_distributed_weight_update_transports( + engine_transports: Sequence[dict | None], +) -> None: + advertised = [transport for transport in engine_transports if transport is not None] + if not advertised: + return + if len(advertised) != len(engine_transports): + raise RuntimeError( + "cannot create one NCCL weight-update group from inference engines " + "with mixed advertised and legacy transport contracts" + ) + + trainer_transport = { + "protocol_version": 1, + "backend": "nccl", + "nccl_version": _nccl_version_string(), + "nccl_cumem_enable": os.environ.get("NCCL_CUMEM_ENABLE", "default"), + } + for index, engine_transport in enumerate(advertised): + if engine_transport != trainer_transport: + raise RuntimeError( + "distributed weight-update transport mismatch before NCCL " + f"rendezvous: trainer={trainer_transport}, " + f"inference_engine[{index}]={engine_transport}. Launch both " + "processes with the same NCCL version and " + "NCCL_CUMEM_ENABLE setting." + ) + + class UpdateWeightFromDistributed(DistBucketedWeightUpdateMixin): """ Update distributed engines via NCCL. Each PP rank: group "miles-pp_{pp_rank}", @@ -79,6 +116,17 @@ def connect_rollout_engines( self.args, self._group_name, rollout_engines ) + def disconnect_rollout_engines(self) -> None: + if not self._is_source or self._model_update_groups is None: + return + disconnect_rollout_engines_from_distributed( + self.args, + self._group_name, + self._model_update_groups, + self.rollout_engines, + ) + self._model_update_groups = None + @property def _is_source(self): """If it's the source gpu that broadcasting weights to rollout side""" @@ -159,6 +207,10 @@ def connect_rollout_engines_from_distributed( """ if engine_gpu_counts is None: engine_gpu_counts = [args.rollout_num_gpus_per_engine] * len(rollout_engines) + engine_transports = ray.get( + [engine.get_distributed_weight_update_transport.remote() for engine in rollout_engines] + ) + _validate_distributed_weight_update_transports(engine_transports) master_address = ray._private.services.get_node_ip_address() with socket.socket() as sock: sock.bind(("", 0)) @@ -209,6 +261,11 @@ def update_weights_from_distributed( """ Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines). """ + # HF conversion commonly returns split or permuted views. NCCL collectives + # require dense tensors, and the receiver allocates from these exact shapes. + converted_named_tensors = [ + (name, tensor if tensor.is_contiguous() else tensor.contiguous()) for name, tensor in converted_named_tensors + ] refs = [ engine.update_weights_from_distributed.remote( names=[name for name, _ in converted_named_tensors], @@ -220,10 +277,20 @@ def update_weights_from_distributed( for engine in rollout_engines ] - handles = [] - for _, param in converted_named_tensors: - handles.append(dist.broadcast(param.data, 0, group=group, async_op=True)) - for handle in handles: - handle.wait() + for name, param in converted_named_tensors: + # Keep only one NCCL collective in flight. Some large dense models have + # hundreds of exported tensors, and enqueueing the complete update at + # once can fail before the receiver drains the first broadcast. + try: + dist.broadcast(param.data, 0, group=group) + except Exception as error: + if hasattr(error, "add_note"): + error.add_note( + "weight broadcast failed for " + f"{name!r}: shape={tuple(param.shape)}, dtype={param.dtype}, " + f"device={param.device}, stride={param.stride()}, " + f"contiguous={param.is_contiguous()}" + ) + raise return refs diff --git a/miles/backends/sglang_utils/arguments.py b/miles/backends/sglang_utils/arguments.py index 71e3f48fb4a..9220a4a383c 100644 --- a/miles/backends/sglang_utils/arguments.py +++ b/miles/backends/sglang_utils/arguments.py @@ -135,7 +135,13 @@ def new_add_argument_wrapper(*name_or_flags, **kwargs): def validate_args(args): + values = vars(args) + # SGLang's CLI uses long-form parallelism destinations in newer releases, + # while ServerArgs and the Miles runtime still use the short field names. args.sglang_tp_size = args.rollout_num_gpus_per_engine + args.sglang_dp_size = values.get("sglang_dp_size", values.get("sglang_data_parallel_size", 1)) + args.sglang_pp_size = values.get("sglang_pp_size", values.get("sglang_pipeline_parallel_size", 1)) + args.sglang_ep_size = values.get("sglang_ep_size", values.get("sglang_expert_parallel_size", 1)) if args.true_on_policy_mode: args.sglang_enable_deterministic_inference = True diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py index 51e44d5e61a..1794114a34b 100644 --- a/miles/backends/sglang_utils/sglang_engine.py +++ b/miles/backends/sglang_utils/sglang_engine.py @@ -22,6 +22,57 @@ logger = logging.getLogger(__name__) +def _extract_omni_distributed_weight_update_transport(model_info: dict) -> dict: + transports = [] + for item in model_info.get("stages", []): + if not isinstance(item, dict) or not isinstance(item.get("data"), dict): + continue + data = item["data"] + if data.get("supports_distributed_weight_update"): + transport = data.get("distributed_weight_update") + if not isinstance(transport, dict): + raise RuntimeError( + "external SGLang-Omni stage supports distributed weight " + "updates but advertises no transport descriptor" + ) + transports.append(transport) + + if not transports: + raise RuntimeError("external SGLang-Omni has no distributed weight-update transport") + if any(transport != transports[0] for transport in transports[1:]): + raise RuntimeError( + f"external SGLang-Omni stages advertise inconsistent weight-update transports: {transports}" + ) + return transports[0] + + +def _validate_omni_server_info(model_info: dict, expect_server_args: dict) -> None: + """Validate the smaller model identity surface exposed by SGLang-Omni.""" + if model_info.get("success") is not True: + raise RuntimeError(f"external SGLang-Omni model_info failed: {model_info}") + + expected_model = str(expect_server_args["model_path"]) + actual_model = model_info.get("model_path") + encoded_hf_model = f"models--{expected_model.replace('/', '--')}" + if not isinstance(actual_model, str) or not (actual_model == expected_model or encoded_hf_model in actual_model): + raise RuntimeError(f"external SGLang-Omni model mismatch: expected {expected_model!r}, got {actual_model!r}") + + expected_tp_size = expect_server_args["tp_size"] + stage_tp_sizes = { + item.get("data", {}).get("tp_size") + for item in model_info.get("stages", []) + if isinstance(item, dict) + and isinstance(item.get("data"), dict) + and item.get("data", {}).get("tp_size") is not None + } + if stage_tp_sizes != {expected_tp_size}: + raise RuntimeError( + f"external SGLang-Omni TP mismatch: expected {expected_tp_size}, got {sorted(stage_tp_sizes)}" + ) + if not isinstance(model_info.get("weight_version"), str) or not model_info["weight_version"]: + raise RuntimeError("external SGLang-Omni model_info has no weight_version") + + def get_base_gpu_id(args, rank): num_gpus = min(args.num_gpus_per_node, args.rollout_num_gpus_per_engine) if args.colocate: @@ -192,6 +243,12 @@ def _init_external(self, expect_server_args, external_engine_need_check_fields): def _get_actual_server_args(): response = requests.get(f"http://{self.server_host}:{self.server_port}/get_server_info") + if response.status_code == 404: + response = requests.get(f"http://{self.server_host}:{self.server_port}/model_info") + response.raise_for_status() + self._omni_model_info = response.json() + _validate_omni_server_info(self._omni_model_info, expect_server_args) + return None response.raise_for_status() return response.json() @@ -209,7 +266,8 @@ def _sanity_check_server_args(actual_server_args, expect_server_args): is_process_alive=lambda: True, ) actual_server_args = _get_actual_server_args() - _sanity_check_server_args(actual_server_args, expect_server_args) + if actual_server_args is not None: + _sanity_check_server_args(actual_server_args, expect_server_args) def _init_normal(self, server_args_dict): logger.info(f"Launch HttpServerEngineAdapter at: {self.server_host}:{self.server_port}") @@ -458,6 +516,12 @@ def get_weight_version(self): return response.json()["weight_version"] response.raise_for_status() + def get_distributed_weight_update_transport(self): + model_info = getattr(self, "_omni_model_info", None) + if model_info is None: + return None + return _extract_omni_distributed_weight_update_transport(model_info) + def unload_lora_adapter(self, lora_name: str): """Unload LoRA adapter.""" return self._make_request( diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index a43e0aecfbf..20e237fa2cf 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -13,6 +13,12 @@ from ...utils.data import process_rollout_data from ...utils.ray_utils import Box from .cp_utils import slice_log_prob_with_cp, slice_with_cp +from .higgs_policy import ( + HiggsPolicyBatch, + collate_higgs_policy_batch, + is_higgs_policy_enabled, + validate_higgs_single_device_config, +) from .mm_data import expand_multimodal_rollout_data_in_place from .parallel import get_parallel_state @@ -38,6 +44,15 @@ def get_rollout_data(args: Namespace, rollout_data_ref: Box) -> RolloutBatch: parallel_state.intra_dp.rank, parallel_state.intra_dp.size, ) + has_action_traces = "action_traces" in rollout_data + if has_action_traces != is_higgs_policy_enabled(args): + raise ValueError( + "structured Higgs action traces and --structured-policy-model-family=higgs_tts must be enabled together" + ) + if has_action_traces: + validate_higgs_single_device_config(args, data_parallel_size=parallel_state.intra_dp.size) + for trace in rollout_data["action_traces"]: + trace.validate() # move tokens to GPU in advance rollout_data["tokens"] = [ torch.tensor(t, dtype=torch.long, device=torch.cuda.current_device()) for t in rollout_data["tokens"] @@ -100,6 +115,51 @@ def get_rollout_data(args: Namespace, rollout_data_ref: Box) -> RolloutBatch: return rollout_data +def get_higgs_batch( + data_iterator: "DataIterator", + args: Namespace, + *, + require_advantages: bool, +) -> HiggsPolicyBatch: + """Fetch and collate one structured Higgs microbatch.""" + + validate_higgs_single_device_config(args, data_parallel_size=get_parallel_state().intra_dp.size) + keys = ["tokens", "action_traces", "advantages"] + if require_advantages: + keys.append("log_probs") + raw_batch = data_iterator.get_next(keys) + if raw_batch["tokens"] is None or raw_batch["action_traces"] is None: + raise ValueError("Higgs microbatches require prompt tokens and typed action traces") + advantages = raw_batch["advantages"] + if require_advantages and advantages is None: + raise ValueError("Higgs training microbatches require GRPO advantages") + old_policy_joint_logprobs = raw_batch.get("log_probs") + if require_advantages and old_policy_joint_logprobs is None: + raise ValueError("Higgs training requires pre-update Megatron joint logprobs") + device = raw_batch["tokens"][0].device + batch = collate_higgs_policy_batch( + raw_batch["tokens"], + raw_batch["action_traces"], + advantages=advantages, + old_policy_joint_logprobs=old_policy_joint_logprobs, + device=device, + ) + if batch.num_codebooks != args.higgs_num_codebooks: + raise ValueError( + f"Higgs rollout has {batch.num_codebooks} codebooks, model expects {args.higgs_num_codebooks}" + ) + if batch.codebook_vocab_size != args.higgs_codebook_vocab_size: + raise ValueError( + "Higgs rollout codebook vocabulary " + f"{batch.codebook_vocab_size} does not match model {args.higgs_codebook_vocab_size}" + ) + if batch.input_ids.shape[1] > args.seq_length: + raise ValueError( + f"Higgs teacher-forcing sequence length {batch.input_ids.shape[1]} exceeds --seq-length={args.seq_length}" + ) + return batch + + def get_batch( data_iterator: "DataIterator", keys: Sequence[str], @@ -381,6 +441,8 @@ def get_data_iterator( parallel_state = get_parallel_state() dp_size = parallel_state.intra_dp.size dp_group = parallel_state.intra_dp.group + if "action_traces" in rollout_data: + validate_higgs_single_device_config(args, data_parallel_size=dp_size) vpp_size = parallel_state.vpp_size microbatch_group_size_per_vp_stage = parallel_state.microbatch_group_size_per_vp_stage diff --git a/miles/backends/training_utils/higgs_policy.py b/miles/backends/training_utils/higgs_policy.py new file mode 100644 index 00000000000..196f005cad6 --- /dev/null +++ b/miles/backends/training_utils/higgs_policy.py @@ -0,0 +1,628 @@ +"""Structured Higgs policy batching and joint-row GRPO math. + +This module deliberately depends only on PyTorch. The tensor contract is +shared by the Megatron integration and CPU unit tests, while checkpoint +conversion remains a separate concern. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F + +HIGGS_MODEL_FAMILY = "higgs_tts" +HIGGS_STREAM_NAME = "higgs_codes" + + +@dataclass(frozen=True) +class HiggsPolicyBatch: + """One padded, unpacked BSHD batch for Higgs teacher forcing.""" + + input_ids: torch.Tensor + prior_codes: torch.Tensor + codec_position_mask: torch.Tensor + sequence_mask: torch.Tensor + prediction_positions: torch.Tensor + actions: torch.Tensor + action_mask: torch.Tensor + old_cell_logprobs: torch.Tensor + old_joint_logprobs: torch.Tensor + row_mask: torch.Tensor + advantages: torch.Tensor | None + prompt_lengths: torch.Tensor + action_lengths: torch.Tensor + num_codebooks: int + codebook_vocab_size: int + + +def is_higgs_policy_enabled(args: Any) -> bool: + return getattr(args, "structured_policy_model_family", None) == HIGGS_MODEL_FAMILY + + +def validate_higgs_single_device_config(args: Any, *, data_parallel_size: int | None = None) -> None: + """Reject configurations outside the first, deliberately naive backend.""" + + if not is_higgs_policy_enabled(args): + return + + required_values = { + "train_backend": "megatron", + "qkv_format": "bshd", + "tensor_model_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + "expert_model_parallel_size": 1, + "actor_num_nodes": 1, + "actor_num_gpus_per_node": 1, + "advantage_estimator": "grpo", + "bf16": True, + "fp16": False, + "true_on_policy_mode": False, + "hidden_dropout": 0.0, + "attention_dropout": 0.0, + "masked_softmax_fusion": False, + "vocab_size": 151936, + "padded_vocab_size": 151936, + "group_query_attention": True, + "num_query_groups": 8, + "kv_channels": 128, + "qk_layernorm": True, + "swiglu": True, + "add_bias_linear": False, + "untie_embeddings_and_output_weights": False, + "normalization": "RMSNorm", + "use_rotary_position_embeddings": True, + "rotary_percent": 1.0, + "rotary_interleaved": False, + "use_rope_scaling": False, + } + errors = [] + for name, expected in required_values.items(): + actual = getattr(args, name, None) + if actual != expected: + errors.append(f"{name}={actual!r} (expected {expected!r})") + # Megatron resolves an unspecified expert TP size to the ordinary TP size + # during its own validation, which is already fixed to one above. + if getattr(args, "expert_tensor_parallel_size", None) not in (None, 1): + errors.append(f"expert_tensor_parallel_size={args.expert_tensor_parallel_size!r} (expected None or 1)") + if getattr(args, "num_experts", None) not in (None, 0): + errors.append("num_experts must be unset for the dense Higgs Qwen3 backbone") + if getattr(args, "mtp_num_layers", None) not in (None, 0): + errors.append("mtp_num_layers must be unset for the initial Higgs policy") + if getattr(args, "spec", None) is not None: + errors.append("custom Megatron layer specs are not supported by the initial Higgs policy") + + false_flags = ( + "sequence_parallel", + "use_dynamic_batch_size", + "use_dynamic_global_batch_size", + "allgather_cp", + "enable_mtp_training", + "use_critic", + "use_tis", + "use_opsm", + "use_opd", + "calculate_per_token_loss", + "get_mismatch_metrics", + "observe_training_entropy", + "use_rollout_entropy", + "use_rollout_routing_replay", + "use_rollout_indexer_replay", + "keep_old_actor", + "debug_disable_optimizer", + ) + for name in false_flags: + if bool(getattr(args, name, False)): + errors.append(f"{name}=True (expected False)") + + if getattr(args, "lora_rank", 0) not in (None, 0): + errors.append("lora_rank must be 0") + if not isinstance(getattr(args, "hf_checkpoint", None), str) or not args.hf_checkpoint: + errors.append("hf_checkpoint must identify the concrete Higgs v3 TTS checkpoint") + if not bool(getattr(args, "use_rollout_logprobs", False)): + errors.append("use_rollout_logprobs must be enabled") + if not bool(getattr(args, "compute_advantages_and_returns", False)): + errors.append("compute_advantages_and_returns must be enabled") + if not bool(getattr(args, "rewards_normalization", False)): + errors.append("rewards_normalization must be enabled for grouped GRPO") + n_samples_per_prompt = getattr(args, "n_samples_per_prompt", None) + if ( + isinstance(n_samples_per_prompt, bool) + or not isinstance(n_samples_per_prompt, int) + or n_samples_per_prompt <= 1 + ): + errors.append("n_samples_per_prompt must be greater than 1 for a nonzero grouped GRPO signal") + kl_coef = getattr(args, "kl_coef", 0.0) + if kl_coef is None: + kl_coef = 0.0 + entropy_coef = getattr(args, "entropy_coef", 0.0) + if entropy_coef is None: + entropy_coef = 0.0 + if kl_coef != 0.0 or bool(getattr(args, "use_kl_loss", False)): + errors.append("reference-policy KL is not implemented for the initial Higgs path") + if entropy_coef != 0.0: + errors.append("entropy_coef must be 0 for the initial Higgs path") + if bool(getattr(args, "normalize_advantages", False)): + errors.append("normalize_advantages must be disabled; GRPO rewards are normalized before batching") + if getattr(args, "megatron_to_hf_mode", None) == "bridge": + errors.append("megatron_to_hf_mode='bridge' has no verified Higgs checkpoint mapping") + if getattr(args, "save_hf", None) is not None: + errors.append( + "save_hf is disabled because standalone Higgs HF snapshot export is not implemented; " + "online raw weight conversion is supported" + ) + parity_atol = getattr(args, "higgs_logprob_parity_atol", None) + if ( + parity_atol is None + or isinstance(parity_atol, bool) + or not isinstance(parity_atol, (int, float)) + or parity_atol <= 0 + ): + errors.append("higgs_logprob_parity_atol must be configured to a positive measured tolerance") + if data_parallel_size is not None and data_parallel_size != 1: + errors.append(f"data_parallel_size={data_parallel_size!r} (expected 1)") + + if errors: + raise ValueError("Higgs structured policy requires the single-device Megatron profile: " + "; ".join(errors)) + + +def _config_dict(config: Any, name: str) -> Mapping[str, Any]: + if isinstance(config, Mapping): + return config + try: + return vars(config) + except TypeError as error: + raise ValueError(f"Higgs {name} must be a configuration mapping or object") from error + + +def validate_higgs_hf_config( + config: Any, + *, + num_codebooks: int, + codebook_vocab_size: int, +) -> None: + """Validate the concrete v3 TTS checkpoint architecture this adapter mirrors.""" + + root = _config_dict(config, "root config") + audio = _config_dict(root.get("audio_encoder_config"), "audio_encoder_config") + text = _config_dict(root.get("text_config"), "text_config") + expected = { + "model_type": (root.get("model_type"), "higgs_multimodal_qwen3"), + "audio.encoder_type": (audio.get("encoder_type"), "discrete"), + "audio.num_codebooks": (audio.get("num_codebooks"), num_codebooks), + "audio.vocab_size": (audio.get("vocab_size"), codebook_vocab_size), + "audio.out_dim": (audio.get("out_dim"), 2560), + "audio.tie_word_embeddings": (audio.get("tie_word_embeddings"), True), + "audio.use_delay_pattern": (audio.get("use_delay_pattern"), True), + "text.model_type": (text.get("model_type"), "qwen3"), + "text.hidden_size": (text.get("hidden_size"), 2560), + "text.num_hidden_layers": (text.get("num_hidden_layers"), 36), + "text.num_attention_heads": (text.get("num_attention_heads"), 32), + "text.num_key_value_heads": (text.get("num_key_value_heads"), 8), + "text.head_dim": (text.get("head_dim"), 128), + "text.intermediate_size": (text.get("intermediate_size"), 9728), + "text.rms_norm_eps": (text.get("rms_norm_eps"), 1e-6), + "text.vocab_size": (text.get("vocab_size"), 151936), + "text.max_position_embeddings": (text.get("max_position_embeddings"), 32768), + "text.hidden_act": (text.get("hidden_act"), "silu"), + "text.tie_word_embeddings": (text.get("tie_word_embeddings"), True), + "text.attention_dropout": (text.get("attention_dropout"), 0.0), + } + errors = [ + f"{name}={actual!r} (expected {wanted!r})" for name, (actual, wanted) in expected.items() if actual != wanted + ] + architectures = root.get("architectures") + if architectures != ["HiggsMultimodalQwen3ForConditionalGeneration"]: + errors.append("architectures must be ['HiggsMultimodalQwen3ForConditionalGeneration']") + if audio.get("out_dim") != text.get("hidden_size"): + errors.append("audio.out_dim must equal text.hidden_size") + rope_parameters = text.get("rope_parameters") + if not isinstance(rope_parameters, Mapping) or rope_parameters.get("rope_theta") != 1_000_000: + errors.append("text.rope_parameters.rope_theta must be 1000000") + dtype = text.get("dtype") + if dtype not in ("bfloat16", torch.bfloat16): + errors.append(f"text.dtype={dtype!r} (expected bfloat16)") + if errors: + raise ValueError("unsupported Higgs checkpoint configuration: " + "; ".join(errors)) + + +def _as_prompt_tensor(prompt: Any, *, device: torch.device | str | None) -> torch.Tensor: + prompt_tensor = torch.as_tensor(prompt, dtype=torch.long, device=device) + if prompt_tensor.ndim != 1 or prompt_tensor.numel() == 0: + raise ValueError("each Higgs prompt must be a nonempty one-dimensional token sequence") + if bool((prompt_tensor < 0).any()): + raise ValueError("Higgs prompt token IDs must be non-negative") + return prompt_tensor + + +def _higgs_stream(trace: Any) -> Any: + validate = getattr(trace, "validate", None) + if not callable(validate): + raise ValueError("Higgs action traces must provide strict validation") + validate() + if trace.model_family != HIGGS_MODEL_FAMILY: + raise ValueError(f"expected model_family={HIGGS_MODEL_FAMILY!r}, got {trace.model_family!r}") + if len(trace.action_streams) != 1: + raise ValueError("the initial Higgs policy path requires exactly one action stream") + stream = trace.action_streams[0] + if stream.name != HIGGS_STREAM_NAME: + raise ValueError(f"expected action stream {HIGGS_STREAM_NAME!r}, got {stream.name!r}") + if stream.action_type != "multi_discrete" or stream.layout != "time_codebook": + raise ValueError("Higgs actions must use multi_discrete/time_codebook layout") + return stream + + +def collate_higgs_policy_batch( + prompts: Sequence[Any], + action_traces: Sequence[Any], + *, + advantages: Sequence[Any] | None = None, + old_policy_joint_logprobs: Sequence[Any] | None = None, + pad_token_id: int = 0, + device: torch.device | str | None = None, +) -> HiggsPolicyBatch: + """Collate prompt IDs and complete prior codebook rows without packing. + + For action row ``t``, the model reads the prompt plus complete rows + ``[0, t)``. Consequently row zero is predicted from the last prompt + position and row ``t > 0`` from the position containing row ``t - 1``. + Forced BOC/EOC cells remain in ``prior_codes`` even though they are masked + out of the policy loss. Server per-cell logprobs remain in + ``old_cell_logprobs`` for diagnostics; when supplied, the pre-update + Megatron joint logprobs are the GRPO old-policy baseline. + """ + + if len(prompts) == 0 or len(prompts) != len(action_traces): + raise ValueError("prompts and action_traces must have the same nonzero batch size") + if advantages is not None and len(advantages) != len(prompts): + raise ValueError("advantages must have one value or row vector per sample") + if old_policy_joint_logprobs is not None and len(old_policy_joint_logprobs) != len(prompts): + raise ValueError("old-policy joint logprobs must have one row vector per sample") + if type(pad_token_id) is not int or pad_token_id < 0: + raise ValueError("pad_token_id must be a non-negative integer") + + prompt_tensors = [_as_prompt_tensor(prompt, device=device) for prompt in prompts] + streams = [_higgs_stream(trace) for trace in action_traces] + num_codebooks = streams[0].shape[1] + codebook_vocab_size = streams[0].vocab_size + for stream in streams: + if stream.shape[0] <= 0: + raise ValueError("each Higgs action stream must contain at least one row") + if stream.shape[1] != num_codebooks or stream.vocab_size != codebook_vocab_size: + raise ValueError("all Higgs streams in a batch must share codebook shape and vocabulary") + + batch_size = len(prompts) + prompt_lengths = torch.tensor([prompt.numel() for prompt in prompt_tensors], dtype=torch.long, device=device) + action_lengths = torch.tensor([stream.shape[0] for stream in streams], dtype=torch.long, device=device) + max_actions = int(action_lengths.max().item()) + # The final action is a label only; every earlier complete row is an input. + input_lengths = prompt_lengths + action_lengths - 1 + max_input = int(input_lengths.max().item()) + + input_ids = torch.full((batch_size, max_input), pad_token_id, dtype=torch.long, device=device) + prior_codes = torch.zeros((batch_size, max_input, num_codebooks), dtype=torch.long, device=device) + codec_position_mask = torch.zeros((batch_size, max_input), dtype=torch.bool, device=device) + sequence_mask = torch.zeros((batch_size, max_input), dtype=torch.bool, device=device) + prediction_positions = torch.zeros((batch_size, max_actions), dtype=torch.long, device=device) + actions = torch.zeros((batch_size, max_actions, num_codebooks), dtype=torch.long, device=device) + action_mask = torch.zeros((batch_size, max_actions, num_codebooks), dtype=torch.bool, device=device) + old_cell_logprobs = torch.zeros((batch_size, max_actions, num_codebooks), dtype=torch.float32, device=device) + advantage_rows = ( + torch.zeros((batch_size, max_actions), dtype=torch.float32, device=device) if advantages is not None else None + ) + + for batch_index, (prompt, stream) in enumerate(zip(prompt_tensors, streams, strict=True)): + prompt_length = prompt.numel() + action_length = stream.shape[0] + input_length = prompt_length + action_length - 1 + input_ids[batch_index, :prompt_length] = prompt + sequence_mask[batch_index, :input_length] = True + + stream_actions = torch.as_tensor(stream.actions, dtype=torch.long, device=device) + stream_mask = torch.as_tensor(stream.action_mask, dtype=torch.bool, device=device) + stream_logprobs = torch.as_tensor(stream.policy_logprobs, dtype=torch.float32, device=device) + if action_length > 1: + prior_slice = slice(prompt_length, prompt_length + action_length - 1) + prior_codes[batch_index, prior_slice] = stream_actions[:-1] + codec_position_mask[batch_index, prior_slice] = True + + actions[batch_index, :action_length] = stream_actions + action_mask[batch_index, :action_length] = stream_mask + old_cell_logprobs[batch_index, :action_length] = torch.where( + stream_mask, stream_logprobs, torch.zeros_like(stream_logprobs) + ) + prediction_positions[batch_index, :action_length] = torch.arange( + prompt_length - 1, + prompt_length - 1 + action_length, + dtype=torch.long, + device=device, + ) + + if advantage_rows is not None: + advantage = torch.as_tensor(advantages[batch_index], dtype=torch.float32, device=device) + if advantage.ndim == 0 or advantage.numel() == 1: + advantage_rows[batch_index, :action_length] = advantage.reshape(()) + elif advantage.ndim == 1 and advantage.numel() == action_length: + advantage_rows[batch_index, :action_length] = advantage + else: + raise ValueError("each Higgs advantage must be scalar or match its action-row count") + + row_mask = action_mask.any(dim=-1) + old_joint_logprobs = old_cell_logprobs.sum(dim=-1) + if old_policy_joint_logprobs is not None: + old_joint_logprobs.zero_() + for batch_index, (values, action_length) in enumerate( + zip(old_policy_joint_logprobs, action_lengths.tolist(), strict=True) + ): + values = torch.as_tensor(values, dtype=torch.float32, device=device) + if values.ndim != 1 or values.numel() != action_length: + raise ValueError("each old-policy joint logprob vector must match its action-row count") + if not bool(torch.isfinite(values).all()): + raise ValueError("old-policy joint logprobs must be finite") + old_joint_logprobs[batch_index, :action_length] = values + return HiggsPolicyBatch( + input_ids=input_ids, + prior_codes=prior_codes, + codec_position_mask=codec_position_mask, + sequence_mask=sequence_mask, + prediction_positions=prediction_positions, + actions=actions, + action_mask=action_mask, + old_cell_logprobs=old_cell_logprobs, + old_joint_logprobs=old_joint_logprobs, + row_mask=row_mask, + advantages=advantage_rows, + prompt_lengths=prompt_lengths, + action_lengths=action_lengths, + num_codebooks=num_codebooks, + codebook_vocab_size=codebook_vocab_size, + ) + + +def build_higgs_teacher_embeddings( + text_embeddings: torch.Tensor, + codec_weight: torch.Tensor, + prior_codes: torch.Tensor, + codec_position_mask: torch.Tensor, +) -> torch.Tensor: + """Overlay summed, channel-offset codebook embeddings on text embeddings.""" + + if text_embeddings.ndim != 3: + raise ValueError("text_embeddings must have shape [batch, sequence, hidden]") + if prior_codes.ndim != 3 or prior_codes.shape[:2] != text_embeddings.shape[:2]: + raise ValueError("prior_codes must have shape [batch, sequence, codebooks]") + if codec_position_mask.shape != text_embeddings.shape[:2]: + raise ValueError("codec_position_mask must have shape [batch, sequence]") + num_codebooks = prior_codes.shape[-1] + if codec_weight.ndim != 2 or codec_weight.shape[0] % num_codebooks != 0: + raise ValueError("codec_weight rows must be divisible by the number of codebooks") + vocab_size = codec_weight.shape[0] // num_codebooks + active_codes = prior_codes[codec_position_mask] + if active_codes.numel() and (bool((active_codes < 0).any()) or bool((active_codes >= vocab_size).any())): + raise ValueError("prior codebook IDs are outside the codec vocabulary") + + offsets = torch.arange(num_codebooks, device=prior_codes.device, dtype=prior_codes.dtype) * vocab_size + codec_embeddings = F.embedding(prior_codes + offsets, codec_weight).sum(dim=-2) + return torch.where(codec_position_mask.unsqueeze(-1), codec_embeddings, text_embeddings) + + +def selected_higgs_logprobs( + logits: torch.Tensor, + actions: torch.Tensor, + action_mask: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return masked per-cell and joint-row logprobs from ``[B,L,Q,V]`` logits.""" + + if logits.ndim != 4: + raise ValueError("Higgs logits must have shape [batch, time, codebook, vocabulary]") + if actions.shape != logits.shape[:-1] or action_mask.shape != actions.shape: + raise ValueError("Higgs actions and masks must match logits [batch, time, codebook]") + if actions.dtype != torch.long: + actions = actions.long() + if action_mask.dtype != torch.bool: + action_mask = action_mask.bool() + if bool((actions[action_mask] < 0).any()) or bool((actions[action_mask] >= logits.shape[-1]).any()): + raise ValueError("sampled Higgs action is outside the model vocabulary") + + # The policy contract is defined against full-vocabulary fp32 softmax. + logprobs = torch.log_softmax(logits.float(), dim=-1) + selected = logprobs.gather(dim=-1, index=actions.unsqueeze(-1)).squeeze(-1) + selected = torch.where(action_mask, selected, torch.zeros_like(selected)) + if not bool(torch.isfinite(selected[action_mask]).all()): + raise RuntimeError("Higgs model produced a non-finite sampled-action logprob") + row_mask = action_mask.any(dim=-1) + return selected, selected.sum(dim=-1), row_mask + + +def _sum_of_sample_row_means(values: torch.Tensor, row_mask: torch.Tensor) -> torch.Tensor: + counts = row_mask.sum(dim=-1) + if bool((counts == 0).any()): + raise ValueError("every Higgs sample must contain at least one active action row") + masked = torch.where(row_mask, values, torch.zeros_like(values)) + return (masked.sum(dim=-1) / counts.to(values.dtype)).sum() + + +def higgs_joint_policy_loss( + current_joint_logprobs: torch.Tensor, + old_joint_logprobs: torch.Tensor, + advantages: torch.Tensor, + row_mask: torch.Tensor, + *, + eps_clip: float, + eps_clip_high: float | None = None, +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Apply GRPO's clipped surrogate once to each joint multi-codebook row ratio.""" + + if current_joint_logprobs.shape != old_joint_logprobs.shape or row_mask.shape != current_joint_logprobs.shape: + raise ValueError("joint logprobs and row_mask must have identical [batch, time] shapes") + if advantages.ndim == 1: + advantages = advantages.unsqueeze(-1).expand_as(current_joint_logprobs) + if advantages.shape != current_joint_logprobs.shape: + raise ValueError("Higgs advantages must have shape [batch] or [batch, time]") + if eps_clip < 0 or (eps_clip_high is not None and eps_clip_high < 0): + raise ValueError("GRPO clipping thresholds must be non-negative") + eps_clip_high = eps_clip if eps_clip_high is None else eps_clip_high + + zeros = current_joint_logprobs.new_zeros(()) + if not bool(torch.isfinite(current_joint_logprobs[row_mask]).all()): + raise RuntimeError("current Higgs joint logprobs must be finite on active rows") + if not bool(torch.isfinite(old_joint_logprobs[row_mask]).all()): + raise RuntimeError("rollout Higgs joint logprobs must be finite on active rows") + if not bool(torch.isfinite(advantages[row_mask]).all()): + raise RuntimeError("Higgs advantages must be finite on active rows") + log_ratio = torch.where( + row_mask, + current_joint_logprobs - old_joint_logprobs, + zeros, + ) + ratio = torch.exp(log_ratio) + if not bool(torch.isfinite(ratio[row_mask]).all()): + raise RuntimeError("Higgs joint importance ratio overflowed on an active row") + clean_advantages = torch.where(row_mask, advantages, zeros) + unclipped = ratio * clean_advantages + clipped_ratio = ratio.clamp(1.0 - eps_clip, 1.0 + eps_clip_high) + clipped = clipped_ratio * clean_advantages + per_row_loss = -torch.minimum(unclipped, clipped) + clipfrac = ((ratio < 1.0 - eps_clip) | (ratio > 1.0 + eps_clip_high)).to(current_joint_logprobs.dtype) + loss = _sum_of_sample_row_means(per_row_loss, row_mask) + metrics = { + "loss": loss.detach(), + "pg_loss": loss.detach(), + "pg_clipfrac": _sum_of_sample_row_means(clipfrac, row_mask).detach(), + } + return loss, metrics + + +def higgs_policy_loss_from_logits( + logits: torch.Tensor, + batch: HiggsPolicyBatch, + *, + eps_clip: float, + eps_clip_high: float | None = None, +) -> tuple[torch.Tensor, dict[str, torch.Tensor], torch.Tensor]: + if batch.advantages is None: + raise ValueError("Higgs training batches require advantages") + cell_logprobs, joint_logprobs, row_mask = selected_higgs_logprobs(logits, batch.actions, batch.action_mask) + if not torch.equal(row_mask, batch.row_mask): + raise ValueError("model and rollout Higgs row masks disagree") + loss, metrics = higgs_joint_policy_loss( + joint_logprobs, + batch.old_joint_logprobs, + batch.advantages, + row_mask, + eps_clip=eps_clip, + eps_clip_high=eps_clip_high, + ) + abs_diff = (joint_logprobs.detach() - batch.old_joint_logprobs).abs() + metrics["train_old_policy_logprob_abs_diff"] = _sum_of_sample_row_means(abs_diff, row_mask).detach() + return loss, metrics, cell_logprobs + + +def get_higgs_joint_log_probs( + logits: torch.Tensor, + *, + higgs_batch: HiggsPolicyBatch, + non_loss_data: bool = True, + **_: Any, +) -> dict[str, list[torch.Tensor]]: + if not non_loss_data: + raise ValueError("Higgs logprob collection is only valid as non-loss data") + _, joint_logprobs, _ = selected_higgs_logprobs(logits, higgs_batch.actions, higgs_batch.action_mask) + return { + "log_probs": [ + joint_logprobs[index, : int(length.item())] for index, length in enumerate(higgs_batch.action_lengths) + ] + } + + +def validate_higgs_logprob_parity( + action_traces: Sequence[Any], + recomputed_joint_logprobs: Sequence[torch.Tensor], + *, + atol: float, +) -> dict[str, float | bool]: + """Measure active-row parity while hard-gating only malformed/non-finite data.""" + + if isinstance(atol, bool) or not isinstance(atol, (int, float)) or atol <= 0: + raise ValueError("Higgs parity atol must be a positive number") + if len(action_traces) == 0 or len(action_traces) != len(recomputed_joint_logprobs): + raise ValueError("Higgs parity inputs must have the same nonzero sample count") + + diffs = [] + for sample_index, (trace, recomputed) in enumerate(zip(action_traces, recomputed_joint_logprobs, strict=True)): + stream = _higgs_stream(trace) + device = recomputed.device + action_mask = torch.as_tensor(stream.action_mask, dtype=torch.bool, device=device) + old_cell_logprobs = torch.as_tensor(stream.policy_logprobs, dtype=torch.float32, device=device) + row_mask = action_mask.any(dim=-1) + expected = torch.where(action_mask, old_cell_logprobs, 0.0).sum(dim=-1) + recomputed = recomputed.float() + if recomputed.ndim != 1 or recomputed.numel() != stream.shape[0]: + raise ValueError( + f"Higgs parity sample {sample_index} recomputed shape {tuple(recomputed.shape)} " + f"does not match {stream.shape[0]} action rows" + ) + if not bool(torch.isfinite(recomputed[row_mask]).all()): + raise RuntimeError(f"Higgs parity sample {sample_index} contains non-finite trainer logprobs") + diffs.append((recomputed[row_mask] - expected[row_mask]).abs()) + + all_diffs = torch.cat(diffs) + max_abs_diff = float(all_diffs.max().item()) + mean_abs_diff = float(all_diffs.mean().item()) + return { + "max_abs_diff": max_abs_diff, + "mean_abs_diff": mean_abs_diff, + "within_tolerance": max_abs_diff <= float(atol), + } + + +def validate_higgs_weight_versions( + sample_weight_versions: Sequence[Sequence[str]] | None, + *, + trainer_weight_version: Any, +) -> str: + """Reject missing, resumed, mixed, or stale structured trajectories.""" + + if not sample_weight_versions: + raise ValueError("Higgs training requires one rollout weight version per sample") + versions = [] + for sample_index, sample_versions in enumerate(sample_weight_versions): + if not isinstance(sample_versions, (list, tuple)) or len(sample_versions) != 1: + raise ValueError(f"Higgs sample {sample_index} must contain exactly one rollout weight version") + version = sample_versions[0] + if not isinstance(version, str) or not version: + raise ValueError(f"Higgs sample {sample_index} weight version must be a nonempty string") + versions.append(version) + if len(set(versions)) != 1: + raise RuntimeError(f"Higgs training batch mixes rollout weight versions: {sorted(set(versions))}") + trainer_version = str(trainer_weight_version) + # SGLang's untouched startup checkpoint is reported as ``default`` while + # Miles' updater starts its monotonic version counter at zero. This alias + # is valid only before the first update; every later version must match + # exactly. + initial_version_alias = versions[0] == "default" and trainer_version == "0" + if versions[0] != trainer_version and not initial_version_alias: + raise RuntimeError(f"stale Higgs rollout weight version {versions[0]!r}; trainer expects {trainer_version!r}") + return versions[0] + + +__all__ = [ + "HIGGS_MODEL_FAMILY", + "HIGGS_STREAM_NAME", + "HiggsPolicyBatch", + "build_higgs_teacher_embeddings", + "collate_higgs_policy_batch", + "get_higgs_joint_log_probs", + "higgs_joint_policy_loss", + "higgs_policy_loss_from_logits", + "is_higgs_policy_enabled", + "selected_higgs_logprobs", + "validate_higgs_hf_config", + "validate_higgs_single_device_config", + "validate_higgs_logprob_parity", + "validate_higgs_weight_versions", +] diff --git a/miles/backends/training_utils/log_utils.py b/miles/backends/training_utils/log_utils.py index 0ea3b538b13..5a47c1d08f1 100644 --- a/miles/backends/training_utils/log_utils.py +++ b/miles/backends/training_utils/log_utils.py @@ -20,6 +20,67 @@ logger = logging.getLogger(__name__) +def _get_higgs_rollout_log_dict(rollout_data: RolloutBatch) -> dict[str, float]: + """Summarize structured action rows without text-response alignment.""" + + traces = rollout_data["action_traces"] + if not traces: + raise ValueError("structured rollout logging requires at least one action trace") + + row_masks = [] + joint_logprobs = [] + action_rows = [] + active_rows = [] + sampled_cells = [] + for trace in traces: + trace.validate() + if trace.model_family != "higgs_tts" or len(trace.action_streams) != 1: + raise ValueError("structured rollout logging currently supports one Higgs action stream") + stream = trace.action_streams[0] + action_mask = torch.as_tensor(stream.action_mask, dtype=torch.bool) + policy_logprobs = torch.as_tensor(stream.policy_logprobs, dtype=torch.float32) + row_mask = action_mask.any(dim=-1) + if not bool(row_mask.any()): + raise ValueError("each Higgs rollout must contain at least one sampled action row") + row_masks.append(row_mask) + joint_logprobs.append(torch.where(action_mask, policy_logprobs, 0.0).sum(dim=-1)) + action_rows.append(float(stream.shape[0])) + active_rows.append(float(row_mask.sum().item())) + sampled_cells.append(float(action_mask.sum().item())) + + log_dict = { + "action_rows": sum(action_rows) / len(action_rows), + "active_action_rows": sum(active_rows) / len(active_rows), + "sampled_action_cells": sum(sampled_cells) / len(sampled_cells), + "rollout_joint_log_probs": sum( + values[mask].mean().item() for values, mask in zip(joint_logprobs, row_masks, strict=True) + ) + / len(row_masks), + } + + for key in ("advantages", "returns", "log_probs", "ref_log_probs"): + values = rollout_data.get(key) + if values is None: + continue + if len(values) != len(row_masks): + raise ValueError(f"structured rollout field {key!r} has the wrong sample count") + sample_means = [] + for value, mask in zip(values, row_masks, strict=True): + value_tensor = torch.as_tensor(value, dtype=torch.float32, device=mask.device) + if value_tensor.ndim != 1 or value_tensor.numel() != mask.numel(): + raise ValueError(f"structured rollout field {key!r} must align with action rows") + sample_means.append(value_tensor[mask].mean().item()) + log_dict[key] = sum(sample_means) / len(sample_means) + + for key in ("rewards", "raw_reward", "truncated", "response_lengths", "total_lengths"): + values = rollout_data.get(key) + if values is None or len(values) == 0: + continue + if all(isinstance(value, (int, float)) for value in values): + log_dict[key] = sum(float(value) for value in values) / len(values) + return log_dict + + def gather_log_data( metric_name: str, args: Namespace, @@ -104,6 +165,13 @@ def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatc - Scalars are converted to Python numbers. """ parallel_state = get_parallel_state() + if "action_traces" in rollout_data: + if parallel_state.tp.rank == 0 and parallel_state.is_pp_last_stage: + gather_log_data("rollout", args, rollout_id, _get_higgs_rollout_log_dict(rollout_data)) + if args.log_passrate: + log_passrate(rollout_id, args, rollout_data) + return + if parallel_state.tp.rank == 0 and parallel_state.is_pp_last_stage: cp_size = parallel_state.cp.size log_dict = {} diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index c693febebae..c873fc37593 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -4,6 +4,11 @@ from torch.utils.checkpoint import checkpoint from miles.backends.training_utils.cp_utils import get_sum_of_sample_mean +from miles.backends.training_utils.higgs_policy import ( + higgs_policy_loss_from_logits, + is_higgs_policy_enabled, + validate_higgs_single_device_config, +) from miles.backends.training_utils.loss_hub.advantages import compute_advantages, normalize_advantages from miles.backends.training_utils.loss_hub.logit_processors import get_log_probs_and_entropy, get_values # noqa: F401 from miles.backends.training_utils.loss_hub.losses import get_loss_function @@ -35,6 +40,31 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) "total_lengths"). Modified in-place to add "advantages" and "returns" keys, each mapping to lists of tensors per sample. """ + if "action_traces" in rollout_data: + validate_higgs_single_device_config(args, data_parallel_size=get_parallel_state().intra_dp.size) + if not is_higgs_policy_enabled(args): + raise ValueError("structured action traces require the Higgs policy adapter") + traces = rollout_data["action_traces"] + rewards = rollout_data["rewards"] + if len(traces) != len(rewards): + raise ValueError("Higgs action traces and rewards must have identical batch size") + device = rollout_data["tokens"][0].device + advantages = [] + for trace, reward in zip(traces, rewards, strict=True): + trace.validate() + if trace.model_family != "higgs_tts" or len(trace.action_streams) != 1: + raise ValueError("the initial Higgs path requires one higgs_tts action stream") + row_count = trace.action_streams[0].shape[0] + reward_tensor = torch.as_tensor(reward, dtype=torch.float32, device=device) + if reward_tensor.ndim != 0: + raise ValueError("each Higgs rollout reward must be scalar") + if not bool(torch.isfinite(reward_tensor)): + raise ValueError("each Higgs rollout reward must be finite") + advantages.append(reward_tensor.expand(row_count).clone()) + rollout_data["advantages"] = advantages + rollout_data["returns"] = [advantage.clone() for advantage in advantages] + return + log_probs: list[torch.Tensor] = rollout_data.get("rollout_log_probs" if args.use_rollout_logprobs else "log_probs") ref_log_probs: list[torch.Tensor] = rollout_data.get("ref_log_probs") rewards: list[float] = rollout_data.get("rewards") @@ -120,6 +150,33 @@ def loss_function( "values" (1D tensor: [count, metric1, metric2, ...]). """ parallel_state = get_parallel_state() + if "higgs_policy_batch" in batch: + validate_higgs_single_device_config(args, data_parallel_size=parallel_state.intra_dp.size) + higgs_batch = batch["higgs_policy_batch"] + loss, log, _ = higgs_policy_loss_from_logits( + logits, + higgs_batch, + eps_clip=args.eps_clip, + eps_clip_high=args.eps_clip_high, + ) + if apply_megatron_loss_scaling: + assert not args.use_dynamic_global_batch_size + loss = loss * num_microbatches / args.global_batch_size + else: + loss = loss / args.global_batch_size + metric_values = [ + torch.tensor(float(higgs_batch.input_ids.shape[0]), device=logits.device), + *(value.to(device=logits.device) for value in log.values()), + ] + return ( + loss, + torch.tensor(1, device=logits.device), + { + "keys": list(log.keys()), + "values": torch.stack(metric_values), + }, + ) + num_tokens = sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in batch["loss_masks"]]) num_samples = len(batch["response_lengths"]) diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index f729d948aed..daca98cb5b0 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -8,6 +8,25 @@ from miles.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST +def _build_train_actor_env_vars(train_env_vars: dict[str, str]) -> dict[str, str]: + env_vars = { + "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": "1", + # DeepEP/NVSHMEM's internal NCCL conflicts with our NCCL and hangs under CUDA graphs. + "NVSHMEM_DISABLE_NCCL": os.environ.get("NVSHMEM_DISABLE_NCCL", "1"), + **{name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST}, + } + + # Both ranks of a custom NCCL communicator must use the same memory + # registration mode. Preserve an explicit operator choice, but do not + # invent one for training actors when an external rollout server may use + # NCCL's default. + if "NCCL_CUMEM_ENABLE" in os.environ: + env_vars["NCCL_CUMEM_ENABLE"] = os.environ["NCCL_CUMEM_ENABLE"] + + env_vars.update(train_env_vars) + return env_vars + + class RayTrainGroup: """ A group of ray actors @@ -51,16 +70,7 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): assert pg is not None pg, reordered_bundle_indices, _reordered_gpu_ids = pg - env_vars = { - # because sglang will always set NCCL_CUMEM_ENABLE to 0 - # we need also set it to 0 to prevent nccl error. - "NCCL_CUMEM_ENABLE": os.environ.get("NCCL_CUMEM_ENABLE", "0"), - "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": "1", - # DeepEP/NVSHMEM's internal NCCL conflicts with our NCCL and hangs under CUDA graphs. - "NVSHMEM_DISABLE_NCCL": os.environ.get("NVSHMEM_DISABLE_NCCL", "1"), - **{name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST}, - **self.args.train_env_vars, - } + env_vars = _build_train_actor_env_vars(self.args.train_env_vars) if source_patcher_config := self.args.dumper_source_patcher_config_train: env_vars["DUMPER_SOURCE_PATCHER_CONFIG"] = source_patcher_config @@ -137,6 +147,11 @@ async def update_weights(self): await self._broadcast("update_weights", info=info) + async def disconnect_rollout_engines(self): + if self.args.train_backend != "megatron" or self.args.debug_train_only or self.args.debug_rollout_only: + return + await self._broadcast("disconnect_rollout_engines") + async def onload(self): await self._broadcast("wake_up") diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index a86a4541cc3..1983ed6df1f 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -98,6 +98,14 @@ def create_placement_groups(args): if args.use_critic: num_gpus += args.critic_num_nodes * args.critic_num_gpus_per_node critic_offset = args.actor_num_nodes * args.actor_num_gpus_per_node + elif args.rollout_external: + # External servers own their GPUs outside this Ray cluster. Only reserve + # local bundles for trainable models; the HTTP proxy actors are CPU-only. + num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + if args.use_critic: + critic_offset = num_gpus + num_gpus += args.critic_num_nodes * args.critic_num_gpus_per_node + rollout_offset = num_gpus else: num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + args.rollout_num_gpus rollout_offset = args.actor_num_nodes * args.actor_num_gpus_per_node diff --git a/miles/ray/rollout/server_group.py b/miles/ray/rollout/server_group.py index 65b065963e2..2af303040db 100644 --- a/miles/ray/rollout/server_group.py +++ b/miles/ray/rollout/server_group.py @@ -84,17 +84,24 @@ def start_engines( continue global_rank = self.rank_offset + i - num_gpus = 0.2 - num_cpus = num_gpus - - gpu_index = self.gpu_offset + i * num_gpu_per_engine - base_gpu_id = int(reordered_gpu_ids[gpu_index]) - - scheduling_strategy = PlacementGroupSchedulingStrategy( - placement_group=pg, - placement_group_capture_child_tasks=True, - placement_group_bundle_index=reordered_bundle_indices[gpu_index], - ) + num_cpus = 0.2 + actor_options: dict[str, Any] = {"num_cpus": num_cpus} + if self.args.rollout_external: + # The actor only proxies HTTP/admin calls to an already-running + # server. It must not reserve one of the trainer's local GPUs. + base_gpu_id = 0 + actor_options["num_gpus"] = 0 + else: + gpu_index = self.gpu_offset + i * num_gpu_per_engine + base_gpu_id = int(reordered_gpu_ids[gpu_index]) + actor_options.update( + num_gpus=0.2, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_capture_child_tasks=True, + placement_group_bundle_index=reordered_bundle_indices[gpu_index], + ), + ) env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} | { key: os.environ.get(key, default_val) @@ -118,9 +125,7 @@ def start_engines( env_vars.update(dumper_utils.get_sglang_env(self.args)) rollout_engine = RolloutRayActor.options( - num_cpus=num_cpus, - num_gpus=num_gpus, - scheduling_strategy=scheduling_strategy, + **actor_options, runtime_env={ "env_vars": env_vars, }, diff --git a/miles/ray/rollout/train_data_conversion.py b/miles/ray/rollout/train_data_conversion.py index 65bc8d4b6db..62d050b70db 100644 --- a/miles/ray/rollout/train_data_conversion.py +++ b/miles/ray/rollout/train_data_conversion.py @@ -39,6 +39,16 @@ def convert_samples_to_train_data( "sample_indices": [sample.index for sample in samples], } + trace_presence = [sample.action_trace is not None for sample in samples] + if any(trace_presence) and not all(trace_presence): + raise ValueError("cannot mix samples with and without structured action traces") + if trace_presence and all(trace_presence): + action_traces = [] + for sample in samples: + sample.action_trace.validate() + action_traces.append(sample.action_trace) + train_data["action_traces"] = action_traces + # loss mask # TODO: compress the loss mask loss_masks = [] @@ -103,7 +113,10 @@ def _post_process_rewards(args, samples: list[Sample] | list[list[Sample]], cust raw_rewards = [sample.get_reward_value(args) for sample in samples] if args.advantage_estimator in ["grpo", "gspo", "reinforce_plus_plus_baseline"] and args.rewards_normalization: # group norm - rewards = torch.tensor(raw_rewards, dtype=torch.float) + # Center in float64. Float32 reduction can move the mean of an identical + # non-representable reward (for example, eight 0.9 values) by one ULP, + # which the GRPO epsilon then amplifies into a spurious policy signal. + rewards = torch.tensor(raw_rewards, dtype=torch.float64) if rewards.shape[-1] == args.n_samples_per_prompt * args.rollout_batch_size: rewards = rewards.reshape(-1, args.n_samples_per_prompt) else: @@ -152,6 +165,7 @@ def split_train_data_by_dp(args, data, dp_size): "round_number", "sample_indices", "rollout_log_probs", + "action_traces", "rollout_routed_experts", "rollout_indexer_topk", "prompt", diff --git a/miles/rollout/generate_utils/sample_utils.py b/miles/rollout/generate_utils/sample_utils.py index 68eb57e7267..6a515a10f6f 100644 --- a/miles/rollout/generate_utils/sample_utils.py +++ b/miles/rollout/generate_utils/sample_utils.py @@ -22,6 +22,9 @@ def _merge_sample_pair(a: Sample, b: Sample, tokenizer) -> Sample: """Merge two samples generated from sibling inference engine calls.""" a, b = deepcopy(a), deepcopy(b) + if a.action_trace is not None or b.action_trace is not None: + raise ValueError("structured action traces cannot be merged") + def _merge_equal_value(field): x = getattr(a, field) y = getattr(b, field) @@ -118,6 +121,8 @@ def _merge_metadata(): loss_mask=a.loss_mask + [0] * obs_len + b.loss_mask, weight_versions=a.weight_versions + b.weight_versions, rollout_log_probs=a.rollout_log_probs + [0.0] * obs_len + b.rollout_log_probs, + action_trace=None, + decoded_audio=_merge_equal_value("decoded_audio"), teacher_log_probs=_merge_optional_per_token("teacher_log_probs"), opd_reverse_kl=_merge_optional_per_token("opd_reverse_kl"), rollout_routed_experts=b.rollout_routed_experts, diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 7d7858b91fe..f1914803761 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -139,6 +139,34 @@ def add_train_arguments(parser): default="thd", help="The qkv layout.", ) + parser.add_argument( + "--structured-policy-model-family", + choices=["higgs_tts"], + default=None, + help="Enable a typed non-text policy adapter alongside the legacy causal-text path.", + ) + parser.add_argument( + "--higgs-num-codebooks", + type=int, + default=8, + help="Number of channels in the Higgs delayed codebook action stream.", + ) + parser.add_argument( + "--higgs-codebook-vocab-size", + type=int, + default=1026, + help="Per-codebook Higgs vocabulary, including BOC and EOC.", + ) + parser.add_argument( + "--higgs-logprob-parity-atol", + type=float, + default=None, + help=( + "Measured maximum absolute joint-row logprob difference above which Higgs emits a warning. " + "Malformed shapes and non-finite values remain fatal; finite kernel-level drift is logged " + "rather than blocking training." + ), + ) parser.add_argument( "--linear-attention-backend", type=str, @@ -2009,6 +2037,14 @@ def parse_args(add_custom_arguments=None): args.compress_ratios = None if args.hf_checkpoint: hf_config = load_hf_config(args.hf_checkpoint) + if args.structured_policy_model_family == "higgs_tts": + from miles.backends.training_utils.higgs_policy import validate_higgs_hf_config + + validate_higgs_hf_config( + hf_config, + num_codebooks=args.higgs_num_codebooks, + codebook_vocab_size=args.higgs_codebook_vocab_size, + ) args.compress_ratios = getattr(hf_config, "compress_ratios", None) hf_validate_args(args, hf_config) @@ -2266,7 +2302,13 @@ def miles_validate_args(args): args.no_load_optim = True args.no_load_rng = True args.finetune = True - args.load = args.ref_load + if args.structured_policy_model_family == "higgs_tts" and args.ref_load is None: + # The strict raw Higgs checkpoint loader maps the original HF + # checkpoint directly into the TP1 Megatron model on a fresh + # run. Resumes still take the valid Megatron --load path above. + args.load = args.hf_checkpoint + else: + args.load = args.ref_load if args.ref_ckpt_step is not None: args.ckpt_step = args.ref_ckpt_step args.start_rollout_id = 0 @@ -2557,6 +2599,13 @@ def miles_validate_args(args): args.use_dynamic_batch_size is False ), "Dynamic batch size is not supported for bshd format. Please specify --micro-batch-size instead." + if args.structured_policy_model_family == "higgs_tts": + from miles.backends.training_utils.higgs_policy import validate_higgs_single_device_config + + if args.higgs_num_codebooks <= 0 or args.higgs_codebook_vocab_size <= 0: + raise ValueError("Higgs codebook dimensions must be positive") + validate_higgs_single_device_config(args) + _maybe_apply_dumper_overrides(args) diff --git a/miles/utils/hf_config.py b/miles/utils/hf_config.py index cd081798fa2..1a61957665d 100644 --- a/miles/utils/hf_config.py +++ b/miles/utils/hf_config.py @@ -30,6 +30,13 @@ class _HFConfigAlias: _CONFIG_ALIASES: tuple[_HFConfigAlias, ...] = ( + _HFConfigAlias( + model_type="higgs_multimodal_qwen3", + base_module="miles_plugins.models.higgs_tts", + base_class="HiggsMultimodalQwen3Config", + compat_class_name="MilesHiggsMultimodalQwen3Config", + auto_model_classes=(), + ), _HFConfigAlias( model_type="deepseek_v32", base_module="transformers.models.deepseek_v3.configuration_deepseek_v3", @@ -98,6 +105,15 @@ def load_hf_config( """ register_hf_config_aliases() config = AutoConfig.from_pretrained(checkpoint_path, trust_remote_code=trust_remote_code, **autoconfig_kwargs) + from miles_plugins.models.higgs_tts import HiggsMultimodalQwen3Config + + if isinstance(config, HiggsMultimodalQwen3Config) and isinstance(config.text_config, dict): + # Transformers' dynamic compatibility subclass serializes composition + # sub-configs back to dictionaries after construction. Miles' generic + # shape validator expects a normal HF config object. + from miles_plugins.models.higgs_tts import build_higgs_text_config + + config.text_config = build_higgs_text_config(config.text_config) if overrides: for key, value in overrides.items(): diff --git a/miles/utils/types.py b/miles/utils/types.py index cd7637d4c44..ebc697b03b6 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -1,3 +1,4 @@ +import math from dataclasses import dataclass, field from enum import Enum from typing import Any @@ -6,6 +7,245 @@ import torch +def _check_dict_keys(data: dict, *, required: set[str], optional: set[str], type_name: str) -> None: + keys = set(data) + missing = required - keys + extra = keys - required - optional + if missing or extra: + raise ValueError(f"{type_name} fields mismatch; missing={sorted(missing)}, extra={sorted(extra)}") + + +@dataclass +class DiscreteActionStream: + """One aligned time-by-channel multi-discrete policy action stream.""" + + name: str + stage: str + modality: str + shape: list[int] + vocab_size: int + actions: list[list[int]] + policy_logprobs: list[list[float]] + action_mask: list[list[bool]] + channel_ids: list[int] + codec_content_mask: list[list[bool]] | None = None + action_type: str = "multi_discrete" + layout: str = "time_codebook" + + def __post_init__(self) -> None: + self.validate() + + def validate(self) -> None: + for field_name, value in ( + ("name", self.name), + ("stage", self.stage), + ("modality", self.modality), + ): + if not isinstance(value, str) or not value: + raise ValueError(f"{field_name} must be a nonempty string") + if self.action_type != "multi_discrete": + raise ValueError("action_type must be 'multi_discrete'") + if self.layout != "time_codebook": + raise ValueError("layout must be 'time_codebook'") + if not isinstance(self.shape, list) or len(self.shape) != 2 or any(type(x) is not int for x in self.shape): + raise ValueError("shape must be a two-element integer list [time, channels]") + + length, channels = self.shape + if length < 0 or channels <= 0: + raise ValueError("action stream dimensions must be non-negative") + if type(self.vocab_size) is not int or self.vocab_size <= 0: + raise ValueError("vocab_size must be a positive integer") + if ( + not isinstance(self.channel_ids, list) + or any(type(channel_id) is not int for channel_id in self.channel_ids) + or self.channel_ids != list(range(channels)) + ): + raise ValueError("channel_ids must be the ordered channel indices") + + matrices = { + "actions": self.actions, + "policy_logprobs": self.policy_logprobs, + "action_mask": self.action_mask, + } + if self.codec_content_mask is not None: + matrices["codec_content_mask"] = self.codec_content_mask + for matrix_name, matrix in matrices.items(): + if not isinstance(matrix, list) or len(matrix) != length: + raise ValueError(f"{matrix_name} must have declared shape {self.shape}") + if any(not isinstance(row, list) or len(row) != channels for row in matrix): + raise ValueError(f"{matrix_name} must have declared shape {self.shape}") + + for row in range(length): + for channel in range(channels): + action = self.actions[row][channel] + logprob = self.policy_logprobs[row][channel] + sampled = self.action_mask[row][channel] + if type(action) is not int or not 0 <= action < self.vocab_size: + raise ValueError("action is outside the declared vocabulary") + if isinstance(logprob, bool) or not isinstance(logprob, (int, float)): + raise ValueError("policy_logprobs must contain numeric values") + if type(sampled) is not bool: + raise ValueError("action_mask must contain boolean values") + if sampled and not math.isfinite(logprob): + raise ValueError("sampled action has a non-finite policy logprob") + if not sampled and logprob != 0.0: + raise ValueError("forced action policy logprob must be zero") + if self.codec_content_mask is not None and type(self.codec_content_mask[row][channel]) is not bool: + raise ValueError("codec_content_mask must contain boolean values") + + def to_dict(self) -> dict[str, Any]: + self.validate() + return { + "name": self.name, + "stage": self.stage, + "modality": self.modality, + "action_type": self.action_type, + "layout": self.layout, + "shape": list(self.shape), + "vocab_size": self.vocab_size, + "actions": [list(row) for row in self.actions], + "policy_logprobs": [list(row) for row in self.policy_logprobs], + "action_mask": [list(row) for row in self.action_mask], + "codec_content_mask": ( + [list(row) for row in self.codec_content_mask] if self.codec_content_mask is not None else None + ), + "channel_ids": list(self.channel_ids), + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "DiscreteActionStream": + if not isinstance(data, dict): + raise ValueError("DiscreteActionStream must be constructed from a dictionary") + _check_dict_keys( + data, + required={ + "name", + "stage", + "modality", + "action_type", + "layout", + "shape", + "vocab_size", + "actions", + "policy_logprobs", + "action_mask", + "channel_ids", + }, + optional={"codec_content_mask"}, + type_name="DiscreteActionStream", + ) + return cls( + name=data["name"], + stage=data["stage"], + modality=data["modality"], + action_type=data["action_type"], + layout=data["layout"], + shape=data["shape"], + vocab_size=data["vocab_size"], + actions=data["actions"], + policy_logprobs=data["policy_logprobs"], + action_mask=data["action_mask"], + codec_content_mask=data.get("codec_content_mask"), + channel_ids=data["channel_ids"], + ) + + +@dataclass +class RolloutActionTrace: + """Versioned collection of structured policy action streams.""" + + version: int + model_family: str + total_action_count: int + action_streams: list[DiscreteActionStream] + + def __post_init__(self) -> None: + self.validate() + + def validate(self) -> None: + if type(self.version) is not int or self.version != 2: + raise ValueError("only rollout action trace version 2 is supported") + if not isinstance(self.model_family, str) or not self.model_family: + raise ValueError("model_family must be a nonempty string") + if type(self.total_action_count) is not int or self.total_action_count < 0: + raise ValueError("total_action_count must be a non-negative integer") + if not isinstance(self.action_streams, list) or not self.action_streams: + raise ValueError("action_streams must contain at least one stream") + if any(not isinstance(stream, DiscreteActionStream) for stream in self.action_streams): + raise ValueError("action_streams must contain DiscreteActionStream values") + for stream in self.action_streams: + stream.validate() + actual_count = sum( + int(sampled) for stream in self.action_streams for row in stream.action_mask for sampled in row + ) + if actual_count != self.total_action_count: + raise ValueError("total_action_count does not match action masks") + + def to_dict(self) -> dict[str, Any]: + self.validate() + return { + "version": self.version, + "model_family": self.model_family, + "total_action_count": self.total_action_count, + "action_streams": [stream.to_dict() for stream in self.action_streams], + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "RolloutActionTrace": + if not isinstance(data, dict): + raise ValueError("RolloutActionTrace must be constructed from a dictionary") + _check_dict_keys( + data, + required={"version", "model_family", "total_action_count", "action_streams"}, + optional=set(), + type_name="RolloutActionTrace", + ) + if not isinstance(data["action_streams"], list): + raise ValueError("action_streams must be a list") + return cls( + version=data["version"], + model_family=data["model_family"], + total_action_count=data["total_action_count"], + action_streams=[DiscreteActionStream.from_dict(stream) for stream in data["action_streams"]], + ) + + +@dataclass +class DecodedAudio: + """Decoded waveform returned for reward evaluation, not policy training.""" + + data: str + format: str + sample_rate: int + + def __post_init__(self) -> None: + self.validate() + + def validate(self) -> None: + if not isinstance(self.data, str) or not self.data: + raise ValueError("decoded audio data must be a nonempty string") + if self.format != "wav": + raise ValueError("decoded audio format must be 'wav'") + if type(self.sample_rate) is not int or self.sample_rate <= 0: + raise ValueError("decoded audio sample_rate must be a positive integer") + + def to_dict(self) -> dict[str, Any]: + self.validate() + return {"data": self.data, "format": self.format, "sample_rate": self.sample_rate} + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "DecodedAudio": + if not isinstance(data, dict): + raise ValueError("DecodedAudio must be constructed from a dictionary") + _check_dict_keys( + data, + required={"data", "format", "sample_rate"}, + optional=set(), + type_name="DecodedAudio", + ) + return cls(data=data["data"], format=data["format"], sample_rate=data["sample_rate"]) + + @dataclass class Sample: """The sample generated""" @@ -127,11 +367,17 @@ def from_dict(data: dict): prefix_cache_info: PrefixCacheInfo = field(default_factory=PrefixCacheInfo) + # Structured policy outputs are independent of the legacy text-token fields. + action_trace: RolloutActionTrace | None = None + decoded_audio: DecodedAudio | None = None + def to_dict(self): value = self.__dict__.copy() value["status"] = self.status.value value["spec_info"] = self.spec_info.to_dict() value["prefix_cache_info"] = self.prefix_cache_info.to_dict() + value["action_trace"] = self.action_trace.to_dict() if self.action_trace is not None else None + value["decoded_audio"] = self.decoded_audio.to_dict() if self.decoded_audio is not None else None return value @staticmethod @@ -140,6 +386,10 @@ def from_dict(data: dict): data["status"] = Sample.Status(data["status"]) data["spec_info"] = Sample.SpecInfo.from_dict(data.get("spec_info", {})) data["prefix_cache_info"] = Sample.PrefixCacheInfo.from_dict(data.get("prefix_cache_info", {})) + if data.get("action_trace") is not None: + data["action_trace"] = RolloutActionTrace.from_dict(data["action_trace"]) + if data.get("decoded_audio") is not None: + data["decoded_audio"] = DecodedAudio.from_dict(data["decoded_audio"]) field_names = set(Sample.__dataclass_fields__.keys()) init_data = {k: v for k, v in data.items() if k in field_names} @@ -179,6 +429,10 @@ def validate(self): assert ( len(self.opd_reverse_kl) == self.response_length ), f"opd_reverse_kl length ({len(self.opd_reverse_kl)}) != response_length ({self.response_length})" + if self.action_trace is not None: + self.action_trace.validate() + if self.decoded_audio is not None: + self.decoded_audio.validate() if self.rollout_routed_experts is not None: actual = len(self.rollout_routed_experts) expect = len(self.tokens) - 1 @@ -228,6 +482,8 @@ def reset_for_retry(self) -> None: self.loss_mask = None self.weight_versions = [] self.rollout_log_probs = None + self.action_trace = None + self.decoded_audio = None self.rollout_routed_experts = None self.rollout_indexer_topk = None self.status = Sample.Status.ABORTED @@ -280,7 +536,10 @@ class ParamInfo: # A dict-based batch produced along the rollout -> training path # In Megatron backend, several fields are converted to torch.Tensor lists on GPU # before being consumed by data iterators (see megatron_utils.actor._get_rollout_data). -RolloutBatch = dict[str, list[torch.Tensor] | list[int] | list[float] | list[str]] +RolloutBatch = dict[ + str, + list[torch.Tensor] | list[int] | list[float] | list[str] | list[RolloutActionTrace], +] @dataclass diff --git a/miles_plugins/models/higgs_tts.py b/miles_plugins/models/higgs_tts.py new file mode 100644 index 00000000000..0709e37e8a3 --- /dev/null +++ b/miles_plugins/models/higgs_tts.py @@ -0,0 +1,59 @@ +"""Hugging Face config registration for the Higgs v3 discrete TTS policy.""" + +from __future__ import annotations + +from typing import Any + +from transformers import PretrainedConfig, Qwen3Config + + +def build_higgs_text_config(text_config: Qwen3Config | dict[str, Any]) -> Qwen3Config: + """Normalize the v3 checkpoint's Qwen3 RoPE theta.""" + + if isinstance(text_config, Qwen3Config): + rope_parameters = dict(text_config.rope_parameters or {}) + if rope_parameters.get("rope_theta") is None: + rope_parameters["rope_theta"] = 1_000_000 + text_config.rope_parameters = rope_parameters + if vars(text_config).get("rope_theta") is None: + text_config.rope_theta = 1_000_000 + return text_config + values = dict(text_config) + rope_parameters = dict(values.get("rope_parameters") or {}) + if rope_parameters.get("rope_theta") is None: + rope_parameters["rope_theta"] = 1_000_000 + values["rope_parameters"] = rope_parameters + if values.get("rope_theta") is None: + values["rope_theta"] = 1_000_000 + normalized = Qwen3Config(**values) + if vars(normalized).get("rope_theta") is None: + normalized.rope_theta = 1_000_000 + return normalized + + +class HiggsMultimodalQwen3Config(PretrainedConfig): + """Minimal composition config used by Miles' Megatron provider. + + The audio codec implementation is not loaded into the trainer. Keeping its + configuration as a mapping is sufficient to validate and construct the + tied discrete codebook policy. + """ + + model_type = "higgs_multimodal_qwen3" + sub_configs = {"text_config": Qwen3Config} + is_composition = True + + def __init__( + self, + text_config: Qwen3Config | dict[str, Any] | None = None, + audio_encoder_config: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + if text_config is None: + text_config = {} + self.text_config = build_higgs_text_config(text_config) + self.audio_encoder_config = dict(audio_encoder_config or {}) + + +__all__ = ["HiggsMultimodalQwen3Config", "build_higgs_text_config"] diff --git a/miles_plugins/omni/__init__.py b/miles_plugins/omni/__init__.py new file mode 100644 index 00000000000..5bcc6784518 --- /dev/null +++ b/miles_plugins/omni/__init__.py @@ -0,0 +1 @@ +"""Structured rollout integrations for sglang-omni.""" diff --git a/miles_plugins/omni/omni_generate_fn.py b/miles_plugins/omni/omni_generate_fn.py new file mode 100644 index 00000000000..036c576007c --- /dev/null +++ b/miles_plugins/omni/omni_generate_fn.py @@ -0,0 +1,241 @@ +"""Single-turn Higgs audio rollout client for sglang-omni.""" + +from __future__ import annotations + +import os +from typing import Any + +from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput +from miles.utils.http_utils import post +from miles.utils.types import Sample + +from .rollout_contract import parse_higgs_generate_response + +_FILTER_DEFAULTS: dict[str, float] = { + "temperature": 1.0, + "top_p": 1.0, + "min_p": 0.0, + "repetition_penalty": 1.0, +} +_PASSTHROUGH_SAMPLING_KEYS = frozenset({"max_new_tokens", "max_tokens", "seed", "sampling_seed"}) +_MILES_PRESENTATION_KEYS = frozenset( + {"skip_special_tokens", "no_stop_trim", "spaces_between_special_tokens", "stop", "stop_token_ids"} +) +_ZERO_SHOT_SPECIALS = ("<|tts|>", "<|text|>", "<|audio|>") + + +def neutral_higgs_sampling_params(sampling_params: dict[str, Any]) -> dict[str, Any]: + """Validate and normalize the initial unfiltered Higgs RL sampling profile.""" + params = dict(sampling_params) + allowed = set(_FILTER_DEFAULTS) | {"top_k"} | _PASSTHROUGH_SAMPLING_KEYS | _MILES_PRESENTATION_KEYS + unknown = set(params) - allowed + if unknown: + raise ValueError(f"unsupported Higgs rollout sampling parameters: {sorted(unknown)}") + if params.get("stop") not in (None, "", []): + raise ValueError("Higgs RL requires text stop strings to be disabled") + if params.get("stop_token_ids") not in (None, []): + raise ValueError("Higgs RL requires text stop token IDs to be disabled") + + normalized: dict[str, Any] = {} + for name, expected in _FILTER_DEFAULTS.items(): + value = params.get(name, expected) + if value is None: + value = expected + if isinstance(value, bool) or not isinstance(value, (int, float)) or float(value) != expected: + raise ValueError(f"Higgs RL requires {name}={expected}") + normalized[name] = expected + + top_k = params.get("top_k") + if isinstance(top_k, bool) or top_k not in (None, 0, -1): + raise ValueError("Higgs RL requires top_k filtering to be disabled") + + if "seed" in params and "sampling_seed" in params: + raise ValueError("set only one of seed and sampling_seed") + if "max_new_tokens" in params and "max_tokens" in params: + raise ValueError("set only one of max_new_tokens and max_tokens") + for name in ("max_new_tokens", "max_tokens"): + value = params.get(name) + if value is not None: + if type(value) is not int or value <= 0: + raise ValueError(f"Higgs rollout {name} must be a positive integer") + normalized[name] = value + seed = params.get("seed", params.get("sampling_seed")) + if seed is not None: + if type(seed) is not int: + raise ValueError("Higgs rollout seed must be an integer") + normalized["seed"] = seed + return normalized + + +def build_higgs_generate_payload( + prompt_ids: list[int], + sampling_params: dict[str, Any], +) -> dict[str, Any]: + """Build one non-streaming, audio-only structured rollout request.""" + if not prompt_ids or any(type(token_id) is not int for token_id in prompt_ids): + raise ValueError("Higgs rollout prompt_ids must be a nonempty integer list") + + return { + "input_ids": list(prompt_ids), + "sampling_params": neutral_higgs_sampling_params(sampling_params), + "stream": False, + "output_modalities": ["audio"], + "return_logprob": True, + "return_omni_rollout": True, + } + + +def build_zero_shot_higgs_prompt_ids(tokenizer: Any, prompt_text: str) -> list[int]: + """Build the exact Higgs text-to-audio prompt expected by the server.""" + if not isinstance(prompt_text, str) or not prompt_text.strip(): + raise ValueError("Higgs zero-shot TTS requires nonempty prompt text") + + vocab = dict(tokenizer.get_added_vocab()) + missing = [token for token in _ZERO_SHOT_SPECIALS if token not in vocab] + if missing: + raise ValueError(f"tokenizer is missing Higgs TTS specials: {missing}") + + text_ids = list(tokenizer.encode(prompt_text, add_special_tokens=False)) + if any(type(token_id) is not int for token_id in text_ids): + raise ValueError("Higgs tokenizer returned non-integer text token IDs") + return [vocab["<|tts|>"], vocab["<|text|>"], *text_ids, vocab["<|audio|>"]] + + +class OmniGenerateFn: + """Miles custom generate function for a fresh Higgs audio trajectory.""" + + @staticmethod + def add_arguments(parser: Any) -> None: + group = parser.add_argument_group("Higgs TTS reward") + group.add_argument( + "--tts-asr-backend", + choices=("local", "sglang_omni"), + default=os.environ.get("MILES_TTS_ASR_BACKEND", "local"), + ) + group.add_argument( + "--tts-asr-url", + default=os.environ.get("MILES_TTS_ASR_URL", "http://127.0.0.1:8080"), + ) + group.add_argument( + "--tts-asr-model", + default=os.environ.get("MILES_TTS_ASR_MODEL"), + ) + group.add_argument( + "--tts-asr-device", + default=os.environ.get("MILES_TTS_ASR_DEVICE", "cpu"), + ) + group.add_argument( + "--tts-asr-language", + default=os.environ.get("MILES_TTS_ASR_LANGUAGE", "en"), + ) + group.add_argument( + "--tts-asr-concurrency", + type=int, + default=int(os.environ.get("MILES_TTS_ASR_CONCURRENCY", "32")), + ) + group.add_argument( + "--tts-asr-timeout", + type=float, + default=float(os.environ.get("MILES_TTS_ASR_TIMEOUT", "300")), + ) + + async def __call__(self, input: GenerateFnInput) -> GenerateFnOutput: + sample = input.sample + _validate_fresh_sample(sample) + + prompt_ids = _prompt_ids(input) + sampling_params = dict(input.sampling_params) + _set_generation_budget(input.args, sampling_params, len(prompt_ids)) + payload = build_higgs_generate_payload(prompt_ids, sampling_params) + + url = _generate_url(input.args) + response = await post(url, payload) + result = parse_higgs_generate_response( + response, + expected_prompt_tokens=len(prompt_ids), + ) + + # Audio codebooks are a separate action stream, never text response tokens. + sample.tokens = list(prompt_ids) + sample.action_trace = result.action_trace + sample.decoded_audio = result.decoded_audio + sample.weight_versions.append(result.weight_version) + sample.status = Sample.Status.COMPLETED if result.finish_type == "stop" else Sample.Status.TRUNCATED + sample.prefix_cache_info.cached_tokens += result.cached_tokens + sample.prefix_cache_info.total_prompt_tokens += result.prompt_tokens + return GenerateFnOutput(samples=sample) + + +def _validate_fresh_sample(sample: Sample) -> None: + if sample.status not in {Sample.Status.PENDING, Sample.Status.ABORTED}: + raise ValueError("Higgs structured rollouts require a pending sample or a clean retry") + if sample.status is Sample.Status.ABORTED and sample.tokens: + raise ValueError("Higgs structured rollouts cannot resume partial token state") + if sample.response or sample.response_length != 0: + raise ValueError("Higgs structured rollouts do not support partial text responses") + if sample.loss_mask is not None or sample.rollout_log_probs is not None: + raise ValueError("Higgs audio actions must not use text loss/logprob fields") + if sample.action_trace is not None or sample.decoded_audio is not None or sample.weight_versions: + raise ValueError("Higgs structured rollouts require a fresh sample") + if sample.multimodal_inputs: + raise ValueError("Higgs RL currently supports zero-shot text-to-audio only; reference media is not supported") + + +def _generate_url(args: Any) -> str: + values = vars(args) + if values.get("rollout_external", False): + addresses = values.get("rollout_external_engine_addrs") + if not isinstance(addresses, list) or len(addresses) != 1 or not isinstance(addresses[0], str): + raise ValueError("the initial Higgs external rollout path requires exactly one engine address") + return f"http://{addresses[0].rstrip('/')}/generate" + return f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate" + + +def _prompt_ids(input: GenerateFnInput) -> list[int]: + sample = input.sample + canonical_ids = build_zero_shot_higgs_prompt_ids(input.state.tokenizer, _prompt_text(sample)) + if sample.tokens: + prompt_ids = list(sample.tokens) + if prompt_ids != canonical_ids: + raise ValueError("pretokenized Higgs prompt does not match the canonical zero-shot encoding") + else: + prompt_ids = canonical_ids + if not prompt_ids or any(type(token_id) is not int for token_id in prompt_ids): + raise ValueError("tokenization did not produce a nonempty integer prompt") + return prompt_ids + + +def _prompt_text(sample: Sample) -> str: + if isinstance(sample.prompt, str): + return sample.prompt + for message in reversed(sample.prompt): + if message.get("role") == "user" and isinstance(message.get("content"), str): + return message["content"] + raise ValueError("Higgs zero-shot TTS requires a string prompt or a user text message") + + +def _set_generation_budget(args: Any, sampling_params: dict[str, Any], prompt_length: int) -> None: + arg_values = vars(args) + max_tokens = sampling_params.pop("max_tokens", None) + max_new_tokens = sampling_params.get("max_new_tokens", max_tokens) + if max_tokens is not None and "max_new_tokens" in sampling_params: + raise ValueError("set only one of max_new_tokens and max_tokens") + if max_new_tokens is None: + max_new_tokens = arg_values.get("rollout_max_response_len") + if type(max_new_tokens) is not int or max_new_tokens <= 0: + raise ValueError("Higgs rollout max_new_tokens must be a positive integer") + + max_context = arg_values.get("rollout_max_context_len") + if max_context is not None: + if type(max_context) is not int or max_context <= prompt_length: + raise ValueError("Higgs prompt leaves no context budget for an audio rollout") + max_new_tokens = min(max_new_tokens, max_context - prompt_length) + sampling_params["max_new_tokens"] = max_new_tokens + + +__all__ = [ + "OmniGenerateFn", + "build_higgs_generate_payload", + "build_zero_shot_higgs_prompt_ids", + "neutral_higgs_sampling_params", +] diff --git a/miles_plugins/omni/rollout_contract.py b/miles_plugins/omni/rollout_contract.py new file mode 100644 index 00000000000..679da035d85 --- /dev/null +++ b/miles_plugins/omni/rollout_contract.py @@ -0,0 +1,250 @@ +"""Strict Higgs rollout contract for sglang-omni ``POST /generate``.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Literal + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + StrictBool, + StrictFloat, + StrictInt, + StrictStr, + field_validator, + model_validator, +) + +from miles.utils.types import DecodedAudio, DiscreteActionStream, RolloutActionTrace + +HIGGS_ROLLOUT_VERSION = 2 +HIGGS_MODEL_FAMILY = "higgs_tts" +HIGGS_STREAM_NAME = "higgs_codes" +HIGGS_STREAM_STAGE = "tts_engine" +HIGGS_NUM_CODEBOOKS = 8 +HIGGS_CODEBOOK_VOCAB_SIZE = 1026 + + +class _StrictWireModel(BaseModel): + model_config = ConfigDict(extra="forbid", strict=True) + + +class HiggsActionStreamResponse(_StrictWireModel): + name: Literal["higgs_codes"] + stage: Literal["tts_engine"] + modality: Literal["audio"] + action_type: Literal["multi_discrete"] + layout: Literal["time_codebook"] + shape: list[StrictInt] + vocab_size: Literal[1026] + actions: list[list[StrictInt]] + policy_logprobs: list[list[StrictFloat]] + action_mask: list[list[StrictBool]] + codec_content_mask: list[list[StrictBool]] | None = None + channel_ids: list[StrictInt] + + @model_validator(mode="after") + def validate_lattice(self) -> HiggsActionStreamResponse: + if len(self.shape) != 2: + raise ValueError("Higgs action stream shape must be [time, codebooks]") + length, codebooks = self.shape + if length <= 0: + raise ValueError("Higgs action stream must contain at least one row") + if codebooks != HIGGS_NUM_CODEBOOKS: + raise ValueError(f"Higgs action stream must contain {HIGGS_NUM_CODEBOOKS} codebooks") + if self.channel_ids != list(range(HIGGS_NUM_CODEBOOKS)): + raise ValueError("Higgs channel_ids must be the ordered codebook indices") + + matrices: dict[str, list[list[Any]]] = { + "actions": self.actions, + "policy_logprobs": self.policy_logprobs, + "action_mask": self.action_mask, + } + if self.codec_content_mask is not None: + matrices["codec_content_mask"] = self.codec_content_mask + for name, matrix in matrices.items(): + if len(matrix) != length or any(len(row) != codebooks for row in matrix): + raise ValueError(f"{name} must have declared shape {self.shape}") + + for row in range(length): + for codebook in range(codebooks): + action = self.actions[row][codebook] + logprob = self.policy_logprobs[row][codebook] + sampled = self.action_mask[row][codebook] + if not 0 <= action < HIGGS_CODEBOOK_VOCAB_SIZE: + raise ValueError("Higgs action is outside the codebook vocabulary") + if sampled and not math.isfinite(logprob): + raise ValueError("sampled Higgs action has a non-finite policy logprob") + if not sampled and logprob != 0.0: + raise ValueError("forced Higgs action policy logprob must be zero") + return self + + def to_domain(self) -> DiscreteActionStream: + return DiscreteActionStream( + name=self.name, + stage=self.stage, + modality=self.modality, + action_type=self.action_type, + layout=self.layout, + shape=list(self.shape), + vocab_size=self.vocab_size, + actions=[list(row) for row in self.actions], + policy_logprobs=[list(row) for row in self.policy_logprobs], + action_mask=[list(row) for row in self.action_mask], + codec_content_mask=( + [list(row) for row in self.codec_content_mask] if self.codec_content_mask is not None else None + ), + channel_ids=list(self.channel_ids), + ) + + +class HiggsRolloutTraceResponse(_StrictWireModel): + version: Literal[2] + model_family: Literal["higgs_tts"] + total_action_count: StrictInt = Field(ge=1) + action_streams: list[HiggsActionStreamResponse] = Field(min_length=1, max_length=1) + + @model_validator(mode="after") + def validate_action_count(self) -> HiggsRolloutTraceResponse: + count = sum(int(sampled) for stream in self.action_streams for row in stream.action_mask for sampled in row) + if count != self.total_action_count: + raise ValueError("total_action_count does not match the Higgs action mask") + return self + + def to_domain(self) -> RolloutActionTrace: + return RolloutActionTrace( + version=self.version, + model_family=self.model_family, + total_action_count=self.total_action_count, + action_streams=[stream.to_domain() for stream in self.action_streams], + ) + + +class _FinishReasonResponse(_StrictWireModel): + type: Literal["stop", "length"] + length: StrictInt | None = Field(default=None, ge=0) + + @model_validator(mode="after") + def validate_length(self) -> _FinishReasonResponse: + if self.type == "length" and self.length is None: + raise ValueError("length finish_reason requires its emitted row count") + if self.type != "length" and self.length is not None: + raise ValueError("only a length finish_reason may carry a length") + return self + + +class _DecodedAudioResponse(_StrictWireModel): + data: StrictStr = Field(min_length=1) + path: None = None + format: Literal["wav"] + sample_rate: StrictInt = Field(gt=0) + + @field_validator("data") + @classmethod + def validate_nonblank_data(cls, data: str) -> str: + if not data.strip(): + raise ValueError("decoded WAV data must not be blank") + return data + + def to_domain(self) -> DecodedAudio: + return DecodedAudio(data=self.data, format=self.format, sample_rate=self.sample_rate) + + +class _MetaInfoResponse(_StrictWireModel): + finish_reason: _FinishReasonResponse + prompt_tokens: StrictInt = Field(ge=0) + completion_tokens: StrictInt = Field(gt=0) + cached_tokens: StrictInt = Field(ge=0) + weight_version: StrictStr = Field(min_length=1) + request_metadata: None = None + output_token_logprobs: list[list[StrictFloat | StrictInt]] | None = None + output_codebook_tokens: list[list[StrictInt]] | None = None + omni_rollout: HiggsRolloutTraceResponse + + @field_validator("weight_version") + @classmethod + def validate_nonblank_weight_version(cls, weight_version: str) -> str: + if not weight_version.strip(): + raise ValueError("weight_version must not be blank") + return weight_version + + @model_validator(mode="after") + def validate_codebook_zero_diagnostics(self) -> _MetaInfoResponse: + diagnostics = self.output_token_logprobs + if diagnostics is None: + return self + stream = self.omni_rollout.action_streams[0] + if len(diagnostics) != stream.shape[0]: + raise ValueError("output_token_logprobs length does not match the Higgs action row count") + for row, item in enumerate(diagnostics): + if len(item) != 2: + raise ValueError("output_token_logprobs entries must be [logprob, token_id]") + logprob, token_id = item + if isinstance(logprob, bool) or not isinstance(logprob, (int, float)) or not math.isfinite(logprob): + raise ValueError("output_token_logprobs contains a non-finite diagnostic logprob") + if type(token_id) is not int or token_id != stream.actions[row][0]: + raise ValueError("output_token_logprobs token does not match the codebook-0 action") + return self + + +class HiggsGenerateResponse(_StrictWireModel): + text: Literal[""] + audio: _DecodedAudioResponse + meta_info: _MetaInfoResponse + + @model_validator(mode="after") + def validate_compatibility_fields(self) -> HiggsGenerateResponse: + stream = self.meta_info.omni_rollout.action_streams[0] + if self.meta_info.completion_tokens != stream.shape[0]: + raise ValueError("completion_tokens does not match the Higgs action row count") + compatibility_codes = self.meta_info.output_codebook_tokens + if compatibility_codes is not None and compatibility_codes != stream.actions: + raise ValueError("output_codebook_tokens does not match the structured Higgs actions") + if self.meta_info.finish_reason.type == "length" and self.meta_info.finish_reason.length != stream.shape[0]: + raise ValueError("finish_reason.length does not match the Higgs action row count") + return self + + +@dataclass(frozen=True) +class HiggsRolloutResult: + action_trace: RolloutActionTrace + decoded_audio: DecodedAudio + finish_type: Literal["stop", "length"] + weight_version: str + prompt_tokens: int + cached_tokens: int + + +def parse_higgs_generate_response( + response: Any, + *, + expected_prompt_tokens: int | None = None, +) -> HiggsRolloutResult: + """Parse one complete Higgs rollout and return domain objects plus status data.""" + parsed = HiggsGenerateResponse.model_validate(response, strict=True) + if expected_prompt_tokens is not None and parsed.meta_info.prompt_tokens != expected_prompt_tokens: + raise ValueError("prompt_tokens does not match the exact prompt IDs sent to the inference server") + return HiggsRolloutResult( + action_trace=parsed.meta_info.omni_rollout.to_domain(), + decoded_audio=parsed.audio.to_domain(), + finish_type=parsed.meta_info.finish_reason.type, + weight_version=parsed.meta_info.weight_version, + prompt_tokens=parsed.meta_info.prompt_tokens, + cached_tokens=parsed.meta_info.cached_tokens, + ) + + +__all__ = [ + "HIGGS_CODEBOOK_VOCAB_SIZE", + "HIGGS_MODEL_FAMILY", + "HIGGS_NUM_CODEBOOKS", + "HIGGS_ROLLOUT_VERSION", + "HIGGS_STREAM_NAME", + "HIGGS_STREAM_STAGE", + "HiggsGenerateResponse", + "HiggsRolloutResult", + "parse_higgs_generate_response", +] diff --git a/miles_plugins/omni/tts_reward.py b/miles_plugins/omni/tts_reward.py new file mode 100644 index 00000000000..c4e06dbc19c --- /dev/null +++ b/miles_plugins/omni/tts_reward.py @@ -0,0 +1,315 @@ +"""ASR round-trip reward for decoded Higgs WAV output.""" + +from __future__ import annotations + +import asyncio +import base64 +import binascii +import io +import math +import os +import re +import wave +from dataclasses import dataclass, field +from typing import Any + +import aiohttp +import numpy as np + +from miles.utils.types import DecodedAudio, Sample + +INVALID_AUDIO_REWARD = -1.0 +_TEXT_CHARS = re.compile(r"[^\w]+", flags=re.UNICODE) + + +def normalize_asr_text(text: str) -> str: + return _TEXT_CHARS.sub("", text.casefold()) + + +def character_error_rate(reference: str, hypothesis: str) -> float: + reference = normalize_asr_text(reference) + hypothesis = normalize_asr_text(hypothesis) + if not reference: + return 0.0 if not hypothesis else 1.0 + + previous = list(range(len(hypothesis) + 1)) + for ref_index, reference_char in enumerate(reference, start=1): + current = [ref_index] + for hyp_index, hypothesis_char in enumerate(hypothesis, start=1): + current.append( + min( + previous[hyp_index] + 1, + current[hyp_index - 1] + 1, + previous[hyp_index - 1] + (reference_char != hypothesis_char), + ) + ) + previous = current + return min(1.0, previous[-1] / len(reference)) + + +def decode_wav(audio: DecodedAudio) -> tuple[np.ndarray, int]: + """Decode the typed server artifact and verify its declared sample rate.""" + raw = _decode_audio_bytes(audio) + try: + with wave.open(io.BytesIO(raw), "rb") as wav_file: + if wav_file.getsampwidth() != 2: + raise ValueError("decoded WAV must use 16-bit PCM") + sample_rate = wav_file.getframerate() + channels = wav_file.getnchannels() + frame_count = wav_file.getnframes() + pcm = np.frombuffer(wav_file.readframes(frame_count), dtype=" bytes: + encoded = audio.data.split(",", 1)[1] if audio.data.startswith("data:") and "," in audio.data else audio.data + try: + raw = base64.b64decode(encoded, validate=True) + except (binascii.Error, ValueError) as error: + raise ValueError("decoded audio is not valid base64") from error + if not raw: + raise ValueError("decoded WAV is empty") + return raw + + +def _validate_audio( + audio: DecodedAudio, target_text: str, reward: TtsRoundTripReward +) -> tuple[np.ndarray, int] | None: + try: + waveform, sample_rate = decode_wav(audio) + except ValueError: + return None + duration = waveform.size / sample_rate + if not reward.min_duration_seconds <= duration <= reward.max_duration_seconds: + return None + if not bool(np.isfinite(waveform).all()): + return None + rms = float(np.sqrt(np.mean(np.square(waveform, dtype=np.float64)))) + if not math.isfinite(rms) or rms < reward.silence_rms_floor: + return None + clipped_fraction = float(np.mean(np.abs(waveform) >= (32767.0 / 32768.0))) + if clipped_fraction > reward.max_clipped_fraction: + return None + if not normalize_asr_text(target_text): + return None + return waveform, sample_rate + + +@dataclass +class TtsRoundTripReward: + asr_model_path: str = field(default_factory=lambda: os.environ.get("MILES_TTS_ASR_MODEL", "openai/whisper-base")) + device: str = field(default_factory=lambda: os.environ.get("MILES_TTS_ASR_DEVICE", "cpu")) + min_duration_seconds: float = 0.3 + max_duration_seconds: float = 30.0 + silence_rms_floor: float = 1e-3 + max_clipped_fraction: float = 0.01 + + _model: Any = field(default=None, init=False, repr=False) + _processor: Any = field(default=None, init=False, repr=False) + + def _load_asr(self) -> None: + if self._model is not None: + return + import torch + from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor + + dtype = torch.float16 if self.device.startswith("cuda") else torch.float32 + self._processor = AutoProcessor.from_pretrained(self.asr_model_path) + self._model = AutoModelForSpeechSeq2Seq.from_pretrained( + self.asr_model_path, + torch_dtype=dtype, + ).to(self.device) + self._model.eval() + + def transcribe(self, waveform: np.ndarray, sample_rate: int) -> str: + import torch + + self._load_asr() + if sample_rate != 16000: + target_length = max(1, round(waveform.size * 16000 / sample_rate)) + source_positions = np.arange(waveform.size, dtype=np.float64) + target_positions = np.linspace(0, waveform.size - 1, target_length) + waveform = np.interp(target_positions, source_positions, waveform).astype(np.float32) + model_inputs = self._processor( + waveform, + sampling_rate=16000, + return_tensors="pt", + return_attention_mask=True, + ) + features = model_inputs.input_features.to(self.device, dtype=self._model.dtype) + attention_mask = model_inputs.attention_mask.to(self.device) + with torch.no_grad(): + token_ids = self._model.generate( + features, + attention_mask=attention_mask, + max_new_tokens=128, + task="transcribe", + ) + return self._processor.batch_decode(token_ids, skip_special_tokens=True)[0] + + def score(self, audio: DecodedAudio, target_text: str) -> float: + validated = _validate_audio(audio, target_text, self) + if validated is None: + return INVALID_AUDIO_REWARD + waveform, sample_rate = validated + + # Model loading and inference failures are infrastructure errors, not bad samples. + transcript = self.transcribe(waveform, sample_rate) + reward = 1.0 - character_error_rate(target_text, transcript) + return float(min(1.0, max(0.0, reward))) + + +@dataclass +class SglangOmniASRReward(TtsRoundTripReward): + """Round-trip reward backed by concurrent OpenAI-compatible ASR requests.""" + + base_url: str = field(default_factory=lambda: os.environ.get("MILES_TTS_ASR_URL", "http://127.0.0.1:8080")) + asr_model_path: str = field(default_factory=lambda: os.environ.get("MILES_TTS_ASR_MODEL", "Qwen/Qwen3-ASR-1.7B")) + language: str = field(default_factory=lambda: os.environ.get("MILES_TTS_ASR_LANGUAGE", "en")) + concurrency: int = field(default_factory=lambda: int(os.environ.get("MILES_TTS_ASR_CONCURRENCY", "32"))) + timeout_seconds: float = field(default_factory=lambda: float(os.environ.get("MILES_TTS_ASR_TIMEOUT", "300"))) + + @property + def transcription_url(self) -> str: + base_url = self.base_url.rstrip("/") + if base_url.endswith("/v1/audio/transcriptions"): + return base_url + return f"{base_url}/v1/audio/transcriptions" + + async def score_batch(self, items: list[tuple[DecodedAudio, str]]) -> list[float]: + if self.concurrency <= 0: + raise ValueError("TTS ASR concurrency must be positive") + + rewards = [INVALID_AUDIO_REWARD] * len(items) + valid: list[tuple[int, DecodedAudio, str]] = [] + for index, (audio, target_text) in enumerate(items): + if _validate_audio(audio, target_text, self) is not None: + valid.append((index, audio, target_text)) + if not valid: + return rewards + + semaphore = asyncio.Semaphore(self.concurrency) + timeout = aiohttp.ClientTimeout(total=self.timeout_seconds) + connector = aiohttp.TCPConnector(limit=self.concurrency) + async with aiohttp.ClientSession(timeout=timeout, connector=connector, trust_env=False) as session: + + async def score_one(index: int, audio: DecodedAudio, target_text: str) -> tuple[int, float]: + form = aiohttp.FormData() + form.add_field("model", self.asr_model_path) + form.add_field("language", self.language) + form.add_field("response_format", "json") + form.add_field( + "file", + _decode_audio_bytes(audio), + filename=f"rollout-{index}.wav", + content_type="audio/wav", + ) + async with semaphore, session.post(self.transcription_url, data=form) as response: + if response.status >= 400: + body = await response.text() + raise RuntimeError(f"ASR request failed with HTTP {response.status}: {body[:500]}") + payload = await response.json() + transcript = payload.get("text") if isinstance(payload, dict) else None + if not isinstance(transcript, str): + raise RuntimeError("ASR response must contain a string 'text' field") + reward = 1.0 - character_error_rate(target_text, transcript) + return index, float(min(1.0, max(0.0, reward))) + + results = await asyncio.gather(*(score_one(*item) for item in valid)) + for index, reward in results: + rewards[index] = reward + return rewards + + +_SHARED_REWARD: TtsRoundTripReward | SglangOmniASRReward | None = None + + +def _reward_model(args: Any) -> TtsRoundTripReward | SglangOmniASRReward: + global _SHARED_REWARD + if _SHARED_REWARD is None: + values = vars(args) + backend = values.get("tts_asr_backend", os.environ.get("MILES_TTS_ASR_BACKEND", "local")) + if backend == "sglang_omni": + _SHARED_REWARD = SglangOmniASRReward( + base_url=values.get("tts_asr_url", os.environ.get("MILES_TTS_ASR_URL", "http://127.0.0.1:8080")), + asr_model_path=values.get("tts_asr_model") + or os.environ.get("MILES_TTS_ASR_MODEL", "Qwen/Qwen3-ASR-1.7B"), + language=values.get("tts_asr_language", os.environ.get("MILES_TTS_ASR_LANGUAGE", "en")), + concurrency=values.get("tts_asr_concurrency", int(os.environ.get("MILES_TTS_ASR_CONCURRENCY", "32"))), + timeout_seconds=values.get("tts_asr_timeout", float(os.environ.get("MILES_TTS_ASR_TIMEOUT", "300"))), + ) + elif backend == "local": + _SHARED_REWARD = TtsRoundTripReward( + asr_model_path=values.get("tts_asr_model") + or os.environ.get("MILES_TTS_ASR_MODEL", "openai/whisper-base"), + device=values.get("tts_asr_device", os.environ.get("MILES_TTS_ASR_DEVICE", "cpu")), + ) + else: + raise ValueError(f"unsupported TTS ASR backend: {backend!r}") + return _SHARED_REWARD + + +def _target_text(sample: Sample) -> str: + if isinstance(sample.prompt, str): + return sample.prompt + for message in reversed(sample.prompt): + if message.get("role") == "user" and isinstance(message.get("content"), str): + return message["content"] + return "" + + +def _score_and_release(reward_model: TtsRoundTripReward, sample: Sample) -> float: + try: + if sample.decoded_audio is None: + return INVALID_AUDIO_REWARD + return reward_model.score(sample.decoded_audio, _target_text(sample)) + finally: + sample.decoded_audio = None + + +async def compute_tts_reward( + args: Any, + sample: Sample | list[Sample], + **_: Any, +) -> float | list[float]: + """Miles custom reward hook; decoded waveform bytes never enter training data.""" + reward_model = _reward_model(args) + if isinstance(reward_model, SglangOmniASRReward): + samples = sample if isinstance(sample, list) else [sample] + try: + items = [(item.decoded_audio, _target_text(item)) for item in samples if item.decoded_audio is not None] + valid_indices = [index for index, item in enumerate(samples) if item.decoded_audio is not None] + scored = await reward_model.score_batch(items) + rewards = [INVALID_AUDIO_REWARD] * len(samples) + for index, reward in zip(valid_indices, scored, strict=True): + rewards[index] = reward + return rewards if isinstance(sample, list) else rewards[0] + finally: + for item in samples: + item.decoded_audio = None + if isinstance(sample, list): + try: + return [_score_and_release(reward_model, item) for item in sample] + finally: + for item in sample: + item.decoded_audio = None + return _score_and_release(reward_model, sample) + + +__all__ = [ + "INVALID_AUDIO_REWARD", + "SglangOmniASRReward", + "TtsRoundTripReward", + "character_error_rate", + "compute_tts_reward", + "decode_wav", + "normalize_asr_text", +] diff --git a/scripts/omni/serve-qwen3-asr.sh b/scripts/omni/serve-qwen3-asr.sh new file mode 100755 index 00000000000..7dee371f917 --- /dev/null +++ b/scripts/omni/serve-qwen3-asr.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +set -euo pipefail + +MODEL_PATH="${MILES_TTS_ASR_MODEL:-Qwen/Qwen3-ASR-1.7B}" +HOST="${MILES_TTS_ASR_HOST:-127.0.0.1}" +PORT="${MILES_TTS_ASR_PORT:-8080}" +GPU="${MILES_TTS_ASR_GPU:-0}" +MEM_FRACTION="${MILES_TTS_ASR_MEM_FRACTION:-0.1}" +MAX_RUNNING_REQUESTS="${MILES_TTS_ASR_MAX_RUNNING_REQUESTS:-16}" + +if [[ -n "${SGLANG_OMNI_ROOT:-}" ]]; then + export PYTHONPATH="${SGLANG_OMNI_ROOT}:${PYTHONPATH:-}" +fi + +export CUDA_VISIBLE_DEVICES="${GPU}" +exec python -m sglang_omni.cli serve \ + --model-path "${MODEL_PATH}" \ + --model-name "${MODEL_PATH}" \ + --host "${HOST}" \ + --port "${PORT}" \ + --mem-fraction-static "${MEM_FRACTION}" \ + --max-running-requests "${MAX_RUNNING_REQUESTS}" diff --git a/tests/fast/backends/megatron_utils/test_broadcast_weight_update.py b/tests/fast/backends/megatron_utils/test_broadcast_weight_update.py new file mode 100644 index 00000000000..87af8285677 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_broadcast_weight_update.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +import pytest +import torch + +from miles.backends.megatron_utils.update_weight.update_weight_from_distributed import broadcast + + +def test_broadcast_materializes_noncontiguous_converted_views(monkeypatch) -> None: + base = torch.arange(12).reshape(3, 4) + converted = [("weight", base.transpose(0, 1))] + received: dict = {} + + class Engine: + class Method: + @staticmethod + def remote(**kwargs): + received.update(kwargs) + return "ref" + + update_weights_from_distributed = Method() + + def fake_broadcast(tensor, src, *, group): + assert tensor.is_contiguous() + received["tensor"] = tensor + + monkeypatch.setattr(broadcast.dist, "broadcast", fake_broadcast) + + refs = broadcast.update_weights_from_distributed( + "group", + object(), + 3, + [Engine()], + converted, + ) + + assert refs == ["ref"] + assert received["names"] == ["weight"] + assert received["shapes"] == [torch.Size([4, 3])] + assert received["tensor"].is_contiguous() + assert torch.equal(received["tensor"], base.transpose(0, 1)) + + +def test_disconnect_releases_custom_group_once(monkeypatch) -> None: + updater = object.__new__(broadcast.UpdateWeightFromDistributed) + updater.args = object() + updater._group_name = "miles-pp_0" + updater._model_update_groups = object() + updater.rollout_engines = [object()] + calls = [] + + monkeypatch.setattr( + broadcast.UpdateWeightFromDistributed, + "_is_source", + property(lambda _self: True), + ) + monkeypatch.setattr( + broadcast, + "disconnect_rollout_engines_from_distributed", + lambda *args: calls.append(args), + ) + + updater.disconnect_rollout_engines() + updater.disconnect_rollout_engines() + + assert len(calls) == 1 + assert calls[0][1] == "miles-pp_0" + assert updater._model_update_groups is None + + +def test_weight_update_transport_matches_trainer(monkeypatch) -> None: + monkeypatch.delenv("NCCL_CUMEM_ENABLE", raising=False) + monkeypatch.setattr(broadcast.torch.cuda.nccl, "version", lambda: (2, 28, 9)) + + broadcast._validate_distributed_weight_update_transports( + [ + { + "protocol_version": 1, + "backend": "nccl", + "nccl_version": "2.28.9", + "nccl_cumem_enable": "default", + } + ] + ) + + +def test_weight_update_transport_rejects_cumem_mismatch(monkeypatch) -> None: + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "0") + monkeypatch.setattr(broadcast.torch.cuda.nccl, "version", lambda: (2, 28, 9)) + + with pytest.raises(RuntimeError, match="transport mismatch"): + broadcast._validate_distributed_weight_update_transports( + [ + { + "protocol_version": 1, + "backend": "nccl", + "nccl_version": "2.28.9", + "nccl_cumem_enable": "default", + } + ] + ) + + +def test_weight_update_transport_allows_legacy_engines() -> None: + broadcast._validate_distributed_weight_update_transports([None, None]) diff --git a/tests/fast/backends/megatron_utils/test_higgs_checkpoint.py b/tests/fast/backends/megatron_utils/test_higgs_checkpoint.py new file mode 100644 index 00000000000..ae87f5485ff --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_higgs_checkpoint.py @@ -0,0 +1,376 @@ +import importlib.util +import json +import sys +import types +from argparse import Namespace +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file + +from tests.ci.ci_register import register_cpu_ci + +from miles.backends.megatron_utils import higgs_checkpoint + +register_cpu_ci(est_time=10, suite="stage-a-cpu", labels=[]) + +_CACHED_HIGGS_REVISION = Path( + "/root/.cache/huggingface/hub/models--bosonai--higgs-audio-v3-tts-4b/" + "snapshots/7556c17e05201fccd9c8cc120bc216dcc7b5d561" +) + + +def _load_higgs_converter(): + package_name = "_miles_higgs_converter_test" + package = types.ModuleType(package_name) + package.__path__ = [] + sys.modules[package_name] = package + base = Path(__file__).resolve().parents[4] / "miles" / "backends" / "megatron_utils" / "megatron_to_hf" + for module_name in ("qwen2", "higgs_tts"): + qualified_name = f"{package_name}.{module_name}" + spec = importlib.util.spec_from_file_location(qualified_name, base / f"{module_name}.py") + module = importlib.util.module_from_spec(spec) + sys.modules[qualified_name] = module + spec.loader.exec_module(module) + return sys.modules[f"{package_name}.higgs_tts"].convert_higgs_to_hf + + +@pytest.fixture(scope="module") +def convert_higgs_to_hf(): + return _load_higgs_converter() + + +def _converter_args() -> Namespace: + return Namespace( + hidden_size=4, + kv_channels=1, + num_attention_heads=4, + num_query_groups=2, + higgs_num_codebooks=2, + higgs_codebook_vocab_size=3, + ) + + +def test_higgs_converter_emits_canonical_embedding_names(convert_higgs_to_hf): + args = _converter_args() + text_embedding = torch.arange(28, dtype=torch.float32).reshape(7, 4) + codec_embedding = torch.arange(24, dtype=torch.float32).reshape(6, 4) + + assert convert_higgs_to_hf( + args, + "module.module.embedding.word_embeddings.weight", + text_embedding, + ) == [("tied.embedding.text_embedding.weight", text_embedding)] + assert convert_higgs_to_hf( + args, + "module.module.codec_embeddings.weight", + codec_embedding, + ) == [("tied.embedding.modality_embeddings.0.embedding.weight", codec_embedding)] + + +def test_higgs_converter_round_trips_grouped_qkv_and_fc1(convert_higgs_to_hf): + args = _converter_args() + q = torch.arange(16, dtype=torch.float32).reshape(4, 4) + k = torch.arange(8, dtype=torch.float32).reshape(2, 4) + 100 + v = torch.arange(8, dtype=torch.float32).reshape(2, 4) + 200 + fused_qkv = torch.cat( + ( + q.view(2, 2, 1, 4), + k.view(2, 1, 1, 4), + v.view(2, 1, 1, 4), + ), + dim=1, + ).reshape(8, 4) + + converted_qkv = convert_higgs_to_hf( + args, + "module.module.decoder.layers.3.self_attention.linear_qkv.weight", + fused_qkv, + ) + assert [name for name, _ in converted_qkv] == [ + "body.layers.3.self_attn.q_proj.weight", + "body.layers.3.self_attn.k_proj.weight", + "body.layers.3.self_attn.v_proj.weight", + ] + assert torch.equal(converted_qkv[0][1], q) + assert torch.equal(converted_qkv[1][1], k) + assert torch.equal(converted_qkv[2][1], v) + + gate = torch.arange(20, dtype=torch.float32).reshape(5, 4) + up = gate + 100 + converted_fc1 = convert_higgs_to_hf( + args, + "module.module.decoder.layers.3.mlp.linear_fc1.weight", + torch.cat((gate, up), dim=0), + ) + assert [name for name, _ in converted_fc1] == [ + "body.layers.3.mlp.gate_proj.weight", + "body.layers.3.mlp.up_proj.weight", + ] + assert torch.equal(converted_fc1[0][1], gate) + assert torch.equal(converted_fc1[1][1], up) + + +def test_higgs_converter_handles_both_norm_layouts(convert_higgs_to_hf): + args = _converter_args() + norm = torch.ones(4) + mappings = { + "module.module.decoder.layers.1.self_attention.linear_qkv.layer_norm_weight": ( + "body.layers.1.input_layernorm.weight" + ), + "module.module.decoder.layers.1.input_layernorm.weight": "body.layers.1.input_layernorm.weight", + "module.module.decoder.layers.1.mlp.linear_fc1.layer_norm_weight": ( + "body.layers.1.post_attention_layernorm.weight" + ), + "module.module.decoder.layers.1.pre_mlp_layernorm.weight": ("body.layers.1.post_attention_layernorm.weight"), + } + for source_name, expected_name in mappings.items(): + assert convert_higgs_to_hf(args, source_name, norm) == [(expected_name, norm)] + + +def test_higgs_converter_covers_exact_399_weight_surface(convert_higgs_to_hf): + args = Namespace( + hidden_size=higgs_checkpoint.HIGGS_HIDDEN_SIZE, + kv_channels=higgs_checkpoint.HIGGS_HEAD_DIM, + num_attention_heads=higgs_checkpoint.HIGGS_NUM_ATTENTION_HEADS, + num_query_groups=higgs_checkpoint.HIGGS_NUM_QUERY_GROUPS, + higgs_num_codebooks=higgs_checkpoint.HIGGS_NUM_CODEBOOKS, + higgs_codebook_vocab_size=higgs_checkpoint.HIGGS_CODEBOOK_VOCAB_SIZE, + ) + expected_names = set(higgs_checkpoint.canonical_higgs_policy_shapes()) + for layout in ("transformer_engine", "local"): + exported_names = [] + for name, shape in higgs_checkpoint._target_parameter_shapes(layout).items(): + parameter = torch.empty(shape, dtype=torch.bfloat16, device="meta") + exported_names.extend( + output_name for output_name, _ in convert_higgs_to_hf(args, f"module.module.{name}", parameter) + ) + assert len(exported_names) == 399 + assert len(set(exported_names)) == 399 + assert set(exported_names) == expected_names + + +def test_higgs_converter_rejects_noncanonical_head_and_codec_shape(convert_higgs_to_hf): + args = _converter_args() + with pytest.raises(ValueError, match="tied Higgs text head"): + convert_higgs_to_hf(args, "module.module.output_layer.weight", torch.empty(7, 4)) + with pytest.raises(ValueError, match="codec embedding shape mismatch"): + convert_higgs_to_hf(args, "module.module.codec_embeddings.weight", torch.empty(5, 4)) + + +def test_converter_dispatch_uses_structured_policy_family_before_model_name(): + package_name = "_miles_higgs_dispatch_test" + base = Path(__file__).resolve().parents[4] / "miles" / "backends" / "megatron_utils" / "megatron_to_hf" + package = types.ModuleType(package_name) + package.__path__ = [str(base)] + sys.modules[package_name] = package + + exports = { + "deepseekv3": ("convert_deepseekv3_to_hf",), + "deepseekv4": ("convert_deepseekv4_to_hf",), + "glm4": ("convert_glm4_to_hf",), + "glm4moe": ("convert_glm4moe_to_hf",), + "kimi_vl": ("convert_kimi_k25_to_hf", "convert_kimivl_to_hf"), + "llama": ("convert_llama_to_hf",), + "mimo": ("convert_mimo_to_hf",), + "qwen2": ("convert_qwen2_to_hf",), + "qwen3_5": ("convert_qwen3_5_to_hf",), + "qwen3_next": ("convert_qwen3_next_to_hf",), + "qwen3moe": ("convert_qwen3moe_to_hf",), + } + for module_name, function_names in exports.items(): + module = types.ModuleType(f"{package_name}.{module_name}") + for function_name in function_names: + setattr( + module, function_name, lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("wrong dispatch")) + ) + sys.modules[module.__name__] = module + + calls = [] + higgs_module = types.ModuleType(f"{package_name}.higgs_tts") + higgs_module.convert_higgs_to_hf = lambda args, name, param: calls.append((name, param)) or [("higgs", param)] + sys.modules[higgs_module.__name__] = higgs_module + processors = types.ModuleType(f"{package_name}.processors") + processors.remove_padding = lambda name, param, vocab_size: param + processors.quantize_params = lambda args, name, tensors, config: tensors + sys.modules[processors.__name__] = processors + + spec = importlib.util.spec_from_file_location( + package_name, + base / "__init__.py", + submodule_search_locations=[str(base)], + ) + dispatch_module = importlib.util.module_from_spec(spec) + sys.modules[package_name] = dispatch_module + spec.loader.exec_module(dispatch_module) + + parameter = torch.ones(1) + assert dispatch_module._convert_to_hf_core( + Namespace(structured_policy_model_family="higgs_tts"), + "higgs-audio-v3-tts-4b", + "codec_embeddings.weight", + parameter, + ) == [("higgs", parameter)] + assert calls == [("codec_embeddings.weight", parameter)] + + assert dispatch_module._convert_to_hf_core( + Namespace(), + "MilesHiggsMultimodalQwen3Config", + "codec_embeddings.weight", + parameter, + ) == [("higgs", parameter)] + assert calls == [ + ("codec_embeddings.weight", parameter), + ("codec_embeddings.weight", parameter), + ] + + +def _write_manifest(tmp_path: Path, tensors: dict[str, torch.Tensor], weight_map: dict[str, str]) -> None: + save_file(tensors, tmp_path / "model.safetensors") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {}, "weight_map": weight_map}), + encoding="utf-8", + ) + + +def test_manifest_validation_is_exact_and_checks_shapes(tmp_path, monkeypatch): + monkeypatch.setattr( + higgs_checkpoint, + "_CANONICAL_HIGGS_POLICY_SHAPES", + {"policy.weight": (2, 3)}, + ) + weight_map = {"policy.weight": "model.safetensors"} + _write_manifest(tmp_path, {"policy.weight": torch.ones(2, 3, dtype=torch.bfloat16)}, weight_map) + assert higgs_checkpoint.validate_higgs_checkpoint_manifest(tmp_path) == weight_map + + _write_manifest(tmp_path, {"policy.weight": torch.ones(2, 4, dtype=torch.bfloat16)}, weight_map) + with pytest.raises(ValueError, match="checkpoint shape mismatch"): + higgs_checkpoint.validate_higgs_checkpoint_manifest(tmp_path) + + unexpected_map = {**weight_map, "policy.extra": "model.safetensors"} + _write_manifest( + tmp_path, + { + "policy.weight": torch.ones(2, 3, dtype=torch.bfloat16), + "policy.extra": torch.ones(1, dtype=torch.bfloat16), + }, + unexpected_map, + ) + with pytest.raises(ValueError, match="unexpected=.*policy.extra"): + higgs_checkpoint.validate_higgs_checkpoint_manifest(tmp_path) + + +def test_target_parameter_validation_accepts_only_complete_bf16_layouts(): + for layout in ("transformer_engine", "local"): + shapes = higgs_checkpoint._target_parameter_shapes(layout) + parameters = {name: torch.empty(shape, dtype=torch.bfloat16, device="meta") for name, shape in shapes.items()} + assert higgs_checkpoint.validate_higgs_target_parameters(parameters) == layout + + missing = dict(parameters) + missing.pop(next(iter(missing))) + with pytest.raises(ValueError, match="parameter manifest"): + higgs_checkpoint.validate_higgs_target_parameters(missing) + + wrong_dtype = dict(parameters) + name = next(iter(wrong_dtype)) + wrong_dtype[name] = torch.empty(shapes[name], dtype=torch.float32, device="meta") + with pytest.raises(ValueError, match="dtype mismatch"): + higgs_checkpoint.validate_higgs_target_parameters(wrong_dtype) + + +def test_direct_loader_copies_and_fuses_complete_policy(tmp_path, monkeypatch): + dimensions = { + "HIGGS_NUM_LAYERS": 1, + "HIGGS_HIDDEN_SIZE": 4, + "HIGGS_TEXT_VOCAB_SIZE": 7, + "HIGGS_NUM_ATTENTION_HEADS": 4, + "HIGGS_NUM_QUERY_GROUPS": 2, + "HIGGS_HEAD_DIM": 1, + "HIGGS_FFN_HIDDEN_SIZE": 5, + "HIGGS_NUM_CODEBOOKS": 2, + "HIGGS_CODEBOOK_VOCAB_SIZE": 3, + "HIGGS_CODEC_ROWS": 6, + } + for name, value in dimensions.items(): + monkeypatch.setattr(higgs_checkpoint, name, value) + policy_shapes = higgs_checkpoint._build_policy_shapes() + monkeypatch.setattr(higgs_checkpoint, "_CANONICAL_HIGGS_POLICY_SHAPES", policy_shapes) + + def values(shape, offset=0): + return torch.arange(offset, offset + torch.tensor(shape).prod().item(), dtype=torch.bfloat16).reshape(shape) + + source_tensors = {name: values(shape) for name, shape in policy_shapes.items()} + source_tensors["body.layers.0.self_attn.q_proj.weight"] = values((4, 4), 10) + source_tensors["body.layers.0.self_attn.k_proj.weight"] = values((2, 4), 100) + source_tensors["body.layers.0.self_attn.v_proj.weight"] = values((2, 4), 200) + source_tensors["body.layers.0.mlp.gate_proj.weight"] = values((5, 4), 300) + source_tensors["body.layers.0.mlp.up_proj.weight"] = values((5, 4), 400) + weight_map = {name: "model.safetensors" for name in policy_shapes} + _write_manifest(tmp_path, source_tensors, weight_map) + + class FakeHiggsModel: + num_codebooks = 2 + codebook_vocab_size = 3 + + def __init__(self): + self.parameters = { + name: torch.nn.Parameter(torch.full(shape, -1, dtype=torch.bfloat16)) + for name, shape in higgs_checkpoint._target_parameter_shapes("local").items() + } + + def named_parameters(self): + return self.parameters.items() + + model = FakeHiggsModel() + loaded = higgs_checkpoint.load_higgs_policy_checkpoint(model, tmp_path) + assert loaded == set(policy_shapes) + assert torch.equal( + model.parameters["embedding.word_embeddings.weight"], + source_tensors["tied.embedding.text_embedding.weight"], + ) + + q = source_tensors["body.layers.0.self_attn.q_proj.weight"].view(2, 2, 1, 4) + k = source_tensors["body.layers.0.self_attn.k_proj.weight"].view(2, 1, 1, 4) + v = source_tensors["body.layers.0.self_attn.v_proj.weight"].view(2, 1, 1, 4) + expected_qkv = torch.cat((q, k, v), dim=1).reshape(8, 4) + assert torch.equal( + model.parameters["decoder.layers.0.self_attention.linear_qkv.weight"], + expected_qkv, + ) + assert torch.equal( + model.parameters["decoder.layers.0.mlp.linear_fc1.weight"], + torch.cat( + ( + source_tensors["body.layers.0.mlp.gate_proj.weight"], + source_tensors["body.layers.0.mlp.up_proj.weight"], + ), + dim=0, + ), + ) + + +def test_repo_id_resolution_uses_snapshot_download(tmp_path, monkeypatch): + import huggingface_hub + + snapshot = tmp_path / "snapshot" + snapshot.mkdir() + calls = [] + monkeypatch.setattr( + huggingface_hub, + "snapshot_download", + lambda *, repo_id: calls.append(repo_id) or str(snapshot), + ) + + assert higgs_checkpoint.resolve_higgs_checkpoint_path("bosonai/higgs-audio-v3-tts-4b") == snapshot + assert calls == ["bosonai/higgs-audio-v3-tts-4b"] + + +def test_latest_cached_checkpoint_has_exact_policy_manifest(): + if not _CACHED_HIGGS_REVISION.is_dir(): + pytest.skip("latest Higgs checkpoint is not cached") + manifest = higgs_checkpoint.validate_higgs_checkpoint_manifest(_CACHED_HIGGS_REVISION) + assert len(manifest) == 399 + assert "tied.embedding.text_embedding.weight" in manifest + assert "tied.embedding.modality_embeddings.0.embedding.weight" in manifest + assert not any(name.startswith("tied.head.") for name in manifest) diff --git a/tests/fast/backends/megatron_utils/test_higgs_model.py b/tests/fast/backends/megatron_utils/test_higgs_model.py new file mode 100644 index 00000000000..6367cb1edb0 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_higgs_model.py @@ -0,0 +1,127 @@ +import sys +from types import ModuleType, SimpleNamespace + +import torch +import torch.nn.functional as F + +from miles.backends.megatron_utils.higgs_model import build_higgs_megatron_model + + +class _FakeVocabParallelEmbedding(torch.nn.Embedding): + def __init__(self, num_embeddings, embedding_dim, **_): + super().__init__(num_embeddings, embedding_dim) + + +class _FakeColumnParallelLinear(torch.nn.Module): + def __init__(self, input_size, output_size, *, skip_weight_param_allocation, **_): + super().__init__() + assert skip_weight_param_allocation + self.input_size = input_size + self.output_size = output_size + self.sequence_parallel = False + + def forward(self, input_, weight=None, runtime_gather_output=None): + assert weight is not None + assert runtime_gather_output is None + return F.linear(input_, weight), None + + +class _FakeGPTModel(torch.nn.Module): + def __init__(self, **kwargs): + super().__init__() + self.config = kwargs["config"] + self.pre_process = kwargs["pre_process"] + self.post_process = kwargs["post_process"] + self.pg_collection = SimpleNamespace(tp=object()) + self.embedding = SimpleNamespace( + word_embeddings=torch.nn.Embedding(kwargs["vocab_size"], self.config.hidden_size), + embedding_dropout=torch.nn.Identity(), + ) + self.output_layer = torch.nn.Linear(self.config.hidden_size, kwargs["vocab_size"], bias=False) + self.share_embeddings_and_output_weights = kwargs["share_embeddings_and_output_weights"] + self.vocab_size = kwargs["vocab_size"] + self.setup_embeddings_and_output_layer() + + def setup_embeddings_and_output_layer(self): + self.embedding.word_embeddings.weight.is_embedding_or_output_parameter = True + if self.share_embeddings_and_output_weights: + self.shared_embedding_or_output_weight().zero_out_wgrad = True + + def forward(self, *, decoder_input, **_): + logits, _ = self.output_layer(decoder_input, weight=self.shared_embedding_or_output_weight()) + return logits.transpose(0, 1).contiguous() + + +def _install_fake_megatron(monkeypatch): + megatron = ModuleType("megatron") + core = ModuleType("megatron.core") + models = ModuleType("megatron.core.models") + gpt = ModuleType("megatron.core.models.gpt") + core.tensor_parallel = SimpleNamespace( + VocabParallelEmbedding=_FakeVocabParallelEmbedding, + ColumnParallelLinear=_FakeColumnParallelLinear, + ) + gpt.GPTModel = _FakeGPTModel + monkeypatch.setitem(sys.modules, "megatron", megatron) + monkeypatch.setitem(sys.modules, "megatron.core", core) + monkeypatch.setitem(sys.modules, "megatron.core.models", models) + monkeypatch.setitem(sys.modules, "megatron.core.models.gpt", gpt) + + +def test_megatron_wrapper_teacher_forces_prior_rows_and_uses_tied_codec_head(monkeypatch): + _install_fake_megatron(monkeypatch) + config = SimpleNamespace( + hidden_size=2, + sequence_parallel=False, + embedding_init_method=lambda weight: None, + init_method=lambda weight: None, + fp32_residual_connection=False, + ) + model = build_higgs_megatron_model( + gpt_model_kwargs={ + "config": config, + "transformer_layer_spec": object(), + "vocab_size": 8, + "max_sequence_length": 16, + "pre_process": True, + "post_process": True, + "share_embeddings_and_output_weights": True, + }, + num_codebooks=2, + codebook_vocab_size=3, + ) + with torch.no_grad(): + model.embedding.word_embeddings.weight.zero_() + model.embedding.word_embeddings.weight[4] = torch.tensor([1.0, 0.0]) + model.embedding.word_embeddings.weight[5] = torch.tensor([0.0, 1.0]) + model.codec_embeddings.weight.copy_( + torch.tensor( + [ + [1.0, 0.0], + [2.0, 0.0], + [3.0, 0.0], + [0.0, 1.0], + [0.0, 2.0], + [0.0, 3.0], + ] + ) + ) + + logits = model( + input_ids=torch.tensor([[4, 5, 0]]), + higgs_prior_codes=torch.tensor([[[0, 0], [0, 0], [1, 2]]]), + higgs_codec_position_mask=torch.tensor([[False, False, True]]), + higgs_sequence_mask=torch.tensor([[True, True, True]]), + higgs_prediction_positions=torch.tensor([[1, 2]]), + ) + + expected_hidden = torch.tensor([[0.0, 1.0], [2.0, 3.0]]) + expected = F.linear(expected_hidden, model.codec_embeddings.weight).reshape(2, 2, 3) + assert logits.shape == (1, 2, 2, 3) + assert torch.allclose(logits[0], expected) + assert model.shared_embedding_or_output_weight() is model.codec_embeddings.weight + assert model.codec_embeddings.weight.is_embedding_or_output_parameter is True + assert model.codec_embeddings.weight.zero_out_wgrad is True + + logits.sum().backward() + assert model.codec_embeddings.weight.grad is not None diff --git a/tests/fast/backends/megatron_utils/test_lora_checkpoint_helpers.py b/tests/fast/backends/megatron_utils/test_lora_checkpoint_helpers.py index 18bfcb65da2..84642d9cfd4 100644 --- a/tests/fast/backends/megatron_utils/test_lora_checkpoint_helpers.py +++ b/tests/fast/backends/megatron_utils/test_lora_checkpoint_helpers.py @@ -9,8 +9,13 @@ from unittest.mock import MagicMock, patch import pytest +import torch -from miles.backends.megatron_utils.checkpoint import _is_megatron_checkpoint, save_checkpoint_with_lora +from miles.backends.megatron_utils.checkpoint import ( + _is_megatron_checkpoint, + _normalize_torch_optimizer_steps_for_checkpoint_load, + save_checkpoint_with_lora, +) # --------------------------------------------------------------------------- # _is_megatron_checkpoint @@ -63,6 +68,41 @@ def test_invalid_iter_patterns(self, tmp_path, name): assert _is_megatron_checkpoint(d) is False +def test_higgs_resume_normalizes_only_divergent_temporary_optimizer_steps(): + torch_optimizer = MagicMock() + torch_optimizer.state = { + "first": {"step": torch.tensor(0.0)}, + "second": {"step": torch.tensor(1.0)}, + "uninitialized": {}, + } + torch_optimizer.state_dict.return_value = {"state": torch_optimizer.state} + optimizer = MagicMock() + optimizer.chained_optimizers = [Namespace(optimizer=torch_optimizer)] + + _normalize_torch_optimizer_steps_for_checkpoint_load(optimizer) + + assert torch_optimizer.state["first"]["step"].item() == 0 + assert torch_optimizer.state["second"]["step"].item() == 0 + assert torch_optimizer.state["uninitialized"] == {} + + +def test_higgs_resume_initializes_an_empty_temporary_optimizer_state(): + torch_optimizer = MagicMock() + torch_optimizer.state_dict.return_value = {"state": {}} + wrapped_optimizer = Namespace(optimizer=torch_optimizer) + + def initialize_states(): + torch_optimizer.state_dict.return_value = {"state": {"first": {"step": torch.tensor(1.0)}}} + + wrapped_optimizer._init_optimizer_states_with_dummy_values = initialize_states + optimizer = MagicMock() + optimizer.chained_optimizers = [wrapped_optimizer] + + _normalize_torch_optimizer_steps_for_checkpoint_load(optimizer) + + assert torch_optimizer.state_dict.call_count == 2 + + # --------------------------------------------------------------------------- # save_checkpoint_with_lora — branch routing # --------------------------------------------------------------------------- diff --git a/tests/fast/backends/megatron_utils/test_update_weight_common.py b/tests/fast/backends/megatron_utils/test_update_weight_common.py new file mode 100644 index 00000000000..54830174639 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_update_weight_common.py @@ -0,0 +1,28 @@ +from argparse import Namespace +from types import SimpleNamespace + +import torch + +from miles.backends.megatron_utils.update_weight import common + + +def test_all_gather_param_skips_collective_for_tp1(monkeypatch) -> None: + parameter = torch.nn.Parameter(torch.arange(6).reshape(2, 3).float()) + parameter.tensor_model_parallel = True + parameter.partition_dim = 0 + parameter.partition_stride = 1 + parallel_state = SimpleNamespace( + tp=SimpleNamespace(size=1, group=object()), + etp=SimpleNamespace(size=1, group=object()), + ) + + monkeypatch.setattr(common, "get_parallel_state", lambda: parallel_state) + + def unexpected_all_gather(*args, **kwargs): + raise AssertionError("TP1 weight export must not enter a collective") + + monkeypatch.setattr(common.dist, "all_gather", unexpected_all_gather) + + gathered = common.all_gather_param(Namespace(), "module.weight", parameter) + + assert gathered.data_ptr() == parameter.data.data_ptr() diff --git a/tests/fast/backends/test_sglang_engine.py b/tests/fast/backends/test_sglang_engine.py new file mode 100644 index 00000000000..4ec2bee7804 --- /dev/null +++ b/tests/fast/backends/test_sglang_engine.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +import pytest + +from miles.backends.sglang_utils.sglang_engine import ( + _extract_omni_distributed_weight_update_transport, + _validate_omni_server_info, +) + + +def _model_info(**updates): + info = { + "success": True, + "model_path": ("/root/.cache/huggingface/hub/models--bosonai--higgs-audio-v3-tts-4b/" "snapshots/revision"), + "weight_version": "default", + "stages": [{"stage": "tts_engine", "data": {"tp_size": 1}}], + } + info.update(updates) + return info + + +def test_omni_external_server_accepts_matching_hf_snapshot() -> None: + _validate_omni_server_info( + _model_info(), + {"model_path": "bosonai/higgs-audio-v3-tts-4b", "tp_size": 1}, + ) + + +@pytest.mark.parametrize( + ("updates", "message"), + [ + ({"model_path": "/models/other"}, "model mismatch"), + ({"stages": [{"data": {"tp_size": 2}}]}, "TP mismatch"), + ({"weight_version": None}, "no weight_version"), + ({"success": False}, "model_info failed"), + ], +) +def test_omni_external_server_rejects_incompatible_identity(updates, message) -> None: + with pytest.raises(RuntimeError, match=message): + _validate_omni_server_info( + _model_info(**updates), + {"model_path": "bosonai/higgs-audio-v3-tts-4b", "tp_size": 1}, + ) + + +def test_extract_omni_distributed_weight_update_transport() -> None: + transport = { + "protocol_version": 1, + "backend": "nccl", + "nccl_version": "2.28.9", + "nccl_cumem_enable": "default", + } + model_info = _model_info( + stages=[ + {"stage": "preprocessing", "data": {}}, + { + "stage": "tts_engine", + "data": { + "supports_distributed_weight_update": True, + "distributed_weight_update": transport, + }, + }, + ] + ) + + assert _extract_omni_distributed_weight_update_transport(model_info) == transport + + +def test_extract_omni_distributed_weight_update_transport_requires_descriptor() -> None: + with pytest.raises(RuntimeError, match="no distributed weight-update transport"): + _extract_omni_distributed_weight_update_transport(_model_info()) diff --git a/tests/fast/backends/training_utils/test_higgs_policy.py b/tests/fast/backends/training_utils/test_higgs_policy.py new file mode 100644 index 00000000000..2bed3c0116b --- /dev/null +++ b/tests/fast/backends/training_utils/test_higgs_policy.py @@ -0,0 +1,574 @@ +import json +from types import SimpleNamespace + +import pytest +import torch + +from miles.backends.training_utils.higgs_policy import ( + build_higgs_teacher_embeddings, + collate_higgs_policy_batch, + get_higgs_joint_log_probs, + higgs_joint_policy_loss, + higgs_policy_loss_from_logits, + selected_higgs_logprobs, + validate_higgs_hf_config, + validate_higgs_logprob_parity, + validate_higgs_single_device_config, + validate_higgs_weight_versions, +) + + +class _Trace: + def __init__(self, actions, logprobs, mask, *, vocab_size=4): + self.model_family = "higgs_tts" + self.action_streams = [ + SimpleNamespace( + name="higgs_codes", + action_type="multi_discrete", + layout="time_codebook", + shape=[len(actions), len(actions[0])], + vocab_size=vocab_size, + actions=actions, + policy_logprobs=logprobs, + action_mask=mask, + ) + ] + + def validate(self): + return None + + +def _config(**overrides): + values = { + "structured_policy_model_family": "higgs_tts", + "hf_checkpoint": "bosonai/higgs-audio-v3-tts-4b", + "train_backend": "megatron", + "qkv_format": "bshd", + "tensor_model_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + "expert_model_parallel_size": 1, + "expert_tensor_parallel_size": None, + "actor_num_nodes": 1, + "actor_num_gpus_per_node": 1, + "advantage_estimator": "grpo", + "bf16": True, + "fp16": False, + "true_on_policy_mode": False, + "hidden_dropout": 0.0, + "attention_dropout": 0.0, + "masked_softmax_fusion": False, + "vocab_size": 151936, + "padded_vocab_size": 151936, + "group_query_attention": True, + "num_query_groups": 8, + "kv_channels": 128, + "qk_layernorm": True, + "swiglu": True, + "add_bias_linear": False, + "untie_embeddings_and_output_weights": False, + "normalization": "RMSNorm", + "use_rotary_position_embeddings": True, + "rotary_percent": 1.0, + "rotary_interleaved": False, + "use_rope_scaling": False, + "use_rollout_logprobs": True, + "megatron_to_hf_mode": "raw", + "lora_rank": 0, + "kl_coef": 0.0, + "compute_advantages_and_returns": True, + "debug_train_only": False, + "higgs_logprob_parity_atol": 0.05, + "rewards_normalization": True, + "n_samples_per_prompt": 2, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _hf_config(**audio_overrides): + audio = { + "encoder_type": "discrete", + "num_codebooks": 8, + "vocab_size": 1026, + "out_dim": 2560, + "tie_word_embeddings": True, + "use_delay_pattern": True, + } + audio.update(audio_overrides) + return { + "model_type": "higgs_multimodal_qwen3", + "architectures": ["HiggsMultimodalQwen3ForConditionalGeneration"], + "audio_encoder_config": audio, + "text_config": { + "model_type": "qwen3", + "hidden_size": 2560, + "num_hidden_layers": 36, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "head_dim": 128, + "intermediate_size": 9728, + "rms_norm_eps": 1e-6, + "vocab_size": 151936, + "max_position_embeddings": 32768, + "hidden_act": "silu", + "tie_word_embeddings": True, + "attention_dropout": 0.0, + "rope_parameters": {"rope_theta": 1_000_000, "rope_type": "default"}, + "dtype": "bfloat16", + }, + } + + +def test_collation_teacher_forces_complete_prior_rows_including_forced_cells(): + trace = _Trace( + actions=[[1, 3], [2, 0], [3, 1]], + logprobs=[[-0.2, 0.0], [-0.3, -0.4], [-0.5, 0.0]], + mask=[[True, False], [True, True], [True, False]], + ) + + batch = collate_higgs_policy_batch([[10, 11, 12]], [trace], advantages=[0.75]) + + assert batch.input_ids.tolist() == [[10, 11, 12, 0, 0]] + assert batch.codec_position_mask.tolist() == [[False, False, False, True, True]] + assert batch.prior_codes[0, 3:].tolist() == [[1, 3], [2, 0]] + assert batch.prediction_positions.tolist() == [[2, 3, 4]] + assert batch.actions.tolist() == [[[1, 3], [2, 0], [3, 1]]] + assert torch.allclose( + batch.old_cell_logprobs, + torch.tensor([[[-0.2, 0.0], [-0.3, -0.4], [-0.5, 0.0]]]), + ) + assert torch.allclose(batch.old_joint_logprobs, torch.tensor([[-0.2, -0.7, -0.5]])) + assert torch.allclose(batch.advantages, torch.full((1, 3), 0.75)) + + +def test_collation_right_pads_variable_prompt_and_action_lengths(): + first = _Trace([[0, 1]], [[-0.1, -0.2]], [[True, True]]) + second = _Trace( + [[1, 2], [2, 3], [0, 1]], + [[-0.1, -0.2], [-0.3, -0.4], [-0.5, -0.6]], + [[True, True], [True, True], [True, True]], + ) + + batch = collate_higgs_policy_batch([[7, 8, 9], [4]], [first, second]) + + assert batch.input_ids.shape == (2, 3) + assert batch.sequence_mask.tolist() == [[True, True, True], [True, True, True]] + assert batch.prediction_positions.tolist() == [[2, 0, 0], [0, 1, 2]] + assert batch.row_mask.tolist() == [[True, False, False], [True, True, True]] + + +def test_collation_uses_recomputed_joint_logprobs_as_old_policy_baseline(): + trace = _Trace( + [[0, 1], [1, 2]], + [[-1.0, -2.0], [-3.0, -4.0]], + [[True, True], [True, True]], + vocab_size=3, + ) + + batch = collate_higgs_policy_batch( + [[7, 8]], + [trace], + advantages=[1.0], + old_policy_joint_logprobs=[torch.tensor([-2.75, -6.5])], + ) + + assert torch.allclose(batch.old_cell_logprobs.sum(-1), torch.tensor([[-3.0, -7.0]])) + assert torch.allclose(batch.old_joint_logprobs, torch.tensor([[-2.75, -6.5]])) + + +@pytest.mark.parametrize( + "baseline", + [ + [torch.tensor([-1.0])], + [torch.tensor([-1.0, float("nan")])], + ], +) +def test_collation_rejects_invalid_recomputed_old_policy_baseline(baseline): + trace = _Trace( + [[0, 1], [1, 2]], + [[-1.0, -2.0], [-3.0, -4.0]], + [[True, True], [True, True]], + vocab_size=3, + ) + + with pytest.raises(ValueError, match="old-policy joint logprob"): + collate_higgs_policy_batch( + [[7, 8]], + [trace], + old_policy_joint_logprobs=baseline, + ) + + +def test_teacher_embedding_uses_channel_offsets_and_sums_codebooks(): + text_embeddings = torch.tensor([[[100.0], [200.0]]]) + # Two codebooks, vocabulary three. Row [1, 2] maps to weights 1 and 5. + codec_weight = torch.arange(6, dtype=torch.float32).unsqueeze(-1) + prior_codes = torch.tensor([[[0, 0], [1, 2]]]) + codec_mask = torch.tensor([[False, True]]) + + actual = build_higgs_teacher_embeddings(text_embeddings, codec_weight, prior_codes, codec_mask) + + assert actual.tolist() == [[[100.0], [6.0]]] + + +def test_selected_logprobs_use_fp32_full_vocabulary_and_zero_forced_cells(): + logits = torch.tensor( + [[[[2.0, 1.0, -1.0], [0.0, 3.0, 1.0]], [[-1.0, 0.0, 2.0], [1.0, 2.0, 3.0]]]], + dtype=torch.bfloat16, + ) + actions = torch.tensor([[[0, 1], [2, 0]]]) + mask = torch.tensor([[[True, True], [True, False]]]) + + cell, joint, row_mask = selected_higgs_logprobs(logits, actions, mask) + + expected = torch.log_softmax(logits.float(), dim=-1).gather(-1, actions.unsqueeze(-1)).squeeze(-1) + expected = torch.where(mask, expected, torch.zeros_like(expected)) + assert cell.dtype == torch.float32 + assert torch.allclose(cell, expected) + assert torch.allclose(joint, expected.sum(-1)) + assert row_mask.tolist() == [[True, True]] + assert cell[0, 1, 1].item() == 0.0 + + +def test_grpo_clips_the_joint_row_ratio_not_individual_codebook_ratios(): + # Factor ratios 2.0 and 0.55 would be clipped independently, but their + # joint ratio is 1.1 and must remain unclipped. + current = torch.tensor([[torch.log(torch.tensor(1.1))]]) + old = torch.zeros_like(current) + + loss, metrics = higgs_joint_policy_loss( + current, + old, + advantages=torch.ones_like(current), + row_mask=torch.ones_like(current, dtype=torch.bool), + eps_clip=0.2, + ) + + assert torch.allclose(loss, torch.tensor(-1.1)) + assert metrics["pg_clipfrac"].item() == 0.0 + + +def test_higgs_loss_backpropagates_through_each_active_codebook_only(): + trace = _Trace( + actions=[[0, 1], [2, 0]], + logprobs=[[-1.0, -1.0], [-1.0, 0.0]], + mask=[[True, True], [True, False]], + vocab_size=3, + ) + batch = collate_higgs_policy_batch([[5, 6]], [trace], advantages=[1.0]) + logits = torch.randn(1, 2, 2, 3, requires_grad=True) + _, initial_joint, _ = selected_higgs_logprobs(logits.detach(), batch.actions, batch.action_mask) + batch.old_joint_logprobs.copy_(initial_joint) + + loss, _, _ = higgs_policy_loss_from_logits(logits, batch, eps_clip=0.2) + loss.backward() + + assert logits.grad is not None + assert bool((logits.grad[0, 0, 0] != 0).any()) + assert bool((logits.grad[0, 0, 1] != 0).any()) + assert bool((logits.grad[0, 1, 0] != 0).any()) + assert torch.count_nonzero(logits.grad[0, 1, 1]).item() == 0 + + +@pytest.mark.parametrize("advantage", [float("nan"), float("inf"), float("-inf")]) +def test_higgs_loss_rejects_nonfinite_active_advantages(advantage): + current = torch.zeros((1, 1)) + with pytest.raises(RuntimeError, match="advantages must be finite"): + higgs_joint_policy_loss( + current, + current, + advantages=torch.tensor([[advantage]]), + row_mask=torch.ones_like(current, dtype=torch.bool), + eps_clip=0.2, + ) + + +def test_forward_logprob_collection_returns_joint_rows_per_sample(): + trace = _Trace([[0, 1], [1, 0]], [[-1.0, -1.0], [-1.0, -1.0]], [[True, True], [True, True]], vocab_size=2) + batch = collate_higgs_policy_batch([[4]], [trace]) + logits = torch.tensor([[[[2.0, 0.0], [0.0, 2.0]], [[1.0, 0.0], [2.0, 0.0]]]]) + + result = get_higgs_joint_log_probs(logits, higgs_batch=batch) + + _, expected, _ = selected_higgs_logprobs(logits, batch.actions, batch.action_mask) + assert len(result["log_probs"]) == 1 + assert torch.allclose(result["log_probs"][0], expected[0]) + + +def test_training_batch_requires_and_uses_preupdate_megatron_logprobs(monkeypatch): + from miles.backends.training_utils import data as data_module + + trace = _Trace( + [[0, 1], [1, 2]], + [[-1.0, -2.0], [-3.0, -4.0]], + [[True, True], [True, True]], + vocab_size=3, + ) + rollout_data = { + "tokens": [torch.tensor([7, 8])], + "action_traces": [trace], + "advantages": [torch.ones(2)], + "log_probs": [torch.tensor([-2.75, -6.5])], + } + + class Iterator: + def get_next(self, keys): + return {key: rollout_data.get(key) for key in keys} + + monkeypatch.setattr( + data_module, + "get_parallel_state", + lambda: SimpleNamespace(intra_dp=SimpleNamespace(size=1)), + ) + args = _config( + higgs_num_codebooks=2, + higgs_codebook_vocab_size=3, + seq_length=16, + ) + + batch = data_module.get_higgs_batch(Iterator(), args, require_advantages=True) + + assert torch.allclose(batch.old_joint_logprobs, torch.tensor([[-2.75, -6.5]])) + + rollout_data["log_probs"] = None + with pytest.raises(ValueError, match="pre-update Megatron joint logprobs"): + data_module.get_higgs_batch(Iterator(), args, require_advantages=True) + + +def test_structured_grpo_advantage_is_broadcast_to_action_rows(monkeypatch): + from miles.backends.training_utils import loss as loss_module + + trace = _Trace( + [[0, 1], [1, 2], [2, 0]], + [[-1.0, -1.0], [-1.0, -1.0], [-1.0, -1.0]], + [[True, True], [True, True], [True, True]], + vocab_size=3, + ) + monkeypatch.setattr( + loss_module, + "get_parallel_state", + lambda: SimpleNamespace(intra_dp=SimpleNamespace(size=1)), + ) + rollout_data = { + "action_traces": [trace], + "rewards": [0.625], + "tokens": [torch.tensor([4, 5])], + } + + loss_module.compute_advantages_and_returns(_config(), rollout_data) + + assert torch.allclose(rollout_data["advantages"][0], torch.full((3,), 0.625)) + assert torch.equal(rollout_data["returns"][0], rollout_data["advantages"][0]) + + +@pytest.mark.parametrize("reward", [float("nan"), float("inf"), float("-inf")]) +def test_structured_grpo_rejects_nonfinite_rewards(monkeypatch, reward): + from miles.backends.training_utils import loss as loss_module + + trace = _Trace([[0, 1]], [[-1.0, -1.0]], [[True, True]]) + monkeypatch.setattr( + loss_module, + "get_parallel_state", + lambda: SimpleNamespace(intra_dp=SimpleNamespace(size=1)), + ) + with pytest.raises(ValueError, match="reward must be finite"): + loss_module.compute_advantages_and_returns( + _config(), + { + "action_traces": [trace], + "rewards": [reward], + "tokens": [torch.tensor([4, 5])], + }, + ) + + +def test_structured_rollout_logging_uses_action_rows_not_empty_text_masks(): + from miles.backends.training_utils.log_utils import _get_higgs_rollout_log_dict + + trace = _Trace( + [[0, 1], [1, 2], [2, 0]], + [[-1.0, -2.0], [-3.0, 0.0], [0.0, 0.0]], + [[True, True], [True, False], [False, False]], + vocab_size=3, + ) + metrics = _get_higgs_rollout_log_dict( + { + "action_traces": [trace], + "advantages": [torch.tensor([0.5, 1.5, 999.0])], + "returns": [torch.tensor([0.5, 1.5, 999.0])], + "rewards": [0.75], + "response_lengths": [0], + "total_lengths": [2], + "loss_masks": [torch.empty(0)], + } + ) + + assert metrics["action_rows"] == 3.0 + assert metrics["active_action_rows"] == 2.0 + assert metrics["sampled_action_cells"] == 3.0 + assert metrics["rollout_joint_log_probs"] == -3.0 + assert metrics["advantages"] == 1.0 + + +@pytest.mark.parametrize( + ("override", "message"), + [ + ({"tensor_model_parallel_size": 2}, "tensor_model_parallel_size"), + ({"actor_num_gpus_per_node": 2}, "actor_num_gpus_per_node"), + ({"qkv_format": "thd"}, "qkv_format"), + ({"bf16": False}, "bf16"), + ({"fp16": True}, "fp16"), + ({"true_on_policy_mode": True}, "true_on_policy_mode"), + ({"hidden_dropout": 0.1}, "hidden_dropout"), + ({"attention_dropout": 0.1}, "attention_dropout"), + ({"vocab_size": 151727}, "vocab_size"), + ({"padded_vocab_size": 151744}, "padded_vocab_size"), + ({"group_query_attention": False}, "group_query_attention"), + ({"num_query_groups": 32}, "num_query_groups"), + ({"kv_channels": 80}, "kv_channels"), + ({"qk_layernorm": False}, "qk_layernorm"), + ({"swiglu": False}, "swiglu"), + ({"add_bias_linear": True}, "add_bias_linear"), + ({"untie_embeddings_and_output_weights": True}, "untie_embeddings_and_output_weights"), + ({"normalization": "LayerNorm"}, "normalization"), + ({"use_rotary_position_embeddings": False}, "use_rotary_position_embeddings"), + ({"rotary_percent": 0.5}, "rotary_percent"), + ({"rotary_interleaved": True}, "rotary_interleaved"), + ({"use_rope_scaling": True}, "use_rope_scaling"), + ({"masked_softmax_fusion": True}, "masked_softmax_fusion"), + ({"num_experts": 8}, "num_experts"), + ({"mtp_num_layers": 1}, "mtp_num_layers"), + ({"spec": ["custom", "provider"]}, "layer specs"), + ({"sequence_parallel": True}, "sequence_parallel"), + ({"megatron_to_hf_mode": "bridge"}, "checkpoint mapping"), + ({"hf_checkpoint": None}, "concrete Higgs"), + ({"save_hf": "/tmp/higgs-hf"}, "standalone Higgs HF snapshot export"), + ({"lora_rank": 8}, "lora_rank"), + ({"higgs_logprob_parity_atol": None}, "measured tolerance"), + ({"higgs_logprob_parity_atol": "0.05"}, "measured tolerance"), + ({"rewards_normalization": False}, "rewards_normalization"), + ({"n_samples_per_prompt": 1}, "nonzero grouped GRPO signal"), + ], +) +def test_single_device_config_fails_closed(override, message): + with pytest.raises(ValueError, match=message): + validate_higgs_single_device_config(_config(**override), data_parallel_size=1) + + +def test_single_device_config_accepts_naive_megatron_profile(): + validate_higgs_single_device_config(_config(), data_parallel_size=1) + + +def test_higgs_checkpoint_config_matches_v3_tts_policy_architecture(): + validate_higgs_hf_config(_hf_config(), num_codebooks=8, codebook_vocab_size=1026) + + +def test_higgs_qwen3_config_normalizes_explicit_null_rope_theta(): + from miles_plugins.models.higgs_tts import build_higgs_text_config + + text_config = dict(_hf_config()["text_config"]) + text_config["rope_parameters"] = {"rope_theta": None, "rope_type": "default"} + + normalized = build_higgs_text_config(text_config) + + assert normalized.rope_parameters["rope_theta"] == 1_000_000 + + +def test_higgs_qwen3_config_normalizes_legacy_top_level_null_rope_theta(): + from miles_plugins.models.higgs_tts import build_higgs_text_config + + text_config = dict(_hf_config()["text_config"]) + text_config.pop("rope_parameters") + text_config["rope_theta"] = None + + normalized = build_higgs_text_config(text_config) + + assert normalized.rope_theta == 1_000_000 + assert normalized.rope_parameters["rope_theta"] == 1_000_000 + + +def test_miles_hf_loader_registers_higgs_composition_config(tmp_path): + from miles.utils.hf_config import load_hf_config + + config = _hf_config() + config["text_config"]["rope_parameters"]["rope_theta"] = None + (tmp_path / "config.json").write_text(json.dumps(config)) + + loaded = load_hf_config(str(tmp_path)) + + assert loaded.model_type == "higgs_multimodal_qwen3" + assert loaded.text_config.hidden_size == 2560 + assert loaded.text_config.rope_parameters["rope_theta"] == 1_000_000 + validate_higgs_hf_config(loaded, num_codebooks=8, codebook_vocab_size=1026) + + +@pytest.mark.parametrize( + "audio_override", + [ + {"num_codebooks": 12}, + {"vocab_size": 1024}, + {"tie_word_embeddings": False}, + {"encoder_type": "whisper"}, + ], +) +def test_higgs_checkpoint_config_rejects_a_different_audio_policy(audio_override): + with pytest.raises(ValueError, match="unsupported Higgs checkpoint"): + validate_higgs_hf_config(_hf_config(**audio_override), num_codebooks=8, codebook_vocab_size=1026) + + +def test_logprob_parity_compares_joint_active_rows_and_accepts_within_tolerance(): + trace = _Trace( + [[0, 1], [1, 2], [2, 0]], + [[-1.0, -2.0], [-3.0, 0.0], [0.0, 0.0]], + [[True, True], [True, False], [False, False]], + vocab_size=3, + ) + + metrics = validate_higgs_logprob_parity( + [trace], + [torch.tensor([-2.99, -3.02, 123.0])], + atol=0.03, + ) + + assert metrics["max_abs_diff"] == pytest.approx(0.02) + assert metrics["mean_abs_diff"] == pytest.approx(0.015) + assert metrics["within_tolerance"] is True + + +def test_logprob_parity_reports_finite_joint_row_mismatch_without_blocking_training(): + trace = _Trace([[0, 1]], [[-1.0, -2.0]], [[True, True]], vocab_size=2) + + metrics = validate_higgs_logprob_parity([trace], [torch.tensor([-2.8])], atol=0.1) + + assert metrics == { + "max_abs_diff": pytest.approx(0.2), + "mean_abs_diff": pytest.approx(0.2), + "within_tolerance": False, + } + + +def test_weight_provenance_requires_one_current_equal_version_per_sample(): + assert validate_higgs_weight_versions([["7"], ["7"]], trainer_weight_version=7) == "7" + + +def test_weight_provenance_accepts_server_default_only_for_initial_trainer_version(): + assert validate_higgs_weight_versions([["default"], ["default"]], trainer_weight_version=0) == "default" + + with pytest.raises(RuntimeError, match="trainer expects '1'"): + validate_higgs_weight_versions([["default"]], trainer_weight_version=1) + + +@pytest.mark.parametrize( + ("versions", "message"), + [ + (None, "requires one rollout"), + ([[]], "exactly one"), + ([["1", "2"]], "exactly one"), + ([["1"], ["2"]], "mixes rollout"), + ([["1"]], "trainer expects '2'"), + ], +) +def test_weight_provenance_rejects_missing_mixed_or_stale_versions(versions, message): + with pytest.raises((ValueError, RuntimeError), match=message): + validate_higgs_weight_versions(versions, trainer_weight_version=2) diff --git a/tests/fast/plugins/__init__.py b/tests/fast/plugins/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/fast/plugins/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/fast/plugins/omni/__init__.py b/tests/fast/plugins/omni/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/fast/plugins/omni/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/fast/plugins/omni/conftest.py b/tests/fast/plugins/omni/conftest.py new file mode 100644 index 00000000000..2e8a3cd6037 --- /dev/null +++ b/tests/fast/plugins/omni/conftest.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import base64 +import io +import wave + +import pytest + + +def wav_base64(*, sample_rate: int = 24000, frames: int = 9600, amplitude: int = 4000) -> str: + samples = int(amplitude).to_bytes(2, "little", signed=True) * frames + buffer = io.BytesIO() + with wave.open(buffer, "wb") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(2) + wav_file.setframerate(sample_rate) + wav_file.writeframes(samples) + return base64.b64encode(buffer.getvalue()).decode("ascii") + + +@pytest.fixture +def higgs_response() -> dict: + actions = [ + [10, 1024, 1024, 1024, 1024, 1024, 1024, 1024], + [1025, 11, 12, 13, 14, 15, 16, 17], + ] + action_mask = [ + [True, False, False, False, False, False, False, False], + [True, True, True, True, True, True, True, True], + ] + policy_logprobs = [ + [-1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [-2.0, -3.0, -4.0, -5.0, -6.0, -7.0, -8.0, -9.0], + ] + return { + "text": "", + "audio": { + "data": wav_base64(), + "path": None, + "format": "wav", + "sample_rate": 24000, + }, + "meta_info": { + "finish_reason": {"type": "stop", "length": None}, + "prompt_tokens": 3, + "completion_tokens": 2, + "cached_tokens": 1, + "weight_version": "7", + "request_metadata": None, + "output_token_logprobs": [[-1.0, 10], [-2.0, 1025]], + "output_codebook_tokens": [list(row) for row in actions], + "omni_rollout": { + "version": 2, + "model_family": "higgs_tts", + "total_action_count": 9, + "action_streams": [ + { + "name": "higgs_codes", + "stage": "tts_engine", + "modality": "audio", + "action_type": "multi_discrete", + "layout": "time_codebook", + "shape": [2, 8], + "vocab_size": 1026, + "actions": [list(row) for row in actions], + "policy_logprobs": policy_logprobs, + "action_mask": action_mask, + "codec_content_mask": action_mask, + "channel_ids": list(range(8)), + } + ], + }, + }, + } diff --git a/tests/fast/plugins/omni/test_omni_generate_fn.py b/tests/fast/plugins/omni/test_omni_generate_fn.py new file mode 100644 index 00000000000..1a1ed8d6ea2 --- /dev/null +++ b/tests/fast/plugins/omni/test_omni_generate_fn.py @@ -0,0 +1,254 @@ +from __future__ import annotations + +from copy import deepcopy +from types import SimpleNamespace + +import pytest + +from miles.rollout.base_types import GenerateFnInput +from miles.utils.types import Sample +from miles_plugins.omni import omni_generate_fn +from miles_plugins.omni.omni_generate_fn import ( + OmniGenerateFn, + build_higgs_generate_payload, + build_zero_shot_higgs_prompt_ids, + neutral_higgs_sampling_params, +) + + +class FakeHiggsTokenizer: + def get_added_vocab(self) -> dict[str, int]: + return {"<|tts|>": 100, "<|text|>": 101, "<|audio|>": 102} + + def encode(self, text: str, *, add_special_tokens: bool) -> list[int]: + assert text == "speak" + assert add_special_tokens is False + return [7, 8] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("temperature", 0.7), + ("top_p", 0.9), + ("top_k", 10), + ("min_p", 0.1), + ("repetition_penalty", 1.1), + ], +) +def test_neutral_sampling_rejects_behavior_changes(field: str, value: float) -> None: + with pytest.raises(ValueError, match="Higgs RL requires"): + neutral_higgs_sampling_params({field: value}) + + +def test_payload_requests_only_structured_audio() -> None: + payload = build_higgs_generate_payload( + [1, 2], + { + "temperature": 1, + "top_p": 1.0, + "top_k": -1, + "sampling_seed": 12, + "max_new_tokens": 8, + "skip_special_tokens": False, + }, + ) + + assert payload["input_ids"] == [1, 2] + assert payload["output_modalities"] == ["audio"] + assert payload["return_logprob"] is True + assert payload["return_omni_rollout"] is True + assert payload["stream"] is False + assert payload["sampling_params"] == { + "temperature": 1.0, + "top_p": 1.0, + "min_p": 0.0, + "repetition_penalty": 1.0, + "max_new_tokens": 8, + "seed": 12, + } + assert "metadata" not in payload + + +def test_zero_shot_prompt_matches_higgs_server_contract() -> None: + assert build_zero_shot_higgs_prompt_ids(FakeHiggsTokenizer(), "speak") == [100, 101, 7, 8, 102] + + +def test_zero_shot_prompt_requires_higgs_special_tokens() -> None: + tokenizer = FakeHiggsTokenizer() + tokenizer.get_added_vocab = lambda: {"<|tts|>": 100, "<|text|>": 101} + + with pytest.raises(ValueError, match="missing Higgs TTS specials"): + build_zero_shot_higgs_prompt_ids(tokenizer, "speak") + + +@pytest.mark.asyncio +async def test_generate_keeps_prompt_and_audio_actions_separate(monkeypatch, higgs_response: dict) -> None: + seen: dict = {} + response = deepcopy(higgs_response) + response["meta_info"]["prompt_tokens"] = 5 + + async def fake_post(url: str, payload: dict): + seen.update(url=url, payload=payload) + return response + + monkeypatch.setattr(omni_generate_fn, "post", fake_post) + args = SimpleNamespace( + sglang_router_ip="127.0.0.1", + sglang_router_port=30000, + rollout_max_response_len=64, + rollout_max_context_len=128, + ) + state = SimpleNamespace(args=args, processor=None, tokenizer=FakeHiggsTokenizer()) + sample = Sample(prompt="speak", tokens=[100, 101, 7, 8, 102]) + + output = await OmniGenerateFn()( + GenerateFnInput( + state=state, + sample=sample, + sampling_params={"temperature": 1.0, "top_p": 1.0, "top_k": -1}, + evaluation=False, + ) + ) + + generated = output.samples + assert generated is sample + assert generated.tokens == [100, 101, 7, 8, 102] + assert generated.response == "" + assert generated.response_length == 0 + assert generated.rollout_log_probs is None + assert generated.loss_mask is None + assert generated.action_trace is not None + assert generated.decoded_audio is not None + assert generated.weight_versions == ["7"] + assert generated.status is Sample.Status.COMPLETED + assert seen["payload"]["input_ids"] == [100, 101, 7, 8, 102] + assert "metadata" not in seen["payload"] + + +@pytest.mark.asyncio +async def test_generate_targets_single_external_omni_engine(monkeypatch, higgs_response: dict) -> None: + seen: dict = {} + response = deepcopy(higgs_response) + response["meta_info"]["prompt_tokens"] = 5 + + async def fake_post(url: str, payload: dict): + seen["url"] = url + return response + + monkeypatch.setattr(omni_generate_fn, "post", fake_post) + args = SimpleNamespace( + rollout_external=True, + rollout_external_engine_addrs=["127.0.0.1:8300"], + rollout_max_response_len=64, + rollout_max_context_len=128, + ) + state = SimpleNamespace(args=args, processor=None, tokenizer=FakeHiggsTokenizer()) + + await OmniGenerateFn()( + GenerateFnInput(state=state, sample=Sample(prompt="speak"), sampling_params={}, evaluation=False) + ) + + assert seen["url"] == "http://127.0.0.1:8300/generate" + + +@pytest.mark.asyncio +async def test_generate_builds_zero_shot_prompt_for_normal_empty_token_sample( + monkeypatch, higgs_response: dict +) -> None: + seen: dict = {} + response = deepcopy(higgs_response) + response["meta_info"]["prompt_tokens"] = 5 + + async def fake_post(url: str, payload: dict): + seen.update(url=url, payload=payload) + return response + + monkeypatch.setattr(omni_generate_fn, "post", fake_post) + args = SimpleNamespace( + sglang_router_ip="127.0.0.1", + sglang_router_port=30000, + rollout_max_response_len=64, + rollout_max_context_len=128, + ) + state = SimpleNamespace(args=args, processor=None, tokenizer=FakeHiggsTokenizer()) + sample = Sample(prompt="speak") + + output = await OmniGenerateFn()(GenerateFnInput(state=state, sample=sample, sampling_params={}, evaluation=False)) + + assert output.samples.tokens == [100, 101, 7, 8, 102] + assert seen["payload"]["input_ids"] == [100, 101, 7, 8, 102] + + +@pytest.mark.asyncio +async def test_generate_rejects_reference_media_before_http(monkeypatch) -> None: + async def unexpected_post(url: str, payload: dict): + pytest.fail("unsupported media must be rejected before HTTP") + + monkeypatch.setattr(omni_generate_fn, "post", unexpected_post) + args = SimpleNamespace(sglang_router_ip="localhost", sglang_router_port=1) + state = SimpleNamespace(args=args, processor=None, tokenizer=FakeHiggsTokenizer()) + sample = Sample(prompt="speak", multimodal_inputs={"audios": ["https://example.test/reference.wav"]}) + + with pytest.raises(ValueError, match="zero-shot text-to-audio only"): + await OmniGenerateFn()(GenerateFnInput(state=state, sample=sample, sampling_params={}, evaluation=False)) + + +@pytest.mark.asyncio +async def test_generate_rejects_noncanonical_pretokenized_prompt() -> None: + args = SimpleNamespace(sglang_router_ip="localhost", sglang_router_port=1) + state = SimpleNamespace(args=args, processor=None, tokenizer=FakeHiggsTokenizer()) + sample = Sample(prompt="speak", tokens=[7, 8]) + + with pytest.raises(ValueError, match="canonical zero-shot encoding"): + await OmniGenerateFn()(GenerateFnInput(state=state, sample=sample, sampling_params={}, evaluation=False)) + + +@pytest.mark.parametrize( + "sample", + [ + Sample(response="partial", response_length=1, tokens=[1]), + Sample(rollout_log_probs=[]), + Sample(loss_mask=[]), + Sample(weight_versions=["6"]), + ], +) +@pytest.mark.asyncio +async def test_generate_rejects_partial_or_resumed_samples(sample: Sample) -> None: + args = SimpleNamespace(sglang_router_ip="localhost", sglang_router_port=1) + state = SimpleNamespace(args=args, processor=None, tokenizer=None) + with pytest.raises(ValueError, match="partial token|partial text|fresh sample|text loss"): + await OmniGenerateFn()(GenerateFnInput(state=state, sample=sample, sampling_params={}, evaluation=False)) + + +@pytest.mark.asyncio +async def test_generate_rejects_aborted_sample_with_partial_tokens() -> None: + args = SimpleNamespace(sglang_router_ip="localhost", sglang_router_port=1) + state = SimpleNamespace(args=args, processor=None, tokenizer=FakeHiggsTokenizer()) + sample = Sample(prompt="speak", status=Sample.Status.ABORTED, tokens=[100]) + + with pytest.raises(ValueError, match="partial token"): + await OmniGenerateFn()(GenerateFnInput(state=state, sample=sample, sampling_params={}, evaluation=False)) + + +@pytest.mark.asyncio +async def test_generate_allows_clean_aborted_retry(monkeypatch, higgs_response: dict) -> None: + response = deepcopy(higgs_response) + response["meta_info"]["prompt_tokens"] = 5 + + async def fake_post(url: str, payload: dict): + return response + + monkeypatch.setattr(omni_generate_fn, "post", fake_post) + args = SimpleNamespace( + sglang_router_ip="localhost", + sglang_router_port=1, + rollout_max_response_len=64, + rollout_max_context_len=128, + ) + state = SimpleNamespace(args=args, processor=None, tokenizer=FakeHiggsTokenizer()) + sample = Sample(prompt="speak", status=Sample.Status.ABORTED) + + output = await OmniGenerateFn()(GenerateFnInput(state=state, sample=sample, sampling_params={}, evaluation=False)) + + assert output.samples.status is Sample.Status.COMPLETED diff --git a/tests/fast/plugins/omni/test_rollout_contract.py b/tests/fast/plugins/omni/test_rollout_contract.py new file mode 100644 index 00000000000..44bd9de0c2f --- /dev/null +++ b/tests/fast/plugins/omni/test_rollout_contract.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import copy + +import pytest +from pydantic import ValidationError + +from miles.utils.types import DecodedAudio, RolloutActionTrace +from miles_plugins.omni.rollout_contract import parse_higgs_generate_response + + +def test_parse_higgs_v2_response_maps_to_domain_types(higgs_response: dict) -> None: + result = parse_higgs_generate_response(higgs_response, expected_prompt_tokens=3) + + assert isinstance(result.action_trace, RolloutActionTrace) + assert isinstance(result.decoded_audio, DecodedAudio) + assert result.action_trace.action_streams[0].shape == [2, 8] + assert result.action_trace.total_action_count == 9 + assert result.decoded_audio.sample_rate == 24000 + assert result.weight_version == "7" + + +@pytest.mark.parametrize( + ("mutate", "message"), + [ + ( + lambda body: body["meta_info"]["omni_rollout"]["action_streams"][0].update(shape=[2, 7]), + "8 codebooks", + ), + ( + lambda body: body["meta_info"]["omni_rollout"]["action_streams"][0]["action_mask"][0].__setitem__(1, 0), + "boolean", + ), + ( + lambda body: body["meta_info"]["omni_rollout"]["action_streams"][0]["policy_logprobs"][0].__setitem__( + 1, -1.0 + ), + "forced Higgs action policy logprob must be zero", + ), + ( + lambda body: body["meta_info"]["omni_rollout"]["action_streams"][0]["policy_logprobs"][1].__setitem__( + 0, float("nan") + ), + "non-finite", + ), + ( + lambda body: body["meta_info"]["omni_rollout"].update(total_action_count=8), + "total_action_count", + ), + ( + lambda body: body["meta_info"].update(completion_tokens=3), + "completion_tokens", + ), + ( + lambda body: body["meta_info"]["output_codebook_tokens"][0].__setitem__(0, 99), + "output_codebook_tokens", + ), + ( + lambda body: body["audio"].update(format="pcm"), + "wav", + ), + ( + lambda body: body["audio"].update(data=""), + "at least 1 character", + ), + ], +) +def test_parse_higgs_v2_response_fails_closed(higgs_response: dict, mutate, message: str) -> None: + body = copy.deepcopy(higgs_response) + mutate(body) + with pytest.raises((ValidationError, ValueError), match=message): + parse_higgs_generate_response(body) + + +def test_parse_rejects_prompt_retokenization(higgs_response: dict) -> None: + with pytest.raises(ValueError, match="exact prompt IDs"): + parse_higgs_generate_response(higgs_response, expected_prompt_tokens=4) + + +def test_parse_rejects_unknown_wire_fields(higgs_response: dict) -> None: + higgs_response["meta_info"]["unexpected"] = "silently dropped" + with pytest.raises(ValidationError, match="Extra inputs are not permitted"): + parse_higgs_generate_response(higgs_response) diff --git a/tests/fast/plugins/omni/test_tts_reward.py b/tests/fast/plugins/omni/test_tts_reward.py new file mode 100644 index 00000000000..a684ffad3af --- /dev/null +++ b/tests/fast/plugins/omni/test_tts_reward.py @@ -0,0 +1,258 @@ +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from aiohttp import web + +from miles.utils.types import DecodedAudio, Sample +from miles_plugins.omni import tts_reward +from miles_plugins.omni.tts_reward import ( + INVALID_AUDIO_REWARD, + SglangOmniASRReward, + TtsRoundTripReward, + character_error_rate, + compute_tts_reward, +) + +from .conftest import wav_base64 + + +def test_character_error_rate_is_normalized_and_bounded() -> None: + assert character_error_rate("Hello, world!", "hello world") == 0.0 + assert character_error_rate("abc", "xyzxyz") == 1.0 + + +def test_tts_reward_scores_valid_wav_without_retaining_components(monkeypatch) -> None: + reward = TtsRoundTripReward() + monkeypatch.setattr(reward, "transcribe", lambda waveform, sample_rate: "hello world") + audio = DecodedAudio(data=wav_base64(), format="wav", sample_rate=24000) + + assert reward.score(audio, "Hello, world!") == 1.0 + + +def test_tts_reward_defaults_to_cpu_when_actor_has_no_gpu(monkeypatch) -> None: + monkeypatch.delenv("MILES_TTS_ASR_DEVICE", raising=False) + + assert TtsRoundTripReward().device == "cpu" + + +def test_tts_reward_surfaces_asr_infrastructure_failures(monkeypatch) -> None: + reward = TtsRoundTripReward() + + def fail_transcription(waveform, sample_rate): + raise RuntimeError("ASR model unavailable") + + monkeypatch.setattr(reward, "transcribe", fail_transcription) + audio = DecodedAudio(data=wav_base64(), format="wav", sample_rate=24000) + + with pytest.raises(RuntimeError, match="ASR model unavailable"): + reward.score(audio, "hello") + + +def test_whisper_receives_attention_mask_and_transcribe_task() -> None: + calls: dict = {} + + class FakeProcessor: + def __call__(self, waveform, **kwargs): + calls["processor"] = kwargs + return SimpleNamespace( + input_features=torch.ones(1, 80, 4), + attention_mask=torch.ones(1, 4, dtype=torch.long), + ) + + def batch_decode(self, token_ids, *, skip_special_tokens): + assert skip_special_tokens is True + return ["hello"] + + class FakeModel: + dtype = torch.float32 + + def generate(self, features, **kwargs): + calls["generate"] = kwargs + return torch.tensor([[1]]) + + reward = TtsRoundTripReward(device="cpu") + reward._processor = FakeProcessor() + reward._model = FakeModel() + + assert reward.transcribe(np.ones(16000, dtype=np.float32), 16000) == "hello" + assert calls["processor"]["return_attention_mask"] is True + assert calls["generate"]["task"] == "transcribe" + assert torch.equal(calls["generate"]["attention_mask"], torch.ones(1, 4, dtype=torch.long)) + + +@pytest.mark.asyncio +async def test_reward_releases_audio_after_scoring(monkeypatch) -> None: + reward = TtsRoundTripReward() + monkeypatch.setattr(reward, "transcribe", lambda waveform, sample_rate: "hello") + monkeypatch.setattr(tts_reward, "_SHARED_REWARD", reward) + sample = Sample( + prompt="hello", + decoded_audio=DecodedAudio(data=wav_base64(), format="wav", sample_rate=24000), + ) + + assert await compute_tts_reward(SimpleNamespace(), sample) == 1.0 + assert sample.decoded_audio is None + + +@pytest.mark.asyncio +async def test_reward_compares_asr_to_the_text_sent_for_generation(monkeypatch) -> None: + reward = TtsRoundTripReward() + monkeypatch.setattr(reward, "transcribe", lambda waveform, sample_rate: "spoken prompt") + monkeypatch.setattr(tts_reward, "_SHARED_REWARD", reward) + sample = Sample( + prompt="spoken prompt", + label="unrelated dataset label", + decoded_audio=DecodedAudio(data=wav_base64(), format="wav", sample_rate=24000), + ) + + assert await compute_tts_reward(SimpleNamespace(), sample) == 1.0 + + +@pytest.mark.asyncio +async def test_reward_releases_audio_even_when_scorer_raises(monkeypatch) -> None: + class RaisingReward: + def score(self, audio, target_text): + raise RuntimeError("ASR failed") + + monkeypatch.setattr(tts_reward, "_SHARED_REWARD", RaisingReward()) + sample = Sample( + prompt="hello", + decoded_audio=DecodedAudio(data=wav_base64(), format="wav", sample_rate=24000), + ) + + with pytest.raises(RuntimeError, match="ASR failed"): + await compute_tts_reward(SimpleNamespace(), sample) + assert sample.decoded_audio is None + + +@pytest.mark.asyncio +async def test_batch_reward_releases_unscored_audio_after_failure(monkeypatch) -> None: + class RaisingReward: + def score(self, audio, target_text): + raise RuntimeError("ASR failed") + + monkeypatch.setattr(tts_reward, "_SHARED_REWARD", RaisingReward()) + samples = [ + Sample( + prompt="hello", + decoded_audio=DecodedAudio(data=wav_base64(), format="wav", sample_rate=24000), + ) + for _ in range(2) + ] + + with pytest.raises(RuntimeError, match="ASR failed"): + await compute_tts_reward(SimpleNamespace(), samples) + assert all(sample.decoded_audio is None for sample in samples) + + +@pytest.mark.asyncio +async def test_invalid_audio_is_rejected_and_released(monkeypatch) -> None: + monkeypatch.setattr(tts_reward, "_SHARED_REWARD", TtsRoundTripReward()) + sample = Sample( + prompt="hello", + decoded_audio=DecodedAudio(data="not-base64", format="wav", sample_rate=24000), + ) + + assert await compute_tts_reward(SimpleNamespace(), sample) == INVALID_AUDIO_REWARD + assert sample.decoded_audio is None + + +async def _start_asr_server(handler): + app = web.Application() + app.router.add_post("/v1/audio/transcriptions", handler) + runner = web.AppRunner(app) + await runner.setup() + site = web.TCPSite(runner, "127.0.0.1", 0) + await site.start() + port = site._server.sockets[0].getsockname()[1] + return runner, f"http://127.0.0.1:{port}" + + +@pytest.mark.asyncio +async def test_remote_asr_batches_concurrently_and_preserves_order(monkeypatch) -> None: + active = 0 + max_active = 0 + seen_models: list[str] = [] + + async def transcribe(request): + nonlocal active, max_active + form = await request.post() + seen_models.append(form["model"]) + index = int(form["file"].filename.removeprefix("rollout-").removesuffix(".wav")) + active += 1 + max_active = max(max_active, active) + await asyncio.sleep(0.04 if index == 0 else 0.01) + active -= 1 + return web.json_response({"text": "first" if index == 0 else "not second"}) + + runner, base_url = await _start_asr_server(transcribe) + reward = SglangOmniASRReward(base_url=base_url, concurrency=2) + monkeypatch.setattr(tts_reward, "_SHARED_REWARD", reward) + samples = [ + Sample( + prompt=prompt, + decoded_audio=DecodedAudio(data=wav_base64(), format="wav", sample_rate=24000), + ) + for prompt in ("first", "second") + ] + try: + rewards = await compute_tts_reward(SimpleNamespace(), samples) + finally: + await runner.cleanup() + + assert rewards == [1.0, 0.5] + assert max_active == 2 + assert seen_models == ["Qwen/Qwen3-ASR-1.7B"] * 2 + assert all(sample.decoded_audio is None for sample in samples) + + +@pytest.mark.asyncio +async def test_remote_asr_does_not_send_invalid_audio(monkeypatch) -> None: + request_count = 0 + + async def transcribe(request): + nonlocal request_count + request_count += 1 + return web.json_response({"text": "valid"}) + + runner, base_url = await _start_asr_server(transcribe) + monkeypatch.setattr(tts_reward, "_SHARED_REWARD", SglangOmniASRReward(base_url=base_url)) + samples = [ + Sample(prompt="invalid", decoded_audio=DecodedAudio(data="bad", format="wav", sample_rate=24000)), + Sample( + prompt="valid", + decoded_audio=DecodedAudio(data=wav_base64(), format="wav", sample_rate=24000), + ), + ] + try: + rewards = await compute_tts_reward(SimpleNamespace(), samples) + finally: + await runner.cleanup() + + assert rewards == [INVALID_AUDIO_REWARD, 1.0] + assert request_count == 1 + + +@pytest.mark.asyncio +async def test_remote_asr_surfaces_http_failure_and_releases_audio(monkeypatch) -> None: + async def transcribe(request): + return web.json_response({"detail": "model unavailable"}, status=503) + + runner, base_url = await _start_asr_server(transcribe) + monkeypatch.setattr(tts_reward, "_SHARED_REWARD", SglangOmniASRReward(base_url=base_url)) + sample = Sample( + prompt="hello", + decoded_audio=DecodedAudio(data=wav_base64(), format="wav", sample_rate=24000), + ) + try: + with pytest.raises(RuntimeError, match="HTTP 503.*model unavailable"): + await compute_tts_reward(SimpleNamespace(), sample) + finally: + await runner.cleanup() + + assert sample.decoded_audio is None diff --git a/tests/fast/ray/rollout/test_train_data_conversion.py b/tests/fast/ray/rollout/test_train_data_conversion.py index 4a3da1a2883..9837c81ceb6 100644 --- a/tests/fast/ray/rollout/test_train_data_conversion.py +++ b/tests/fast/ray/rollout/test_train_data_conversion.py @@ -11,7 +11,28 @@ convert_samples_to_train_data, split_train_data_by_dp, ) -from miles.utils.types import Sample +from miles.utils.types import DecodedAudio, DiscreteActionStream, RolloutActionTrace, Sample + + +def _make_action_trace(action: int) -> RolloutActionTrace: + return RolloutActionTrace( + version=2, + model_family="higgs_tts", + total_action_count=1, + action_streams=[ + DiscreteActionStream( + name="higgs_codes", + stage="tts_engine", + modality="audio", + shape=[1, 1], + vocab_size=8, + actions=[[action]], + policy_logprobs=[[-0.5]], + action_mask=[[True]], + channel_ids=[0], + ) + ], + ) @pytest.fixture(scope="module", autouse=True) @@ -173,6 +194,44 @@ def test_dynamic_global_batch_size_metadata_passed_through(self): ) assert out["dynamic_global_batch_size"] == 16 + def test_structured_action_traces_pass_through_without_decoded_audio(self): + args = make_args(rewards_normalization=False) + samples = [ + make_sample( + response_length=0, + tokens=[10, 11], + action_trace=_make_action_trace(action), + decoded_audio=DecodedAudio(data="UklGRg==", format="wav", sample_rate=24000), + ) + for action in (3, 4) + ] + + out = convert_samples_to_train_data( + args, + samples, + metadata={}, + custom_convert_samples_to_train_data_func=None, + custom_reward_post_process_func=None, + ) + + assert out["action_traces"] == [sample.action_trace for sample in samples] + assert "decoded_audio" not in out + assert "decoded_audios" not in out + + def test_mixed_structured_and_text_batch_is_rejected(self): + args = make_args(rewards_normalization=False) + structured = make_sample(action_trace=_make_action_trace(3)) + text = make_sample() + + with pytest.raises(ValueError, match="cannot mix samples"): + convert_samples_to_train_data( + args, + [structured, text], + metadata={}, + custom_convert_samples_to_train_data_func=None, + custom_reward_post_process_func=None, + ) + # ----------------------------- _post_process_rewards ----------------------------- @@ -216,6 +275,20 @@ def test_grpo_with_std_normalization_unit_variance(self): assert abs(np.std(processed, ddof=1) - 1.0) < 1e-4 + def test_grpo_identical_nonrepresentable_rewards_have_exactly_zero_advantage(self): + args = make_args( + advantage_estimator="grpo", + rewards_normalization=True, + grpo_std_normalization=True, + n_samples_per_prompt=8, + rollout_batch_size=1, + ) + samples = make_samples_grouped(1, 8, rewards=[0.9] * 8) + + _, processed = _post_process_rewards(args, samples, custom_reward_post_process_func=None) + + assert processed == [0.0] * 8 + def test_gspo_uses_grpo_normalization_path(self): args = make_args( advantage_estimator="gspo", @@ -437,6 +510,25 @@ def test_optional_keys_propagated_when_present(self): assert "rollout_log_probs" in parts[0] assert "round_number" in parts[0] + def test_action_traces_are_partitioned(self): + args = make_args(balance_data=False) + traces = [_make_action_trace(action) for action in (1, 2, 3, 4)] + data = { + "tokens": [[1], [2], [3], [4]], + "response_lengths": [0, 0, 0, 0], + "rewards": [0, 0, 0, 0], + "truncated": [0, 0, 0, 0], + "loss_masks": [[], [], [], []], + "sample_indices": [0, 1, 2, 3], + "action_traces": traces, + } + + refs = split_train_data_by_dp(args, data, dp_size=2) + parts = [ray.get(r.inner) for r in refs] + + assert parts[0]["action_traces"] == [traces[0], traces[2]] + assert parts[1]["action_traces"] == [traces[1], traces[3]] + def test_shared_keys_not_split(self): """raw_reward, total_lengths, dynamic_global_batch_size are shared, not split.""" args = make_args(balance_data=False) diff --git a/tests/fast/ray/test_actor_group.py b/tests/fast/ray/test_actor_group.py new file mode 100644 index 00000000000..52bcadbff0c --- /dev/null +++ b/tests/fast/ray/test_actor_group.py @@ -0,0 +1,25 @@ +from miles.ray.actor_group import _build_train_actor_env_vars + + +def test_train_actor_does_not_invent_nccl_cumem_setting(monkeypatch) -> None: + monkeypatch.delenv("NCCL_CUMEM_ENABLE", raising=False) + + env_vars = _build_train_actor_env_vars({}) + + assert "NCCL_CUMEM_ENABLE" not in env_vars + + +def test_train_actor_preserves_explicit_nccl_cumem_setting(monkeypatch) -> None: + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1") + + env_vars = _build_train_actor_env_vars({}) + + assert env_vars["NCCL_CUMEM_ENABLE"] == "1" + + +def test_train_env_vars_override_process_nccl_cumem_setting(monkeypatch) -> None: + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1") + + env_vars = _build_train_actor_env_vars({"NCCL_CUMEM_ENABLE": "0"}) + + assert env_vars["NCCL_CUMEM_ENABLE"] == "0" diff --git a/tests/fast/ray/test_placement_group.py b/tests/fast/ray/test_placement_group.py new file mode 100644 index 00000000000..f1bb6258cf5 --- /dev/null +++ b/tests/fast/ray/test_placement_group.py @@ -0,0 +1,30 @@ +from types import SimpleNamespace + +from miles.ray import placement_group + + +def test_external_rollout_reserves_only_training_gpus(monkeypatch) -> None: + requested = [] + fake_pg = object() + + def fake_create(num_gpus): + requested.append(num_gpus) + return fake_pg, [0], [6] + + monkeypatch.setattr(placement_group, "_create_placement_group", fake_create) + args = SimpleNamespace( + debug_train_only=False, + debug_rollout_only=False, + colocate=False, + rollout_external=True, + actor_num_nodes=1, + actor_num_gpus_per_node=1, + rollout_num_gpus=1, + use_critic=False, + ) + + groups = placement_group.create_placement_groups(args) + + assert requested == [1] + assert groups["actor"] == (fake_pg, [0], [6]) + assert groups["rollout"] == (fake_pg, [], []) diff --git a/tests/fast/rollout/generate_utils/test_sample_utils.py b/tests/fast/rollout/generate_utils/test_sample_utils.py index 1feabdb0079..284a9c1e816 100644 --- a/tests/fast/rollout/generate_utils/test_sample_utils.py +++ b/tests/fast/rollout/generate_utils/test_sample_utils.py @@ -4,7 +4,7 @@ import pytest from miles.rollout.generate_utils.sample_utils import _merge_sample_pair -from miles.utils.types import Sample +from miles.utils.types import DiscreteActionStream, RolloutActionTrace, Sample @pytest.fixture @@ -175,3 +175,27 @@ def test_sample_validate_fails_raises(self, mock_tokenizer): with pytest.raises(AssertionError, match="loss_mask length"): _merge_sample_pair(a, b, mock_tokenizer) + + def test_structured_action_trace_cannot_be_merged(self, mock_tokenizer): + a = make_sample(tokens=[1, 2, 10], response_length=1, loss_mask=[1]) + b = make_sample(tokens=[1, 2, 10, 20, 30], response_length=1, loss_mask=[1]) + stream = DiscreteActionStream( + name="higgs_codes", + stage="tts_engine", + modality="audio", + shape=[1, 1], + vocab_size=8, + actions=[[3]], + policy_logprobs=[[-0.2]], + action_mask=[[True]], + channel_ids=[0], + ) + a.action_trace = RolloutActionTrace( + version=2, + model_family="higgs_tts", + total_action_count=1, + action_streams=[stream], + ) + + with pytest.raises(ValueError, match="structured action traces cannot be merged"): + _merge_sample_pair(a, b, mock_tokenizer) diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py index 5fb43454800..4da94220991 100644 --- a/tests/fast/utils/test_arguments.py +++ b/tests/fast/utils/test_arguments.py @@ -5,6 +5,7 @@ import pytest +from miles.backends.sglang_utils.arguments import validate_args as validate_sglang_args from miles.utils.arguments import _maybe_apply_dumper_overrides, get_miles_extra_args_provider from miles.utils.misc import function_registry @@ -141,3 +142,24 @@ def test_recompute_logprobs_via_prefill_flag_is_parsed(): args = parser.parse_args(["--recompute-logprobs-via-prefill"] + REQUIRED_ARGS) assert args.recompute_logprobs_via_prefill is True + + +def test_sglang_parallelism_long_form_cli_fields_are_normalized() -> None: + args = SimpleNamespace( + rollout_num_gpus_per_engine=2, + sglang_data_parallel_size=3, + sglang_pipeline_parallel_size=4, + sglang_expert_parallel_size=5, + sglang_enable_dp_attention=True, + true_on_policy_mode=False, + recompute_logprobs_via_prefill=False, + sglang_router_policy=None, + sglang_router_ip=None, + ) + + validate_sglang_args(args) + + assert args.sglang_tp_size == 2 + assert args.sglang_dp_size == 3 + assert args.sglang_pp_size == 4 + assert args.sglang_ep_size == 5 diff --git a/tests/fast/utils/test_types.py b/tests/fast/utils/test_types.py index 1d45dc93935..4dbfb90ac06 100644 --- a/tests/fast/utils/test_types.py +++ b/tests/fast/utils/test_types.py @@ -5,7 +5,28 @@ import numpy import pytest -from miles.utils.types import Sample +from miles.utils.types import DecodedAudio, DiscreteActionStream, RolloutActionTrace, Sample + + +def _make_action_trace() -> RolloutActionTrace: + stream = DiscreteActionStream( + name="higgs_codes", + stage="tts_engine", + modality="audio", + shape=[2, 2], + vocab_size=8, + actions=[[0, 3], [4, 5]], + policy_logprobs=[[0.0, -0.3], [-0.4, -0.5]], + action_mask=[[False, True], [True, True]], + codec_content_mask=[[False, True], [True, True]], + channel_ids=[0, 1], + ) + return RolloutActionTrace( + version=2, + model_family="higgs_tts", + total_action_count=3, + action_streams=[stream], + ) def _make_sample( @@ -105,3 +126,94 @@ def test_strip_negative_is_noop(self, tokenizer): original_tokens = list(s.tokens) s.strip_last_output_tokens(-1, tokenizer) assert s.tokens == original_tokens + + +class TestStructuredActionTypes: + def test_sample_dict_round_trip_preserves_typed_artifacts(self): + sample = Sample( + tokens=[1, 2], + action_trace=_make_action_trace(), + decoded_audio=DecodedAudio(data="UklGRg==", format="wav", sample_rate=24000), + ) + + restored = Sample.from_dict(sample.to_dict()) + + assert isinstance(restored.action_trace, RolloutActionTrace) + assert isinstance(restored.action_trace.action_streams[0], DiscreteActionStream) + assert isinstance(restored.decoded_audio, DecodedAudio) + assert restored.action_trace == sample.action_trace + assert restored.decoded_audio == sample.decoded_audio + + def test_action_stream_rejects_shape_mismatch(self): + data = _make_action_trace().action_streams[0].to_dict() + data["shape"] = [3, 2] + + with pytest.raises(ValueError, match="declared shape"): + DiscreteActionStream.from_dict(data) + + def test_action_stream_rejects_nonfinite_sampled_logprob(self): + data = _make_action_trace().action_streams[0].to_dict() + data["policy_logprobs"][0][1] = float("nan") + + with pytest.raises(ValueError, match="non-finite"): + DiscreteActionStream.from_dict(data) + + def test_action_stream_rejects_nonzero_forced_logprob(self): + data = _make_action_trace().action_streams[0].to_dict() + data["policy_logprobs"][0][0] = -1.0 + + with pytest.raises(ValueError, match="must be zero"): + DiscreteActionStream.from_dict(data) + + def test_action_stream_rejects_out_of_range_action(self): + data = _make_action_trace().action_streams[0].to_dict() + data["actions"][1][1] = data["vocab_size"] + + with pytest.raises(ValueError, match="outside"): + DiscreteActionStream.from_dict(data) + + def test_action_stream_requires_strict_ordered_channel_ids(self): + data = _make_action_trace().action_streams[0].to_dict() + data["channel_ids"] = [False, 1] + + with pytest.raises(ValueError, match="channel_ids"): + DiscreteActionStream.from_dict(data) + + def test_to_dict_revalidates_mutated_streams(self): + stream = _make_action_trace().action_streams[0] + stream.policy_logprobs[0][0] = -1.0 + + with pytest.raises(ValueError, match="must be zero"): + stream.to_dict() + + def test_trace_rejects_incorrect_action_count(self): + data = _make_action_trace().to_dict() + data["total_action_count"] = 2 + + with pytest.raises(ValueError, match="total_action_count"): + RolloutActionTrace.from_dict(data) + + def test_trace_rejects_unknown_fields(self): + data = _make_action_trace().to_dict() + data["unexpected"] = True + + with pytest.raises(ValueError, match="extra=.*unexpected"): + RolloutActionTrace.from_dict(data) + + def test_decoded_audio_requires_wav_with_positive_sample_rate(self): + with pytest.raises(ValueError, match="format"): + DecodedAudio(data="data", format="mp3", sample_rate=24000) + with pytest.raises(ValueError, match="sample_rate"): + DecodedAudio(data="data", format="wav", sample_rate=0) + + def test_reset_for_retry_clears_structured_outputs(self): + sample = Sample( + tokens=[1, 2], + action_trace=_make_action_trace(), + decoded_audio=DecodedAudio(data="UklGRg==", format="wav", sample_rate=24000), + ) + + sample.reset_for_retry() + + assert sample.action_trace is None + assert sample.decoded_audio is None diff --git a/train.py b/train.py index af296622589..4e55e512002 100644 --- a/train.py +++ b/train.py @@ -1,4 +1,5 @@ import asyncio +import logging from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS @@ -9,6 +10,8 @@ from miles.utils.misc import should_run_periodic_action from miles.utils.tracking_utils import finish_tracking, init_tracking +logger = logging.getLogger(__name__) + async def train(args): configure_logger() @@ -23,6 +26,17 @@ async def train(args): # create the actor and critic models actor_model, critic_model = await create_training_models(args, pgs, rollout_manager) + try: + await _run_training(args, actor_model, critic_model, rollout_manager, num_rollout_per_epoch) + finally: + try: + await actor_model.disconnect_rollout_engines() + except Exception: + logger.exception("Failed to disconnect rollout weight-update groups during shutdown") + await rollout_manager.dispose.remote() + + +async def _run_training(args, actor_model, critic_model, rollout_manager, num_rollout_per_epoch): if args.offload_rollout: await rollout_manager.onload_weights.remote() @@ -106,8 +120,6 @@ async def save(rollout_id): if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch): await rollout_manager.eval.remote(rollout_id) - await rollout_manager.dispose.remote() - if __name__ == "__main__": args = parse_args()