From 0e2d566d4dd4e71e8e7938166c51f02328763f3c Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Wed, 22 Jul 2026 04:30:52 +0000 Subject: [PATCH 01/20] [Feat] Add distillation loss and compute logprobs on train/inference engine --- .../build_teacher_server_commands.py | 418 +++++++++++++ .../config/rl_dapo_math_mopd.py | 397 ++++++++++++ .../config/rl_dapo_math_mopd_async.py | 398 ++++++++++++ .../config/rl_dapo_math_opd.py | 247 ++++++++ .../scripts/launch_teacher_utils.sh | 267 ++++++++ .../scripts/run_sampled_token_mopd.sh | 74 +++ .../scripts/run_sampled_token_opd.sh | 72 +++ tests/rl/test_on_policy_distillation.py | 580 ++++++++++++++++++ tests/rl/test_prepare_train_data.py | 12 +- tests/rl/test_producer.py | 8 +- tests/rl/test_rl_colocate_trainer.py | 2 +- xtuner/v1/data_proto/rl_data.py | 4 + xtuner/v1/loss/__init__.py | 4 +- xtuner/v1/loss/rl_loss.py | 73 +++ xtuner/v1/model/utils/misc.py | 11 + xtuner/v1/rl/agent_loop/agent_loop.py | 84 +++ .../agent_loop_manager/agent_loop_manager.py | 5 + .../disagg_agent_loop_manager.py | 4 + .../rl/agent_loop_manager/disagg_producer.py | 2 +- .../v1/rl/agent_loop_manager/produce_utils.py | 12 +- xtuner/v1/rl/agent_loop_manager/producer.py | 6 +- xtuner/v1/rl/distillation/__init__.py | 30 + xtuner/v1/rl/distillation/config.py | 147 +++++ .../rl/distillation/rollout_teacher_client.py | 339 ++++++++++ .../rl/distillation/train_teacher_manager.py | 239 ++++++++ xtuner/v1/rl/loss/__init__.py | 7 + xtuner/v1/rl/loss/base_loss.py | 1 + xtuner/v1/rl/loss/distillation_loss.py | 447 ++++++++++++++ xtuner/v1/rl/model_utils.py | 54 ++ xtuner/v1/rl/trainer/controller.py | 52 ++ xtuner/v1/rl/trainer/worker.py | 142 +++-- xtuner/v1/train/rl_trainer.py | 178 ++++-- 32 files changed, 4208 insertions(+), 108 deletions(-) create mode 100644 recipe/on_policy_distillation/build_teacher_server_commands.py create mode 100644 recipe/on_policy_distillation/config/rl_dapo_math_mopd.py create mode 100644 recipe/on_policy_distillation/config/rl_dapo_math_mopd_async.py create mode 100644 recipe/on_policy_distillation/config/rl_dapo_math_opd.py create mode 100644 recipe/on_policy_distillation/scripts/launch_teacher_utils.sh create mode 100755 recipe/on_policy_distillation/scripts/run_sampled_token_mopd.sh create mode 100755 recipe/on_policy_distillation/scripts/run_sampled_token_opd.sh create mode 100644 tests/rl/test_on_policy_distillation.py create mode 100644 xtuner/v1/rl/distillation/__init__.py create mode 100644 xtuner/v1/rl/distillation/config.py create mode 100644 xtuner/v1/rl/distillation/rollout_teacher_client.py create mode 100644 xtuner/v1/rl/distillation/train_teacher_manager.py create mode 100644 xtuner/v1/rl/loss/distillation_loss.py create mode 100644 xtuner/v1/rl/model_utils.py diff --git a/recipe/on_policy_distillation/build_teacher_server_commands.py b/recipe/on_policy_distillation/build_teacher_server_commands.py new file mode 100644 index 0000000000..88b3185320 --- /dev/null +++ b/recipe/on_policy_distillation/build_teacher_server_commands.py @@ -0,0 +1,418 @@ +"""Build executable Teacher server commands from an OPD config. + +The NUL-delimited output starts with the Teacher replica count, resolved +endpoint mapping, total Student worker count, and current-node Student worker +count. Each replica record contains its placement, endpoint, health URLs, and +executable command arguments. An externally managed Teacher replica has a +target node rank of ``-1`` and a command-argument count of zero. +""" + +import argparse +import json +import os +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Literal +from urllib.parse import urlparse + + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + +from xtuner.v1.rl.on_policy_distillation import OPDTeacherConfig # noqa: E402 +from xtuner.v1.utils.config import Config # noqa: E402 + + +@dataclass(frozen=True) +class TeacherReplicaRequest: + teacher: OPDTeacherConfig + replica_index: int + + @property + def key(self) -> tuple[str, int]: + return self.teacher.name, self.replica_index + + @property + def display_name(self) -> str: + return f"{self.teacher.name}[{self.replica_index}]" + + +def build_teacher_launch_server_commands( + config_path: str, + backend: Literal["sglang", "lmdeploy"], +) -> tuple[dict[str, list[str]], int, int, list[list[str]]]: + """Build executable Teacher server command records. + + Args: + config_path (str): Path to the XTuner Python config containing + ``opd_config``. + backend (Literal["sglang", "lmdeploy"]): Teacher serving backend. + + Returns: + A four-item tuple containing: + + - ``endpoint_map``: Logical Teacher name to advertised replica + endpoints. + - ``student_num_workers``: Total number of Student GPUs in the cluster. + - ``student_local_num_workers``: Number of Student GPUs on the current + node. + - ``records``: One flattened string record per Teacher replica. Each + record is ``[display_name, target_node_rank, local_devices, endpoint, + health_url, model_info_url, command_arg_count, *command]``. + """ + config = Config.fromfile(config_path) + node_count = int(os.environ.get("NODE_COUNT", "1")) + node_rank = int(os.environ.get("NODE_RANK", "0")) + gpus_per_node = int(os.environ.get("PROC_PER_NODE", "8")) + node_addresses = tuple( + address.strip() + for address in os.environ.get( + "WORKER_ALL_SOCKET_ADDRS", "127.0.0.1" + ).split(",") + ) + assert node_count > 0 + assert 0 <= node_rank < node_count + assert gpus_per_node > 0 + assert len(node_addresses) == node_count + assert all(node_addresses) + + teachers = config.opd_config.teachers + replica_requests = _expand_teacher_replicas(teachers) + health_path, model_info_path = _get_teacher_server_paths(backend) + teacher_placements, student_local_num_workers = _allocate_teacher_devices( + replica_requests, + node_count=node_count, + node_rank=node_rank, + gpus_per_node=gpus_per_node, + ) + teacher_num_workers = sum( + len(local_devices) + for _, local_devices in teacher_placements.values() + ) + student_num_workers = node_count * gpus_per_node - teacher_num_workers + endpoint_map: dict[str, list[str]] = {teacher.name: [] for teacher in teachers} + records: list[list[str]] = [] + used_node_ports: set[tuple[int, int]] = set() + + for replica in replica_requests: + teacher = replica.teacher + launch_config = teacher.launch_config + if launch_config is None: + target_node_rank = -1 + local_cuda_visible_devices = "" + endpoint = teacher.endpoints[replica.replica_index].rstrip("/") + command: list[str] = [] + else: + target_node_rank, local_devices = teacher_placements[replica.key] + server_port = _allocate_teacher_server_port( + launch_config.server_port, + target_node_rank=target_node_rank, + used_node_ports=used_node_ports, + replica_name=replica.display_name, + ) + + local_cuda_visible_devices = ",".join(str(device) for device in local_devices) + endpoint = f"http://{node_addresses[target_node_rank]}:{server_port}" + if backend == "sglang": + command = _build_sglang_command(teacher, server_port=server_port) + elif backend == "lmdeploy": + command = _build_lmdeploy_command(teacher, server_port=server_port) + else: + raise ValueError(f"Unsupported Teacher backend: {backend}") + + endpoint_map[teacher.name].append(endpoint) + records.append( + [ + replica.display_name, + str(target_node_rank), + local_cuda_visible_devices, + endpoint, + f"{endpoint}/{health_path}", + f"{endpoint}/{model_info_path}", + str(len(command)), + *command, + ] + ) + return endpoint_map, student_num_workers, student_local_num_workers, records + + +def _expand_teacher_replicas( + teachers: list[OPDTeacherConfig], +) -> list[TeacherReplicaRequest]: + return [ + TeacherReplicaRequest( + teacher=teacher, + replica_index=replica_index, + ) + for teacher in teachers + for replica_index in range(teacher.num_replicas) + ] + + +def _allocate_teacher_devices( + replica_requests: list[TeacherReplicaRequest], + *, + node_count: int, + node_rank: int, + gpus_per_node: int, +) -> tuple[dict[tuple[str, int], tuple[int, list[int]]], int]: + """Allocate high-rank GPUs to local Teachers and leave the rest to Student. + + Teacher replicas are processed by descending ``num_workers``. Requests + with the same size keep their expanded config order. Each replica is placed + wholly on the highest-rank node that has enough free GPUs, using that + node's highest free local device ordinals. Replicas without + ``launch_config`` are externally managed and do not consume cluster GPUs. + + Args: + replica_requests: Teacher replicas in expanded config order. + node_count: Number of homogeneous nodes in the cluster. + node_rank: Rank of the current node. + gpus_per_node: Number of local GPUs available on every node. + + Returns: + A pair of ``(placements, student_local_num_workers)``: + + - ``placements`` maps each local ``(Teacher name, replica index)`` to + ``(target_node_rank, local_device_ordinals)``. + - ``student_local_num_workers`` is the number of remaining GPUs assigned + to Student on the current node. + + For inputs equivalent to: + + .. code-block:: python + + node_count = 4 + node_rank = 3 + gpus_per_node = 8 + replica_requests = [ + ("teacher1", 0, 4), + ("teacher2", 0, 2), + ] + + the returned values are: + + .. code-block:: python + + placements = { + ("teacher1", 0): (3, [4, 5, 6, 7]), + ("teacher2", 0): (3, [2, 3]), + } + student_local_num_workers = 2 + + Raises: + ValueError: If a Teacher cannot fit wholly on one node, or if Teacher + allocation leaves no GPU for Student. + """ + free_devices_by_node = [ + list(range(gpus_per_node)) for _ in range(node_count) + ] + local_teacher_requests: list[tuple[TeacherReplicaRequest, int]] = [] + for replica in replica_requests: + launch_config = replica.teacher.launch_config + if launch_config is not None: + local_teacher_requests.append((replica, launch_config.num_workers)) + local_teacher_requests.sort(key=lambda request: request[1], reverse=True) + + placements: dict[tuple[str, int], tuple[int, list[int]]] = {} + for replica, num_workers in local_teacher_requests: + if num_workers > gpus_per_node: + raise ValueError( + f"Teacher replica {replica.display_name!r} requests {num_workers} workers, " + f"but each node has only {gpus_per_node} GPUs" + ) + + target_node_rank = next( + ( + node_rank + for node_rank in range(node_count - 1, -1, -1) + if len(free_devices_by_node[node_rank]) >= num_workers + ), + None, + ) + if target_node_rank is None: + remaining_devices = sum( + len(local_devices) for local_devices in free_devices_by_node + ) + raise ValueError( + f"Teacher replica {replica.display_name!r} requests {num_workers} workers, " + "but no single node has enough remaining GPUs; " + f"{remaining_devices} GPUs remain across the cluster" + ) + + free_devices = free_devices_by_node[target_node_rank] + local_devices = free_devices[-num_workers:] + del free_devices[-num_workers:] + placements[replica.key] = (target_node_rank, local_devices) + + if not any(free_devices_by_node): + raise ValueError("Teacher allocation leaves no GPUs for Student workers") + student_local_num_workers = len(free_devices_by_node[node_rank]) + return placements, student_local_num_workers + + +def _allocate_teacher_server_port( + base_port: int, + *, + target_node_rank: int, + used_node_ports: set[tuple[int, int]], + replica_name: str, +) -> int: + server_port = base_port + while (target_node_rank, server_port) in used_node_ports: + server_port += 1 + if server_port > 65535: + raise ValueError( + f"Teacher replica {replica_name!r} cannot allocate a free port " + f"on node {target_node_rank} starting from {base_port}" + ) + used_node_ports.add((target_node_rank, server_port)) + return server_port + + +def _get_teacher_server_paths( + backend: Literal["sglang", "lmdeploy"], +) -> tuple[str, str]: + if backend == "sglang": + return "health_generate", "get_model_info" + return "health", "v1/models" + + +def _build_sglang_command( + teacher: OPDTeacherConfig, + *, + server_port: int, +) -> list[str]: + config = teacher.launch_config + assert config is not None + tensor_parallel_size = config.tensor_parallel_size + if config.expert_parallel_size > 1: + tensor_parallel_size = config.expert_parallel_size + + command = [ + sys.executable, + "-m", + "sglang.launch_server", + "--model-path", + str(config.model_path), + "--host", + "0.0.0.0", + "--port", + str(server_port), + "--dtype", + config.dtype, + "--tp", + str(tensor_parallel_size), + "--ep", + str(config.expert_parallel_size), + "--mem-fraction-static", + str(config.gpu_memory_utilization), + ] + if config.log_level is not None: + command.extend(["--log-level", config.log_level]) + if config.context_length is not None: + command.extend(["--context-length", str(config.context_length)]) + if config.max_batch_size is not None: + command.extend(["--max-running-requests", str(config.max_batch_size)]) + if config.chunked_prefill_size is not None: + command.extend(["--chunked-prefill-size", str(config.chunked_prefill_size)]) + return command + + +def _build_lmdeploy_command( + teacher: OPDTeacherConfig, + *, + server_port: int, +) -> list[str]: + config = teacher.launch_config + assert config is not None + data_parallel_size = ( + config.expert_parallel_size + if config.expert_parallel_size > 1 + else 1 + ) + command = [ + sys.executable, + "-m", + "lmdeploy", + "serve", + "api_server", + str(config.model_path), + "--backend", + "pytorch", + "--role", + "Hybrid", + "--logprobs-mode", + "raw_logprobs", + "--server-name", + "0.0.0.0", + "--server-port", + str(server_port), + "--dtype", + config.dtype, + "--tp", + str(config.tensor_parallel_size), + "--ep", + str(config.expert_parallel_size), + "--dp", + str(data_parallel_size), + "--cache-max-entry-count", + str(config.gpu_memory_utilization), + ] + if config.log_level is not None: + command.extend(["--log-level", config.log_level.upper()]) + if config.context_length is not None: + command.extend(["--session-len", str(config.context_length)]) + if config.max_batch_size is not None: + command.extend(["--max-batch-size", str(config.max_batch_size)]) + if config.max_prefill_token_num is not None: + command.extend(["--max-prefill-token-num", str(config.max_prefill_token_num)]) + if teacher.enable_prefix_caching: + command.append("--enable-prefix-caching") + return command + + +def _write_teacher_records( + endpoint_map: dict[str, list[str]], + student_num_workers: int, + student_local_num_workers: int, + records: list[list[str]], +) -> None: + endpoint_map_json = json.dumps( + endpoint_map, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + fields = [ + str(len(records)), + endpoint_map_json, + str(student_num_workers), + str(student_local_num_workers), + ] + for record in records: + fields.extend(record) + + payload = "\0".join(fields) + "\0" + sys.stdout.buffer.write(payload.encode("utf-8")) + + +def _main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("config_path") + parser.add_argument("backend", choices=("sglang", "lmdeploy")) + args = parser.parse_args() + endpoint_map, student_num_workers, student_local_num_workers, records = ( + build_teacher_launch_server_commands(args.config_path, args.backend) + ) + _write_teacher_records( + endpoint_map, + student_num_workers, + student_local_num_workers, + records, + ) + + +if __name__ == "__main__": + _main() diff --git a/recipe/on_policy_distillation/config/rl_dapo_math_mopd.py b/recipe/on_policy_distillation/config/rl_dapo_math_mopd.py new file mode 100644 index 0000000000..bef113ab89 --- /dev/null +++ b/recipe/on_policy_distillation/config/rl_dapo_math_mopd.py @@ -0,0 +1,397 @@ +import json +import os + +from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig +from xtuner.v1.data_proto.rl_data import SampleParams +from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig +from xtuner.v1.datasets.rl_tokenize_fn import RLQwen3VLTokenizeFnConfig +from xtuner.v1.model import Qwen3VLDense2BConfig +from xtuner.v1.rl.advantage import GRPOAdvantageConfig +from xtuner.v1.rl.agent_loop import SingleTurnAgentLoopConfig +from xtuner.v1.rl.agent_loop_manager import ( + AgentLoopManagerConfig, + SamplerConfig, + SyncProduceStrategyConfig, + TaskSpecConfig, +) +from xtuner.v1.rl.evaluator import EvaluatorConfig +from xtuner.v1.rl.judger import ComposedJudgerConfig, GEO3KJudgerConfig, GSM8KJudgerConfig +from xtuner.v1.rl.loss import GRPOLossConfig +from xtuner.v1.rl.on_policy_distillation import ( + OPDConfig, + OPDTeacherConfig, + OPDTeacherLaunchConfig, +) +from xtuner.v1.rl.replay_buffer import SyncReplayBufferConfig +from xtuner.v1.rl.rollout.worker import RolloutConfig +from xtuner.v1.rl.trainer import WorkerConfig +from xtuner.v1.rl.utils import ( + AcceleratorResourcesConfig, + CPUResourcesConfig, +) +from xtuner.v1.train.rl_trainer import RLColocateTrainerConfig + + +work_dir = os.environ["WORK_DIR"] +model_path = os.environ["MODEL_PATH"] +gsm8k_teacher_model_path = os.environ["GSM8K_TEACHER_MODEL_PATH"] +geo3k_teacher_model_path = os.environ["GEO3K_TEACHER_MODEL_PATH"] +meta_data_path = os.environ["DATA_PATH"] +eval_meta_data_path = os.environ.get("EVAL_DATA_PATH", "") +NNODE = int(os.environ.get("WORLD_SIZE", "1")) + + +def _as_list(value): + return value if isinstance(value, list) else [value] + + +EVAL_DATA_SOURCE_TAGS = { + "openai/gsm8k": "gsm8k", + "hiyouga/geometry3k": "geo3k", +} + + +def mopd_compute_metric(samples): + scores_by_source = {data_source: [] for data_source in EVAL_DATA_SOURCE_TAGS} + for sample in samples: + data_source = sample.extra_fields.get("origin_data_source") + if data_source not in scores_by_source: + raise ValueError(f"Unexpected evaluation data source: {data_source!r}") + + reward = sample.reward or {} + if "score" not in reward: + raise ValueError(f"Missing reward score for evaluation data source: {data_source!r}") + scores_by_source[data_source].append(float(reward["score"])) + + missing_sources = [data_source for data_source, scores in scores_by_source.items() if not scores] + if missing_sources: + raise ValueError(f"Missing evaluation samples for data sources: {missing_sources}") + + all_scores = [score for scores in scores_by_source.values() for score in scores] + metrics = {"score": sum(all_scores) / len(all_scores)} + for data_source, tag in EVAL_DATA_SOURCE_TAGS.items(): + source_scores = scores_by_source[data_source] + metrics[f"{tag}/score"] = sum(source_scores) / len(source_scores) + return metrics + + +# Training shape aligned with verl PR #6051: +# examples/on_policy_distillation_trainer/run_qwen3_mopd_gsm8k_geo3k.sh. +# Teacher roles and model families follow the GSM8K/Geo3K experiment. +experimental_name = "dapo_math_mopd" +total_epochs = 15 +total_train_steps_env = os.environ.get("TOTAL_TRAIN_STEPS") +total_train_steps = int(total_train_steps_env) if total_train_steps_env is not None else None +train_batch_size = 256 +teacher_num_replicas = int(os.environ.get("MOPD_TEACHER_NUM_REPLICAS", "2")) +teacher_replica_num_workers = int(os.environ.get("MOPD_TEACHER_REPLICA_NUM_WORKERS", "1")) +enable_prefix_caching = os.environ.get("MOPD_ENABLE_PREFIX_CACHING", "0") == "1" +prompt_repeat_k = 1 +rollout_tp_size = 1 +rollout_ep_size = 1 +max_prompt_length = 1024 +max_response_length = 2048 +pack_max_length = max_prompt_length + max_response_length +max_num_tokens = pack_max_length +train_optimizer_steps = 1 +enable_evaluate = bool(eval_meta_data_path) +evaluate_step = 5 +eval_prompt_repeat_k = 1 +checkpoint_interval = 200 + +# 1. resources: four colocated Student workers, plus one GPU per Teacher replica. +resources = AcceleratorResourcesConfig( + accelerator="GPU", + num_workers=4 * NNODE, + num_cpus_per_worker=12, + cpu_memory_per_worker=16 * 1024**3, # 16 GB +) + +# 2. rollout +rollout_config = RolloutConfig( + env=experimental_name, + device=resources.accelerator, + model_path=model_path, + dtype="bfloat16", + tensor_parallel_size=rollout_tp_size, + expert_parallel_size=rollout_ep_size, + gpu_memory_utilization=0.6, + context_length=max_response_length + max_prompt_length, + enable_return_routed_experts=False, + rollout_max_batch_size_per_instance=2048, +) + +# 3. train worker +lr_cfg = LRConfig(lr_type="constant", warmup_ratio=0, lr_min=1e-6) +fsdp_cfg = FSDPConfig( + torch_compile=False, + cpu_offload=False, + ep_size=1, + reduce_dtype="float32", +) +model_cfg = Qwen3VLDense2BConfig() +if hasattr(model_cfg, "balancing_loss_cfg"): + model_cfg.balancing_loss_cfg = None +if hasattr(model_cfg, "z_loss_cfg"): + model_cfg.z_loss_cfg = None +optim_cfg = AdamWConfig(lr=1e-6, foreach=False, weight_decay=0.1, betas=(0.9, 0.98)) +loss_cfg = GRPOLossConfig( + policy_loss_cfg={ + "cliprange_high": 0.2, + "cliprange_low": 0.2, + "loss_type": "vanilla", + "clip_ratio_c": 10.0, + "log_prob_diff_min": -20.0, + "log_prob_diff_max": 20.0, + }, + ignore_idx=-100, + use_kl_loss=False, + kl_loss_coef=0.0, + kl_loss_type="low_var_kl", + mode="chunk", + chunk_size=512, +) +train_worker_cfg = WorkerConfig( + model_cfg=model_cfg, + load_from=model_path, + optim_cfg=optim_cfg, + loss_cfg=loss_cfg, + lr_cfg=lr_cfg, + fsdp_cfg=fsdp_cfg, + sp_size=1, + optimizer_steps=train_optimizer_steps, + pack_max_length=pack_max_length, +) + +# 4. train agent loop manager +with open(meta_data_path, "r", encoding="utf-8") as f: + ds_collections = json.load(f) + +train_dataset_cfg = [] +for name, data in ds_collections.items(): + annotations = _as_list(data["annotation"]) + for annotation in annotations: + train_dataset_cfg.append( + { + "dataset": DatasetConfig( + name=name, + anno_path=annotation, + media_root=data.get("media_root", ""), + sample_ratio=data.get("sample_ratio", 1.0), + class_name="VLMJsonlDataset", + ), + "tokenize_fn": RLQwen3VLTokenizeFnConfig( + processor_path=model_path, + max_length=max_prompt_length, + system_message=data.get("system_message", None), + chat_template="qwen3-vl", + add_generation_prompt=True, + enable_thinking=True, + ), + } + ) + +dataloader_cfg = DataloaderConfig( + dataset_config_list=train_dataset_cfg, + num_workers=8, + pack_max_length=pack_max_length, + collator="fake_collator", + pack_level="none", +) +sampler_config = SamplerConfig( + dataloader_cfg=dataloader_cfg, + prompt_repeat_k=prompt_repeat_k, +) +training_sample_params = SampleParams( + max_tokens=max_response_length, + top_k=0, + top_p=1.0, + temperature=1.0, + min_tokens=0, + skip_special_tokens=False, + return_routed_experts=False, +) +agent_loop_config = SingleTurnAgentLoopConfig( + hf_checkpoint=model_path, + sample_params=training_sample_params, +) +produce_strategy_config = SyncProduceStrategyConfig() +agent_loop_manager_cfg = AgentLoopManagerConfig( + tasks=TaskSpecConfig( + task_name="train_task", + agent_loop_config=agent_loop_config, + judger_config=ComposedJudgerConfig( + branches={ + "openai/gsm8k": GSM8KJudgerConfig( + judger_name="openai/gsm8k", + cpu_resources=CPUResourcesConfig( + num_workers=1, + num_cpus_per_worker=1, + ), + ), + "hiyouga/geometry3k": GEO3KJudgerConfig( + cpu_resources=CPUResourcesConfig( + num_workers=1, + num_cpus_per_worker=1, + ), + ), + } + ), + produce_strategy_config=produce_strategy_config, + sampler_config=sampler_config, + ), +) + +# 5. evaluation +eval_agent_loop_manager_cfg = None +evaluator_config = None +if enable_evaluate: + with open(eval_meta_data_path, "r", encoding="utf-8") as f: + eval_ds_collections = json.load(f) + + eval_dataset_cfg = [] + for name, data in eval_ds_collections.items(): + annotations = _as_list(data["annotation"]) + for annotation in annotations: + eval_dataset_cfg.append( + { + "dataset": DatasetConfig( + name=name, + anno_path=annotation, + media_root=data.get("media_root", ""), + sample_ratio=data.get("sample_ratio", 1.0), + class_name="VLMJsonlDataset", + ), + "tokenize_fn": RLQwen3VLTokenizeFnConfig( + processor_path=model_path, + max_length=max_prompt_length, + system_message=data.get("system_message", None), + chat_template="qwen3-vl", + add_generation_prompt=True, + enable_thinking=True, + ignore_multimodal_info=True, + ), + } + ) + + eval_judger_config = ComposedJudgerConfig( + branches={ + "openai/gsm8k": GSM8KJudgerConfig( + judger_name="openai/gsm8k", + cpu_resources=CPUResourcesConfig( + num_workers=1, + num_cpus_per_worker=1, + ), + ), + "hiyouga/geometry3k": GEO3KJudgerConfig( + cpu_resources=CPUResourcesConfig( + num_workers=1, + num_cpus_per_worker=1, + ), + ), + } + ) + eval_dataloader_cfg = DataloaderConfig( + dataset_config_list=eval_dataset_cfg, + num_workers=8, + pack_max_length=pack_max_length, + collator="fake_collator", + pack_level="none", + ) + eval_sampler_config = SamplerConfig( + dataloader_cfg=eval_dataloader_cfg, + prompt_repeat_k=eval_prompt_repeat_k, + ) + evaluation_sample_params = SampleParams( + max_tokens=max_response_length, + top_k=0, + top_p=1.0, + temperature=1.0, + min_tokens=0, + skip_special_tokens=False, + return_routed_experts=False, + ) + eval_agent_loop_config = SingleTurnAgentLoopConfig( + hf_checkpoint=model_path, + sample_params=evaluation_sample_params, + ) + eval_agent_loop_manager_cfg = AgentLoopManagerConfig( + tasks=TaskSpecConfig( + task_name="eval_task", + agent_loop_config=eval_agent_loop_config, + judger_config=eval_judger_config, + sampler_config=eval_sampler_config, + ), + ) + evaluator_config = EvaluatorConfig(compute_metric_func=mopd_compute_metric) + +# 6. multi-teacher pure on-policy distillation +# +# This is the only topology block that needs to be edited for an experiment. +# Every training record's data_source must have an entry in +# data_source_teacher_map. Teacher model paths are read from +# GSM8K_TEACHER_MODEL_PATH and GEO3K_TEACHER_MODEL_PATH. +opd_config = OPDConfig( + mode="pg-opd", + task_adv_weight=0.0, + opd_adv_weight=1.0, + teachers=[ + OPDTeacherConfig( + name="gsm8k_teacher", + num_replicas=teacher_num_replicas, + enable_prefix_caching=enable_prefix_caching, + launch_config=OPDTeacherLaunchConfig( + model_path=gsm8k_teacher_model_path, + num_workers=teacher_replica_num_workers, + server_port=13141, + tensor_parallel_size=teacher_replica_num_workers, + expert_parallel_size=1, + context_length=max_num_tokens, + max_batch_size=max_num_tokens, + gpu_memory_utilization=0.8, + ), + ), + OPDTeacherConfig( + name="geo3k_teacher", + num_replicas=teacher_num_replicas, + enable_prefix_caching=enable_prefix_caching, + launch_config=OPDTeacherLaunchConfig( + model_path=geo3k_teacher_model_path, + num_workers=teacher_replica_num_workers, + server_port=13142, + tensor_parallel_size=teacher_replica_num_workers, + expert_parallel_size=1, + context_length=max_num_tokens, + max_batch_size=max_num_tokens, + gpu_memory_utilization=0.8, + ), + ), + ], + data_source_teacher_map={ + "openai/gsm8k": "gsm8k_teacher", + "hiyouga/geometry3k": "geo3k_teacher", + }, +) + +trainer = RLColocateTrainerConfig( + resources=resources, + train_worker_cfg=train_worker_cfg, # TODO: uniform naming of cfg and config + rollout_config=rollout_config, + tokenizer_path=model_path, + replay_buffer_config=SyncReplayBufferConfig(), + agent_loop_manager_cfg=agent_loop_manager_cfg, + eval_agent_loop_manager_cfg=eval_agent_loop_manager_cfg, + evaluator_config=evaluator_config, + load_from=model_path, + train_batch_size=train_batch_size, + advantage_estimator_config=GRPOAdvantageConfig(eps=1e-8), + opd_config=opd_config, + enable_evaluate=enable_evaluate, + enable_initial_evaluate=enable_evaluate, + evaluate_step=evaluate_step, + total_train_steps=total_train_steps, + total_epochs=total_epochs, + work_dir=work_dir, + seed=1234, + debug_rollout=False, +) diff --git a/recipe/on_policy_distillation/config/rl_dapo_math_mopd_async.py b/recipe/on_policy_distillation/config/rl_dapo_math_mopd_async.py new file mode 100644 index 0000000000..5abf0a1681 --- /dev/null +++ b/recipe/on_policy_distillation/config/rl_dapo_math_mopd_async.py @@ -0,0 +1,398 @@ +import json +import os + +from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig +from xtuner.v1.data_proto.rl_data import SampleParams +from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig +from xtuner.v1.datasets.rl_tokenize_fn import RLQwen3VLTokenizeFnConfig +from xtuner.v1.model import Qwen3VLDense2BConfig +from xtuner.v1.rl.advantage import GRPOAdvantageConfig +from xtuner.v1.rl.agent_loop import SingleTurnAgentLoopConfig +from xtuner.v1.rl.agent_loop_manager import ( + AgentLoopManagerConfig, + AsyncProduceStrategyConfig, + SamplerConfig, + TaskSpecConfig, +) +from xtuner.v1.rl.evaluator import EvaluatorConfig +from xtuner.v1.rl.judger import ComposedJudgerConfig, GEO3KJudgerConfig, GSM8KJudgerConfig +from xtuner.v1.rl.loss import GRPOLossConfig +from xtuner.v1.rl.on_policy_distillation import ( + OPDConfig, + OPDTeacherConfig, + OPDTeacherLaunchConfig, +) +from xtuner.v1.rl.replay_buffer import AsyncReplayBufferConfig +from xtuner.v1.rl.rollout.worker import RolloutConfig +from xtuner.v1.rl.trainer import RolloutImportanceSampling, WorkerConfig +from xtuner.v1.rl.utils import ( + AcceleratorResourcesConfig, + CPUResourcesConfig, +) +from xtuner.v1.train.rl_trainer import RLColocateTrainerConfig + + +work_dir = os.environ["WORK_DIR"] +model_path = os.environ["MODEL_PATH"] +gsm8k_teacher_model_path = os.environ["GSM8K_TEACHER_MODEL_PATH"] +geo3k_teacher_model_path = os.environ["GEO3K_TEACHER_MODEL_PATH"] +meta_data_path = os.environ["DATA_PATH"] +eval_meta_data_path = os.environ.get("EVAL_DATA_PATH", "") +NNODE = int(os.environ.get("WORLD_SIZE", "1")) + + +def _as_list(value): + return value if isinstance(value, list) else [value] + + +EVAL_DATA_SOURCE_TAGS = { + "openai/gsm8k": "gsm8k", + "hiyouga/geometry3k": "geo3k", +} + + +def mopd_compute_metric(samples): + scores_by_source = {data_source: [] for data_source in EVAL_DATA_SOURCE_TAGS} + for sample in samples: + data_source = sample.extra_fields.get("origin_data_source") + if data_source not in scores_by_source: + raise ValueError(f"Unexpected evaluation data source: {data_source!r}") + + reward = sample.reward or {} + if "score" not in reward: + raise ValueError(f"Missing reward score for evaluation data source: {data_source!r}") + scores_by_source[data_source].append(float(reward["score"])) + + missing_sources = [data_source for data_source, scores in scores_by_source.items() if not scores] + if missing_sources: + raise ValueError(f"Missing evaluation samples for data sources: {missing_sources}") + + all_scores = [score for scores in scores_by_source.values() for score in scores] + metrics = {"score": sum(all_scores) / len(all_scores)} + for data_source, tag in EVAL_DATA_SOURCE_TAGS.items(): + source_scores = scores_by_source[data_source] + metrics[f"{tag}/score"] = sum(source_scores) / len(source_scores) + return metrics + + +# Training shape aligned with verl PR #6051: +# examples/on_policy_distillation_trainer/run_qwen3_mopd_gsm8k_geo3k.sh. +# Teacher roles and model families follow the GSM8K/Geo3K experiment. +experimental_name = "dapo_math_mopd_async" +total_epochs = 15 +train_batch_size = 128 +prompt_repeat_k = 1 +rollout_tp_size = 1 +rollout_ep_size = 1 +max_prompt_length = 1024 +max_response_length = 2048 +pack_max_length = max_prompt_length + max_response_length +max_num_tokens = pack_max_length +train_optimizer_steps = 1 +enable_evaluate = bool(eval_meta_data_path) +evaluate_step = 5 +eval_prompt_repeat_k = 1 +checkpoint_interval = 200 + +# 1. resources: four colocated Student workers, plus one GPU per Teacher. +resources = AcceleratorResourcesConfig( + accelerator="GPU", + num_workers=4 * NNODE, + num_cpus_per_worker=12, + cpu_memory_per_worker=16 * 1024**3, # 16 GB +) + +# 2. rollout +rollout_config = RolloutConfig( + env=experimental_name, + device=resources.accelerator, + model_path=model_path, + dtype="bfloat16", + tensor_parallel_size=rollout_tp_size, + expert_parallel_size=rollout_ep_size, + gpu_memory_utilization=0.6, + context_length=max_response_length + max_prompt_length, + enable_return_routed_experts=False, + rollout_max_batch_size_per_instance=2048, +) + +# 3. train worker +lr_cfg = LRConfig(lr_type="constant", warmup_ratio=0, lr_min=1e-6) +fsdp_cfg = FSDPConfig( + torch_compile=False, + cpu_offload=False, + ep_size=1, + reduce_dtype="float32", +) +model_cfg = Qwen3VLDense2BConfig() +if hasattr(model_cfg, "balancing_loss_cfg"): + model_cfg.balancing_loss_cfg = None +if hasattr(model_cfg, "z_loss_cfg"): + model_cfg.z_loss_cfg = None +optim_cfg = AdamWConfig(lr=1e-6, foreach=False, weight_decay=0.1, betas=(0.9, 0.98)) +loss_cfg = GRPOLossConfig( + policy_loss_cfg={ + "cliprange_high": 0.2, + "cliprange_low": 0.2, + "loss_type": "vanilla", + "clip_ratio_c": 10.0, + "log_prob_diff_min": -20.0, + "log_prob_diff_max": 20.0, + }, + ignore_idx=-100, + use_kl_loss=False, + kl_loss_coef=0.0, + kl_loss_type="low_var_kl", + mode="chunk", + chunk_size=512, + rollout_is=RolloutImportanceSampling( + rollout_is_level="token", + rollout_is_mode="both", + rollout_is_threshold=(5, 0.5), + rollout_is_mask_threshold=(5, 0.5), + rollout_is_veto_threshold=(20, 0), + ), +) +train_worker_cfg = WorkerConfig( + model_cfg=model_cfg, + load_from=model_path, + optim_cfg=optim_cfg, + loss_cfg=loss_cfg, + lr_cfg=lr_cfg, + fsdp_cfg=fsdp_cfg, + sp_size=1, + optimizer_steps=train_optimizer_steps, + pack_max_length=pack_max_length, +) + +# 4. train agent loop manager +with open(meta_data_path, "r", encoding="utf-8") as f: + ds_collections = json.load(f) + +train_dataset_cfg = [] +for name, data in ds_collections.items(): + annotations = _as_list(data["annotation"]) + for annotation in annotations: + train_dataset_cfg.append( + { + "dataset": DatasetConfig( + name=name, + anno_path=annotation, + media_root=data.get("media_root", ""), + sample_ratio=data.get("sample_ratio", 1.0), + class_name="VLMJsonlDataset", + ), + "tokenize_fn": RLQwen3VLTokenizeFnConfig( + processor_path=model_path, + max_length=max_prompt_length, + system_message=data.get("system_message", None), + chat_template="qwen3-vl", + add_generation_prompt=True, + enable_thinking=True, + ), + } + ) + +dataloader_cfg = DataloaderConfig( + dataset_config_list=train_dataset_cfg, + num_workers=8, + pack_max_length=pack_max_length, + collator="fake_collator", + pack_level="none", +) +sampler_config = SamplerConfig( + dataloader_cfg=dataloader_cfg, + prompt_repeat_k=prompt_repeat_k, +) +training_sample_params = SampleParams( + max_tokens=max_response_length, + top_k=0, + top_p=1.0, + temperature=1.0, + min_tokens=0, + skip_special_tokens=False, + return_routed_experts=False, +) +agent_loop_config = SingleTurnAgentLoopConfig( + hf_checkpoint=model_path, + sample_params=training_sample_params, +) +produce_strategy_config = AsyncProduceStrategyConfig( + over_sample_threshold=1.0, + enable_partial_rollout=True, + max_staleness=2, +) +agent_loop_manager_cfg = AgentLoopManagerConfig( + tasks=TaskSpecConfig( + task_name="train_task", + agent_loop_config=agent_loop_config, + judger_config=ComposedJudgerConfig( + branches={ + "openai/gsm8k": GSM8KJudgerConfig( + judger_name="openai/gsm8k", + cpu_resources=CPUResourcesConfig( + num_workers=1, + num_cpus_per_worker=1, + ), + ), + "hiyouga/geometry3k": GEO3KJudgerConfig( + cpu_resources=CPUResourcesConfig( + num_workers=1, + num_cpus_per_worker=1, + ), + ), + } + ), + produce_strategy_config=produce_strategy_config, + sampler_config=sampler_config, + ), +) + +# 5. evaluation +eval_agent_loop_manager_cfg = None +evaluator_config = None +if enable_evaluate: + with open(eval_meta_data_path, "r", encoding="utf-8") as f: + eval_ds_collections = json.load(f) + + eval_dataset_cfg = [] + for name, data in eval_ds_collections.items(): + annotations = _as_list(data["annotation"]) + for annotation in annotations: + eval_dataset_cfg.append( + { + "dataset": DatasetConfig( + name=name, + anno_path=annotation, + media_root=data.get("media_root", ""), + sample_ratio=data.get("sample_ratio", 1.0), + class_name="VLMJsonlDataset", + ), + "tokenize_fn": RLQwen3VLTokenizeFnConfig( + processor_path=model_path, + max_length=max_prompt_length, + system_message=data.get("system_message", None), + chat_template="qwen3-vl", + add_generation_prompt=True, + enable_thinking=True, + ignore_multimodal_info=True, + ), + } + ) + + eval_judger_config = ComposedJudgerConfig( + branches={ + "openai/gsm8k": GSM8KJudgerConfig( + judger_name="openai/gsm8k", + cpu_resources=CPUResourcesConfig( + num_workers=1, + num_cpus_per_worker=1, + ), + ), + "hiyouga/geometry3k": GEO3KJudgerConfig( + cpu_resources=CPUResourcesConfig( + num_workers=1, + num_cpus_per_worker=1, + ), + ), + } + ) + eval_dataloader_cfg = DataloaderConfig( + dataset_config_list=eval_dataset_cfg, + num_workers=8, + pack_max_length=pack_max_length, + collator="fake_collator", + pack_level="none", + ) + eval_sampler_config = SamplerConfig( + dataloader_cfg=eval_dataloader_cfg, + prompt_repeat_k=eval_prompt_repeat_k, + ) + evaluation_sample_params = SampleParams( + max_tokens=max_response_length, + top_k=0, + top_p=1.0, + temperature=1.0, + min_tokens=0, + skip_special_tokens=False, + return_routed_experts=False, + ) + eval_agent_loop_config = SingleTurnAgentLoopConfig( + hf_checkpoint=model_path, + sample_params=evaluation_sample_params, + ) + eval_agent_loop_manager_cfg = AgentLoopManagerConfig( + tasks=TaskSpecConfig( + task_name="eval_task", + agent_loop_config=eval_agent_loop_config, + judger_config=eval_judger_config, + sampler_config=eval_sampler_config, + ), + ) + evaluator_config = EvaluatorConfig(compute_metric_func=mopd_compute_metric) + +# 6. multi-teacher pure on-policy distillation +# +# This is the only topology block that needs to be edited for an experiment. +# Every training record's data_source must have an entry in +# data_source_teacher_map. Teacher model paths are read from +# GSM8K_TEACHER_MODEL_PATH and GEO3K_TEACHER_MODEL_PATH. +opd_config = OPDConfig( + mode="pg-opd", + task_adv_weight=0.0, + opd_adv_weight=1.0, + teachers=[ + OPDTeacherConfig( + name="gsm8k_teacher", + launch_config=OPDTeacherLaunchConfig( + model_path=gsm8k_teacher_model_path, + num_workers=1, + server_port=13141, + tensor_parallel_size=1, + expert_parallel_size=1, + context_length=max_num_tokens, + max_batch_size=max_num_tokens, + gpu_memory_utilization=0.8, + ), + ), + OPDTeacherConfig( + name="geo3k_teacher", + launch_config=OPDTeacherLaunchConfig( + model_path=geo3k_teacher_model_path, + num_workers=1, + server_port=13142, + tensor_parallel_size=1, + expert_parallel_size=1, + context_length=max_num_tokens, + max_batch_size=max_num_tokens, + gpu_memory_utilization=0.8, + ), + ), + ], + data_source_teacher_map={ + "openai/gsm8k": "gsm8k_teacher", + "hiyouga/geometry3k": "geo3k_teacher", + }, +) + +trainer = RLColocateTrainerConfig( + resources=resources, + train_worker_cfg=train_worker_cfg, # TODO: uniform naming of cfg and config + rollout_config=rollout_config, + tokenizer_path=model_path, + replay_buffer_config=AsyncReplayBufferConfig(), + agent_loop_manager_cfg=agent_loop_manager_cfg, + eval_agent_loop_manager_cfg=eval_agent_loop_manager_cfg, + evaluator_config=evaluator_config, + load_from=model_path, + train_batch_size=train_batch_size, + advantage_estimator_config=GRPOAdvantageConfig(eps=1e-8), + opd_config=opd_config, + enable_evaluate=enable_evaluate, + enable_initial_evaluate=enable_evaluate, + evaluate_step=evaluate_step, + total_epochs=total_epochs, + work_dir=work_dir, + seed=1234, + debug_rollout=False, +) diff --git a/recipe/on_policy_distillation/config/rl_dapo_math_opd.py b/recipe/on_policy_distillation/config/rl_dapo_math_opd.py new file mode 100644 index 0000000000..da55149dc6 --- /dev/null +++ b/recipe/on_policy_distillation/config/rl_dapo_math_opd.py @@ -0,0 +1,247 @@ +import os +from pathlib import Path + +from transformers import AutoTokenizer +from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig +from xtuner.v1.data_proto.rl_data import SampleParams +from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig +from xtuner.v1.datasets.rl_tokenize_fn import RLTextTokenizeFnConfig +from xtuner.v1.model import get_model_config_from_hf +from xtuner.v1.rl.advantage import GRPOAdvantageConfig +from xtuner.v1.rl.agent_loop import SingleTurnAgentLoopConfig +from xtuner.v1.rl.agent_loop_manager import ( + AgentLoopManagerConfig, + SamplerConfig, + SyncProduceStrategyConfig, + TaskSpecConfig, +) +from xtuner.v1.rl.evaluator import EvaluatorConfig +from xtuner.v1.rl.judger import DapoMathJudgerConfig +from xtuner.v1.rl.loss import GRPOLossConfig +from xtuner.v1.rl.on_policy_distillation import ( + OPDConfig, + OPDTeacherConfig, + OPDTeacherLaunchConfig, +) +from xtuner.v1.rl.replay_buffer import SyncReplayBufferConfig +from xtuner.v1.rl.rollout.worker import RolloutConfig +from xtuner.v1.rl.trainer import WorkerConfig +from xtuner.v1.rl.utils import ( + AcceleratorResourcesConfig, + CPUResourcesConfig, + get_eos_token, +) +from xtuner.v1.train.rl_trainer import RLColocateTrainerConfig + + +work_dir = os.environ["WORK_DIR"] +model_path = os.environ["MODEL_PATH"] +teacher_model_path = os.environ["TEACHER_MODEL_PATH"] +data_path = os.environ["DATA_PATH"] +eval_data_path = os.environ.get("EVAL_DATA_PATH", "") +NNODE = int(os.environ.get("WORLD_SIZE", "1")) + +# basic settings +experimental_name = "dapo_math_opd" +total_train_steps = 300 +train_batch_size = 16 +prompt_repeat_k = 4 +rollout_tp_size = 1 +rollout_ep_size = 1 +max_prompt_length = 2048 +max_response_length = 16384 +pack_max_length = max_prompt_length + max_response_length +train_optimizer_steps = 1 +enable_evaluate = bool(eval_data_path) +evaluate_step = 20 +eval_prompt_repeat_k = 16 + +# 1. resources +resources = AcceleratorResourcesConfig( + accelerator="GPU", + num_workers=4 * NNODE, + num_cpus_per_worker=12, + cpu_memory_per_worker=16 * 1024**3, # 16 GB +) + +# 2. rollout +rollout_config = RolloutConfig( + env=experimental_name, + device=resources.accelerator, + model_path=model_path, + dtype="bfloat16", + tensor_parallel_size=rollout_tp_size, + expert_parallel_size=rollout_ep_size, + gpu_memory_utilization=0.6, + context_length=max_response_length + max_prompt_length, + enable_return_routed_experts=False, + rollout_max_batch_size_per_instance=2048, +) + +# 3. train worker +lr_cfg = LRConfig(lr_type="constant", warmup_ratio=0, lr_min=1e-6) +fsdp_cfg = FSDPConfig( + torch_compile=False, + cpu_offload=False, + ep_size=1, + reduce_dtype="float32", +) +model_cfg = get_model_config_from_hf(Path(model_path)) +if hasattr(model_cfg, "balancing_loss_cfg"): + model_cfg.balancing_loss_cfg = None +if hasattr(model_cfg, "z_loss_cfg"): + model_cfg.z_loss_cfg = None +optim_cfg = AdamWConfig(lr=1e-6, foreach=False, weight_decay=0.1, betas=(0.9, 0.98)) +loss_cfg = GRPOLossConfig( + policy_loss_cfg={ + "cliprange_high": 0.28, + "cliprange_low": 0.2, + "loss_type": "vanilla", + "clip_ratio_c": 10.0, + "log_prob_diff_min": -20.0, + "log_prob_diff_max": 20.0, + }, + ignore_idx=-100, + use_kl_loss=False, + kl_loss_coef=0.0, + kl_loss_type="low_var_kl", + mode="chunk", + chunk_size=512, +) +train_worker_cfg = WorkerConfig( + model_cfg=model_cfg, + load_from=model_path, + optim_cfg=optim_cfg, + loss_cfg=loss_cfg, + lr_cfg=lr_cfg, + fsdp_cfg=fsdp_cfg, + sp_size=1, + optimizer_steps=train_optimizer_steps, + pack_max_length=pack_max_length, +) + +# 4. train agent loop manager +train_dataset = DatasetConfig(name=experimental_name, anno_path=data_path) +tokenizer_config = RLTextTokenizeFnConfig(max_length=max_prompt_length) +train_dataset_cfg = [{"dataset": train_dataset, "tokenize_fn": tokenizer_config}] +dataloader_cfg = DataloaderConfig( + dataset_config_list=train_dataset_cfg, + pack_max_length=pack_max_length, + collator="fake_collator", + pack_level="none", +) +sampler_config = SamplerConfig( + dataloader_cfg=dataloader_cfg, + prompt_repeat_k=prompt_repeat_k, +) +training_sample_params = SampleParams( + max_tokens=max_response_length, + top_k=0, + top_p=1.0, + temperature=1.0, + min_tokens=0, + skip_special_tokens=False, + return_routed_experts=False, +) +agent_loop_config = SingleTurnAgentLoopConfig( + hf_checkpoint=model_path, + sample_params=training_sample_params, +) +produce_strategy_config = SyncProduceStrategyConfig() +agent_loop_manager_cfg = AgentLoopManagerConfig( + tasks=TaskSpecConfig( + task_name="train_task", + agent_loop_config=agent_loop_config, + produce_strategy_config=produce_strategy_config, + sampler_config=sampler_config, + ), +) + +# 5. evaluation +eval_agent_loop_manager_cfg = None +evaluator_config = None +if enable_evaluate: + eos_token_id = get_eos_token(model_path) + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + eos_token = tokenizer.convert_ids_to_tokens(eos_token_id) + eval_judger_config = DapoMathJudgerConfig( + judger_name="aime_math", + cpu_resources=CPUResourcesConfig(num_workers=1, num_cpus_per_worker=1), + eos_token=eos_token, + enable_overlong_buffer=False, + ) + eval_dataset = DatasetConfig(name="aime", anno_path=eval_data_path, sample_ratio=1.0) + eval_dataset_cfg = [{"dataset": eval_dataset, "tokenize_fn": tokenizer_config}] + eval_dataloader_cfg = DataloaderConfig( + dataset_config_list=eval_dataset_cfg, + pack_max_length=pack_max_length, + collator="fake_collator", + pack_level="none", + ) + eval_sampler_config = SamplerConfig( + dataloader_cfg=eval_dataloader_cfg, + prompt_repeat_k=eval_prompt_repeat_k, + ) + evaluation_sample_params = SampleParams( + max_tokens=max_response_length, + top_k=0, + top_p=1.0, + temperature=1.0, + min_tokens=0, + skip_special_tokens=False, + return_routed_experts=False, + ) + eval_agent_loop_config = SingleTurnAgentLoopConfig( + hf_checkpoint=model_path, + sample_params=evaluation_sample_params, + ) + eval_agent_loop_manager_cfg = AgentLoopManagerConfig( + tasks=TaskSpecConfig( + task_name="eval_task", + agent_loop_config=eval_agent_loop_config, + judger_config=eval_judger_config, + sampler_config=eval_sampler_config, + ), + ) + evaluator_config = EvaluatorConfig() + +# 6. pure on-policy distillation +opd_config = OPDConfig( + mode="pg-opd", + task_adv_weight=0.0, + opd_adv_weight=1.0, + teachers=[ + OPDTeacherConfig( + name="teacher", + enable_prefix_caching=True, + launch_config=OPDTeacherLaunchConfig( + model_path=teacher_model_path, + num_workers=1, + server_port=13141, + ), + ) + ], + data_source_teacher_map={"math_dapo": "teacher"}, +) + +trainer = RLColocateTrainerConfig( + resources=resources, + train_worker_cfg=train_worker_cfg, # TODO: uniform naming of cfg and config + rollout_config=rollout_config, + tokenizer_path=model_path, + replay_buffer_config=SyncReplayBufferConfig(), + agent_loop_manager_cfg=agent_loop_manager_cfg, + eval_agent_loop_manager_cfg=eval_agent_loop_manager_cfg, + evaluator_config=evaluator_config, + load_from=model_path, + train_batch_size=train_batch_size, + advantage_estimator_config=GRPOAdvantageConfig(eps=1e-8), + opd_config=opd_config, + enable_evaluate=enable_evaluate, + enable_initial_evaluate=enable_evaluate, + evaluate_step=evaluate_step, + total_train_steps=total_train_steps, + work_dir=work_dir, + seed=1234, + debug_rollout=False, +) diff --git a/recipe/on_policy_distillation/scripts/launch_teacher_utils.sh b/recipe/on_policy_distillation/scripts/launch_teacher_utils.sh new file mode 100644 index 0000000000..ca47c95f78 --- /dev/null +++ b/recipe/on_policy_distillation/scripts/launch_teacher_utils.sh @@ -0,0 +1,267 @@ +#!/usr/bin/env bash + +# Source-only helpers for launching, waiting for, and stopping OPD Teacher servers. + +OPD_RECIPE_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd) + +TEACHER_NAMES=() +TEACHER_ENDPOINTS=() +TEACHER_HEALTH_URLS=() +TEACHER_MODEL_INFO_URLS=() +TEACHER_LOG_FILES=() +TEACHER_PIDS=() +STUDENT_CUDA_VISIBLE_DEVICES= +XTUNER_RL_NUM_WORKERS= + +start_single_teacher_server() { + _start_teacher_servers "$1" "$2" "$3" "1" +} + +start_teacher_servers() { + _start_teacher_servers "$1" "$2" "$3" "" +} + +wait_for_teacher_servers() { + local startup_timeout_s=$1 + local deadline=$((SECONDS + startup_timeout_s)) + local teacher_index + local pid + local log_file + local -a teacher_ready=() + local -a health_check_pids=() + local -a pending_names=() + + if (( ${#TEACHER_NAMES[@]} == 0 )); then + echo "No Teacher servers have been configured." >&2 + return 1 + fi + + for teacher_index in "${!TEACHER_NAMES[@]}"; do + teacher_ready[teacher_index]=0 + done + + while (( SECONDS < deadline )); do + for teacher_index in "${!TEACHER_NAMES[@]}"; do + pid="${TEACHER_PIDS[teacher_index]}" + log_file="${TEACHER_LOG_FILES[teacher_index]}" + if [[ -n "${pid}" ]] && ! kill -0 "${pid}" 2>/dev/null; then + echo "Teacher ${TEACHER_NAMES[teacher_index]} exited before becoming ready." >&2 + if [[ -n "${log_file}" ]]; then + tail -n 50 "${log_file}" >&2 || true + fi + return 1 + fi + done + + health_check_pids=() + for teacher_index in "${!TEACHER_NAMES[@]}"; do + if (( teacher_ready[teacher_index] )); then + continue + fi + curl -sf --max-time 2 \ + "${TEACHER_HEALTH_URLS[teacher_index]}" \ + >/dev/null 2>&1 & + health_check_pids[teacher_index]="$!" + done + + for teacher_index in "${!health_check_pids[@]}"; do + if wait "${health_check_pids[teacher_index]}"; then + teacher_ready[teacher_index]=1 + echo "Teacher ${TEACHER_NAMES[teacher_index]} is ready at ${TEACHER_ENDPOINTS[teacher_index]}" + fi + done + + pending_names=() + for teacher_index in "${!TEACHER_NAMES[@]}"; do + if (( ! teacher_ready[teacher_index] )); then + pending_names+=("${TEACHER_NAMES[teacher_index]}") + fi + done + if (( ${#pending_names[@]} == 0 )); then + return 0 + fi + + echo "Waiting for teachers: ${pending_names[*]}" + sleep 5 + done + + for teacher_index in "${!TEACHER_NAMES[@]}"; do + if (( teacher_ready[teacher_index] )); then + continue + fi + echo "Teacher ${TEACHER_NAMES[teacher_index]} did not become ready within ${startup_timeout_s}s." >&2 + log_file="${TEACHER_LOG_FILES[teacher_index]}" + if [[ -n "${log_file}" ]]; then + tail -n 50 "${log_file}" >&2 || true + fi + done + return 1 +} + +stop_teacher_servers() { + local attempt + local any_alive + local pid + + for pid in "${TEACHER_PIDS[@]}"; do + if [[ -n "${pid}" ]]; then + kill -TERM -- "-${pid}" 2>/dev/null || true + fi + done + + for ((attempt = 0; attempt < 30; attempt++)); do + any_alive=0 + for pid in "${TEACHER_PIDS[@]}"; do + if [[ -n "${pid}" ]] && kill -0 -- "-${pid}" 2>/dev/null; then + any_alive=1 + break + fi + done + if (( ! any_alive )); then + break + fi + sleep 1 + done + + for pid in "${TEACHER_PIDS[@]}"; do + if [[ -n "${pid}" ]]; then + if kill -0 -- "-${pid}" 2>/dev/null; then + kill -KILL -- "-${pid}" 2>/dev/null || true + fi + wait "${pid}" 2>/dev/null || true + fi + done + + TEACHER_PIDS=() +} + +_start_teacher_servers() { + local config_file=$1 + local backend=$2 + local work_dir=$3 + local expected_teacher_count=$4 + local teacher_count + local teacher_index + local teacher_field_offset=4 + local teacher_name + local teacher_target_node_rank + local teacher_local_devices + local teacher_endpoint + local teacher_command_arg_count + local teacher_log_file + local device + local student_local_num_workers + local -a teacher_fields=() + local -a teacher_command=() + + # Local non-rjob launches default to a single node. Preserve topology + # values supplied by distributed launchers. + export NODE_COUNT="${NODE_COUNT:-1}" + export NODE_RANK="${NODE_RANK:-0}" + + TEACHER_NAMES=() + TEACHER_ENDPOINTS=() + TEACHER_HEALTH_URLS=() + TEACHER_MODEL_INFO_URLS=() + TEACHER_LOG_FILES=() + TEACHER_PIDS=() + STUDENT_CUDA_VISIBLE_DEVICES= + XTUNER_RL_NUM_WORKERS= + mkdir -p "${work_dir}" + + # Command builder output: + # replica count, endpoint map JSON, Student worker count, current-node + # Student worker count, then repeated display name, target node rank, local devices, + # endpoint, health URL, model-info URL, + # command-argument count, and command argv. + unset XTUNER_OPD_TEACHER_ENDPOINTS_JSON + mapfile -d "" -t teacher_fields < <( + python "${OPD_RECIPE_DIR}/build_teacher_server_commands.py" \ + "${config_file}" "${backend}" + ) + export XTUNER_OPD_TEACHER_ENDPOINTS_JSON="${teacher_fields[1]}" + export XTUNER_RL_NUM_WORKERS="${teacher_fields[2]}" + student_local_num_workers=${teacher_fields[3]} + STUDENT_CUDA_VISIBLE_DEVICES= + for ((device = 0; device < student_local_num_workers; device++)); do + STUDENT_CUDA_VISIBLE_DEVICES+="${STUDENT_CUDA_VISIBLE_DEVICES:+,}${device}" + done + export STUDENT_CUDA_VISIBLE_DEVICES + + teacher_count=${teacher_fields[0]} + if [[ -n "${expected_teacher_count}" ]] && (( teacher_count != expected_teacher_count )); then + echo "Expected ${expected_teacher_count} Teacher, got ${teacher_count}." >&2 + return 1 + fi + if (( teacher_count == 0 )); then + echo "OPD config does not contain any Teachers: ${config_file}" >&2 + return 1 + fi + + echo "Teacher backend: ${backend}" + echo "Student workers: ${XTUNER_RL_NUM_WORKERS}" + echo "Student local GPUs on node ${NODE_RANK}: ${STUDENT_CUDA_VISIBLE_DEVICES:-none}" + for ((teacher_index = 0; teacher_index < teacher_count; teacher_index++)); do + teacher_name=${teacher_fields[teacher_field_offset]} + teacher_target_node_rank=${teacher_fields[teacher_field_offset + 1]} + teacher_local_devices=${teacher_fields[teacher_field_offset + 2]} + teacher_endpoint=${teacher_fields[teacher_field_offset + 3]} + teacher_command_arg_count=${teacher_fields[teacher_field_offset + 6]} + + TEACHER_NAMES+=("${teacher_name}") + TEACHER_ENDPOINTS+=("${teacher_endpoint}") + TEACHER_HEALTH_URLS+=("${teacher_fields[teacher_field_offset + 4]}") + TEACHER_MODEL_INFO_URLS+=("${teacher_fields[teacher_field_offset + 5]}") + + teacher_field_offset=$((teacher_field_offset + 7)) + teacher_command=( + "${teacher_fields[@]:teacher_field_offset:teacher_command_arg_count}" + ) + teacher_field_offset=$((teacher_field_offset + teacher_command_arg_count)) + + if (( teacher_command_arg_count == 0 )); then + if [[ -n "${expected_teacher_count}" ]]; then + echo "Teacher ${teacher_name} must define launch_config for local startup." >&2 + return 1 + fi + echo "Using externally managed Teacher ${teacher_name} at ${teacher_endpoint}" + TEACHER_LOG_FILES+=("") + TEACHER_PIDS+=("") + continue + fi + + if (( NODE_RANK != teacher_target_node_rank )); then + echo \ + "Teacher ${teacher_name} is assigned to node ${teacher_target_node_rank};" \ + "skipping node ${NODE_RANK}." + TEACHER_LOG_FILES+=("") + TEACHER_PIDS+=("") + continue + fi + + if [[ -n "${expected_teacher_count}" ]]; then + teacher_log_file="${work_dir}/teacher.log" + else + teacher_log_file="${work_dir}/teacher_${teacher_index}_node_${NODE_RANK}.log" + fi + TEACHER_LOG_FILES+=("${teacher_log_file}") + + echo "Starting Teacher ${teacher_name}" + echo "Teacher node: ${NODE_RANK}" + echo "Teacher local GPUs: ${teacher_local_devices}" + echo "Teacher endpoint: ${teacher_endpoint}" + echo "Teacher log: ${teacher_log_file}" + + setsid env \ + PYTHONUNBUFFERED=1 \ + CUDA_VISIBLE_DEVICES="${teacher_local_devices}" \ + "${teacher_command[@]}" \ + >"${teacher_log_file}" 2>&1 & + TEACHER_PIDS+=("$!") + done + + if (( teacher_field_offset != ${#teacher_fields[@]} )); then + echo "Teacher command builder returned malformed records." >&2 + return 1 + fi +} diff --git a/recipe/on_policy_distillation/scripts/run_sampled_token_mopd.sh b/recipe/on_policy_distillation/scripts/run_sampled_token_mopd.sh new file mode 100755 index 0000000000..e9f7e3c927 --- /dev/null +++ b/recipe/on_policy_distillation/scripts/run_sampled_token_mopd.sh @@ -0,0 +1,74 @@ +#!/usr/bin/env bash + +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd -- "${SCRIPT_DIR}/../../.." && pwd) +source "${SCRIPT_DIR}/launch_teacher_utils.sh" + +export STUDENT_MODEL_PATH=${1:?"Usage: $0 STUDENT_MODEL_PATH DATA_PATH"} +export MODEL_PATH="${STUDENT_MODEL_PATH}" +export DATA_PATH=${2:?"Usage: $0 STUDENT_MODEL_PATH DATA_PATH"} +export GSM8K_TEACHER_MODEL_PATH=${GSM8K_TEACHER_MODEL_PATH:?"GSM8K_TEACHER_MODEL_PATH is required"} +export GEO3K_TEACHER_MODEL_PATH=${GEO3K_TEACHER_MODEL_PATH:?"GEO3K_TEACHER_MODEL_PATH is required"} + +export OPD_CONFIG_PATH="${OPD_CONFIG_PATH:-recipe/on_policy_distillation/config/rl_dapo_math_mopd.py}" +export EVAL_DATA_PATH="${EVAL_DATA_PATH:-}" +export TEACHER_STARTUP_TIMEOUT_S="${TEACHER_STARTUP_TIMEOUT_S:-1200}" +export PYTHONUNBUFFERED=1 + +USE_SGLANG=${XTUNER_USE_SGLANG:-0} +USE_LMDEPLOY=${XTUNER_USE_LMDEPLOY:-0} +USE_VLLM=${XTUNER_USE_VLLM:-0} + +if [[ "${USE_SGLANG}" == "1" && "${USE_LMDEPLOY}" == "0" && "${USE_VLLM}" == "0" ]]; then + OPD_BACKEND="sglang" +elif [[ "${USE_SGLANG}" == "0" && "${USE_LMDEPLOY}" == "1" && "${USE_VLLM}" == "0" ]]; then + OPD_BACKEND="lmdeploy" +else + echo "Exactly one of XTUNER_USE_SGLANG and XTUNER_USE_LMDEPLOY must be set to 1; XTUNER_USE_VLLM must be 0." >&2 + exit 1 +fi + +export XTUNER_USE_SGLANG="${USE_SGLANG}" +export XTUNER_USE_LMDEPLOY="${USE_LMDEPLOY}" +export XTUNER_USE_VLLM="${USE_VLLM}" + +export WORK_DIR="${WORK_DIR:-${REPO_ROOT}/work_dirs/dapo_math_mopd}" +export OPD_CONFIG_FILE="${REPO_ROOT}/${OPD_CONFIG_PATH}" + +TRAINING_STARTED=0 + +cleanup() { + local exit_code=$? + + trap - EXIT INT TERM + + stop_teacher_servers + + if (( TRAINING_STARTED )); then + ray stop --force >/dev/null 2>&1 || true + fi + + exit "${exit_code}" +} + +trap cleanup EXIT +trap "exit 130" INT +trap "exit 143" TERM + +start_teacher_servers "${OPD_CONFIG_FILE}" "${OPD_BACKEND}" "${WORK_DIR}" +wait_for_teacher_servers "${TEACHER_STARTUP_TIMEOUT_S}" + +echo "All ${#TEACHER_NAMES[@]} teachers are ready." +echo "Starting MOPD training with student GPUs: ${STUDENT_CUDA_VISIBLE_DEVICES}" + +cd "${REPO_ROOT}" +TRAINING_STARTED=1 +CUDA_VISIBLE_DEVICES="${STUDENT_CUDA_VISIBLE_DEVICES}" \ + bash -o pipefail examples/v1/scripts/run_rl.sh \ + "${OPD_CONFIG_FILE}" \ + "${OPD_BACKEND}" \ + "${STUDENT_MODEL_PATH}" \ + "${DATA_PATH}" \ + "${EVAL_DATA_PATH}" diff --git a/recipe/on_policy_distillation/scripts/run_sampled_token_opd.sh b/recipe/on_policy_distillation/scripts/run_sampled_token_opd.sh new file mode 100755 index 0000000000..212be2a67d --- /dev/null +++ b/recipe/on_policy_distillation/scripts/run_sampled_token_opd.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash + +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd -- "${SCRIPT_DIR}/../../.." && pwd) +source "${SCRIPT_DIR}/launch_teacher_utils.sh" + +export STUDENT_MODEL_PATH=${1:?"Usage: $0 STUDENT_MODEL_PATH TEACHER_MODEL_PATH DATA_PATH"} +export MODEL_PATH="${STUDENT_MODEL_PATH}" +export TEACHER_MODEL_PATH=${2:?"Usage: $0 STUDENT_MODEL_PATH TEACHER_MODEL_PATH DATA_PATH"} +export DATA_PATH=${3:?"Usage: $0 STUDENT_MODEL_PATH TEACHER_MODEL_PATH DATA_PATH"} +export OPD_CONFIG_PATH="${OPD_CONFIG_PATH:-recipe/on_policy_distillation/config/rl_dapo_math_opd.py}" +export EVAL_DATA_PATH="${EVAL_DATA_PATH:-}" +export TEACHER_STARTUP_TIMEOUT_S="${TEACHER_STARTUP_TIMEOUT_S:-1200}" +export PYTHONUNBUFFERED=1 + +USE_SGLANG=${XTUNER_USE_SGLANG:-0} +USE_LMDEPLOY=${XTUNER_USE_LMDEPLOY:-0} +USE_VLLM=${XTUNER_USE_VLLM:-0} + +if [[ "${USE_SGLANG}" == "1" && "${USE_LMDEPLOY}" == "0" && "${USE_VLLM}" == "0" ]]; then + OPD_BACKEND="sglang" +elif [[ "${USE_SGLANG}" == "0" && "${USE_LMDEPLOY}" == "1" && "${USE_VLLM}" == "0" ]]; then + OPD_BACKEND="lmdeploy" +else + echo "Exactly one of XTUNER_USE_SGLANG and XTUNER_USE_LMDEPLOY must be set to 1; XTUNER_USE_VLLM must be 0." >&2 + exit 1 +fi + +export XTUNER_USE_SGLANG="${USE_SGLANG}" +export XTUNER_USE_LMDEPLOY="${USE_LMDEPLOY}" +export XTUNER_USE_VLLM="${USE_VLLM}" + +export WORK_DIR="${WORK_DIR:-${REPO_ROOT}/work_dirs/dapo_math_opd}" +export OPD_CONFIG_FILE="${REPO_ROOT}/${OPD_CONFIG_PATH}" + +TRAINING_STARTED=0 + +cleanup() { + local exit_code=$? + + trap - EXIT INT TERM + + stop_teacher_servers + + if (( TRAINING_STARTED )); then + ray stop --force >/dev/null 2>&1 || true + fi + + exit "${exit_code}" +} + +trap cleanup EXIT +trap "exit 130" INT +trap "exit 143" TERM + +start_teacher_servers "${OPD_CONFIG_FILE}" "${OPD_BACKEND}" "${WORK_DIR}" +wait_for_teacher_servers "${TEACHER_STARTUP_TIMEOUT_S}" +curl -sS --max-time 10 "${TEACHER_MODEL_INFO_URLS[0]}" +echo +echo "Starting Pure PG-OPD training with student GPUs: ${STUDENT_CUDA_VISIBLE_DEVICES}" + +cd "${REPO_ROOT}" +TRAINING_STARTED=1 +CUDA_VISIBLE_DEVICES="${STUDENT_CUDA_VISIBLE_DEVICES}" \ + bash -o pipefail examples/v1/scripts/run_rl.sh \ + "${OPD_CONFIG_FILE}" \ + "${OPD_BACKEND}" \ + "${STUDENT_MODEL_PATH}" \ + "${DATA_PATH}" \ + "${EVAL_DATA_PATH}" diff --git a/tests/rl/test_on_policy_distillation.py b/tests/rl/test_on_policy_distillation.py new file mode 100644 index 0000000000..fa627bbee4 --- /dev/null +++ b/tests/rl/test_on_policy_distillation.py @@ -0,0 +1,580 @@ +import math +import os +import signal +import subprocess +import time +import unittest +from pathlib import Path +from unittest.mock import patch + +import httpx +import torch + +from recipe.on_policy_distillation.build_teacher_server_commands import ( + build_teacher_server_command, +) +from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.rl.loss import GRPOLossConfig +from xtuner.v1.rl.on_policy_distillation import ( + OPDConfig, + OPDTeacherConfig, + OPDTeacherLaunchConfig, + TeacherLogprobClient, + apply_opd_kl_to_advantages, +) +from xtuner.v1.rl.utils import find_free_ports + + +REPO_ROOT = Path(__file__).resolve().parents[2] +BASELINE_PATH = os.getenv("XTUNER_OPD_BASELINE") +TEACHER_MODEL_PATH = os.getenv("XTUNER_OPD_TEACHER_MODEL") +STUDENT_MODEL_PATH = os.getenv("XTUNER_OPD_STUDENT_MODEL") +TEACHER_STARTUP_TIMEOUT_S = float(os.getenv("XTUNER_OPD_TEACHER_STARTUP_TIMEOUT_S", "1200")) +TRAINER_CONFIG_PATH = ( + REPO_ROOT / "recipe/on_policy_distillation/config/rl_dapo_math_opd.py" +) + + +def _wait_for_teacher(process: subprocess.Popen, endpoint: str, backend: str) -> None: + deadline = time.monotonic() + TEACHER_STARTUP_TIMEOUT_S + health_path = "health_generate" if backend == "sglang" else "health" + with httpx.Client(timeout=1.0, trust_env=False) as client: + while time.monotonic() < deadline: + return_code = process.poll() + if return_code is not None: + raise RuntimeError(f"Teacher process exited during startup with code {return_code}") + try: + if client.get(f"{endpoint}/{health_path}").status_code == 200: + return + except httpx.RequestError: + pass + time.sleep(1.0) + raise TimeoutError(f"Teacher did not become ready within {TEACHER_STARTUP_TIMEOUT_S} seconds") + + +def _build_debug_rollout_batch(samples: list[dict]) -> list[list[RolloutState]]: + train_batch = [] + for sample in samples: + sample_index = int(sample["sample_index"]) + group_index = sample.get("group_index") + prompt_ids = torch.as_tensor(sample["prompt_token_ids"], dtype=torch.long).tolist() + response_ids = torch.as_tensor(sample["response_token_ids"], dtype=torch.long).tolist() + rollout_logprobs = torch.as_tensor(sample["rollout_log_probs"], dtype=torch.float32).tolist() + teacher_logprobs = torch.as_tensor(sample["teacher_log_probs"], dtype=torch.float32).tolist() + response_mask = torch.as_tensor(sample["loss_mask"], dtype=torch.bool).int().tolist() + response = str(sample["response"]) + + train_batch.append( + [ + RolloutState( + rollout_id=sample_index, + group_id=sample_index if group_index is None else int(group_index), + message=[], + prompt_ids=prompt_ids, + tokens=prompt_ids, + response=response, + response_ids=response_ids, + logprobs=rollout_logprobs, + teacher_tokens=response_ids, + teacher_logprobs=teacher_logprobs, + response_mask=response_mask, + reward={"score": 0.0}, + finish_reason="stop", + status=Status.COMPLETED, + extra_fields={"origin_data_source": "baseline"}, + ) + ] + ) + return train_batch + + +@unittest.skipUnless(BASELINE_PATH, "XTUNER_OPD_BASELINE is required") +class TestPGOPDAdvantage(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + baseline = torch.load(BASELINE_PATH, map_location="cpu", weights_only=False) + cls.samples = baseline["samples"] + + def test_compute_advantages_matches_baseline(self) -> None: + config = OPDConfig( + teachers=[OPDTeacherConfig(name="teacher", endpoint="http://unused")], + data_source_teacher_map={"baseline": "teacher"}, + ) + loss_cfg = GRPOLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + + for sample in self.samples: + with self.subTest(sample_index=sample["sample_index"]): + old_logprobs = torch.as_tensor(sample["old_log_probs"], dtype=torch.float32) + teacher_logprobs = torch.as_tensor(sample["teacher_log_probs"], dtype=torch.float32) + loss_mask = torch.as_tensor(sample["loss_mask"], dtype=torch.float32) + shifted_labels = torch.where( + loss_mask.bool(), + torch.zeros_like(loss_mask, dtype=torch.long), + torch.full_like(loss_mask, -100, dtype=torch.long), + ) + loss_ctx = loss_cfg.build( + { + "shifted_labels": shifted_labels, + "advantages": torch.zeros_like(old_logprobs), + "old_logprobs": old_logprobs, + "teacher_logprobs": teacher_logprobs, + } + ) + assert loss_ctx is not None + apply_opd_kl_to_advantages(loss_ctx, config=config) + actual_advantages = loss_ctx.loss_kwargs.advantages.cpu() + expected_advantages = torch.as_tensor(sample["advantages"], dtype=torch.float32) * loss_mask + try: + torch.testing.assert_close( + actual_advantages, + expected_advantages, + rtol=1e-5, + atol=1e-5, + ) + except AssertionError as error: + mismatch_indices = torch.nonzero( + ~torch.isclose(actual_advantages, expected_advantages, rtol=1e-5, atol=1e-5) + ).flatten() + mismatch_values = "\n".join( + ( + f"index={index}: " + f"actual={actual_advantages[index].item()!r}, " + f"expected={expected_advantages[index].item()!r}, " + f"abs_diff={abs(actual_advantages[index] - expected_advantages[index]).item()!r}" + ) + for index in mismatch_indices.tolist() + ) + raise AssertionError(f"{error}\n\nMismatched advantages:\n{mismatch_values}") from None + + +def _load_trainer_samples(capture_dir: Path) -> dict[int, dict]: + trainer_samples = {} + for capture_file in sorted(capture_dir.glob("rank_*.pt")): + for batch in torch.load(capture_file, map_location="cpu", weights_only=True): + shifted_labels = torch.as_tensor(batch["shifted_labels"], dtype=torch.long).reshape(-1) + old_log_probs = torch.as_tensor(batch["old_log_probs"], dtype=torch.float32).reshape(-1) + advantages = torch.as_tensor(batch["advantages"], dtype=torch.float32).reshape(-1) + boundaries = torch.as_tensor(batch["cu_seq_lens_q"], dtype=torch.long).tolist() + num_padding = int(batch["num_padding"]) + padding_start = shifted_labels.numel() - num_padding + + assert boundaries[len(batch["rollout_ids"])] == padding_start + torch.testing.assert_close( + shifted_labels[padding_start:], + torch.full_like(shifted_labels[padding_start:], -100), + rtol=0, + atol=0, + ) + + for rollout_id, start, end in zip(batch["rollout_ids"], boundaries[:-1], boundaries[1:]): + response_mask = shifted_labels[start:end] != -100 + trainer_samples[int(rollout_id)] = { + "response_token_ids": shifted_labels[start:end][response_mask], + "old_log_probs": old_log_probs[start:end][response_mask], + "advantages": advantages[start:end][response_mask], + } + + return trainer_samples + + +def _run_trainer_once( + baseline_samples: list[dict], + *, + baseline_path: Path, + student_model_path: Path, + debug_rollout_dir: Path, + capture_dir: Path, + run_dir: Path, +) -> dict[int, dict]: + import ray + + from xtuner.v1.rl.trainer.controller import TrainingController + from xtuner.v1.rl.trainer.worker import TrainingWorker, WorkerConfig + from xtuner.v1.train.rl_trainer import RLColocateTrainer + from xtuner.v1.utils import Config + + max_prompt_length = max(len(sample["prompt_token_ids"]) for sample in baseline_samples) + max_response_length = max(len(sample["response_token_ids"]) for sample in baseline_samples) + max_input_length = max(len(sample["token_ids"]) - 1 for sample in baseline_samples) + pack_max_length = math.ceil(max_input_length / 512) * 512 + num_workers = 8 + + class CapturingRLColocateTrainer(RLColocateTrainer): + def _prepare_train_data( + self, + data_groups, + pack_max_length, + raw_rewards_sum=0.0, + raw_rewards_count=0, + ): + data_batches, data_info = super()._prepare_train_data( + data_groups, + pack_max_length, + raw_rewards_sum=raw_rewards_sum, + raw_rewards_count=raw_rewards_count, + ) + rollout_states = (state for group in data_groups for state in group) + for data_batch, rollout_state in zip(data_batches, rollout_states): + data_batch["rollout_id"] = rollout_state.rollout_id + return data_batches, data_info + + class CapturingTrainingController(TrainingController): + def _packing(self, data_batches, pack_max_length, language_cfg): + pack_infos = self._get_pack_infos( + data_batches, + [data["seq_ctx"].input_ids.numel() for data in data_batches], + pack_max_length, + ) + packed_data_batches = super()._packing(data_batches, pack_max_length, language_cfg) + for packed_data, pack_info in zip(packed_data_batches, pack_infos): + packed_data["rollout_ids"] = [data_batches[index]["rollout_id"] for index in pack_info["indices"]] + return packed_data_batches + + class CapturingTrainingWorker(TrainingWorker): + trainer_capture_dir = capture_dir + + def fit(self, data_batches, rollout_idx): + from unittest.mock import patch as mock_patch + + from xtuner.v1.rl.trainer import worker as worker_module + + captured_batches = [ + { + "rollout_ids": data.get("rollout_ids", []), + "cu_seq_lens_q": data["seq_ctx"].cu_seq_lens_q.detach().cpu(), + "num_padding": data["seq_ctx"].num_padding, + "shifted_labels": data["shifted_labels"].detach().cpu().reshape(-1), + } + for data in data_batches + ] + captured_batch_index = 0 + apply_opd = worker_module.apply_opd_kl_to_advantages + + def capture_opd_result(loss_ctx, *, config): + nonlocal captured_batch_index + reverse_kl_sum = apply_opd(loss_ctx, config=config) + captured_batches[captured_batch_index]["old_log_probs"] = ( + loss_ctx.loss_kwargs.old_logprobs.detach().cpu().reshape(-1) + ) + captured_batches[captured_batch_index]["advantages"] = ( + loss_ctx.loss_kwargs.advantages.detach().cpu().reshape(-1) + ) + captured_batch_index += 1 + return reverse_kl_sum + + with mock_patch.object(worker_module, "apply_opd_kl_to_advantages", capture_opd_result): + worker_log_item = TrainingWorker.fit(self, data_batches, rollout_idx) + + torch.save(captured_batches, self.trainer_capture_dir / f"rank_{self.rank}.pt") + return worker_log_item + + def build_capturing_training_workers(self, placement_group): + from xtuner.v1.rl.utils import AutoAcceleratorWorkers + + capturing_worker_cls = ray.remote( + runtime_env={ + "env_vars": { + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1", + "HCCL_NPU_SOCKET_PORT_RANGE": "auto", + } + } + )(CapturingTrainingWorker) + train_workers, _ = AutoAcceleratorWorkers.from_placement_group( + capturing_worker_cls, + self, + placement_group, + ) + ray.wait([worker.ready.remote() for worker in train_workers]) + return CapturingTrainingController(workers=train_workers) + + trainer_environment = { + "WORK_DIR": str(run_dir / "trainer"), + "MODEL_PATH": str(student_model_path), + "TEACHER_MODEL_PATH": os.getenv("XTUNER_OPD_TEACHER_MODEL", str(student_model_path)), + "DATA_PATH": os.getenv("XTUNER_OPD_DATA_PATH", str(baseline_path)), + "WORLD_SIZE": "1", + "ONLY_CALC_MISMATCH_RATIO": "1", + "XTUNER_DETERMINISTIC": "true", + "XTUNER_USE_FA3": os.getenv("XTUNER_OPD_USE_FA3", "1"), + "TOKENIZERS_PARALLELISM": "false", + "PYTHONUNBUFFERED": "1", + } + + with patch.dict(os.environ, trainer_environment, clear=False): + try: + ray.init( + num_cpus=12 * num_workers + 10, + num_gpus=num_workers, + include_dashboard=False, + _temp_dir="/dev/shm/xtuner-opd-old-logprobs", + ) + cfg = Config.fromfile(TRAINER_CONFIG_PATH) + cfg.trainer.resources.num_workers = num_workers + cfg.trainer.total_train_steps = 1 + cfg.trainer.train_batch_size = len(baseline_samples) + cfg.trainer.train_worker_cfg.optimizer_steps = 1 + cfg.trainer.train_worker_cfg.pack_max_length = pack_max_length + cfg.trainer.debug_train = True + cfg.trainer.debug_rollout_dir = debug_rollout_dir + with ( + patch("xtuner.v1.train.rl_trainer.XTUNER_DETERMINISTIC", True), + patch( + "xtuner.v1.train.rl_trainer.RLColocateTrainer", + CapturingRLColocateTrainer, + ), + patch.object(WorkerConfig, "build", build_capturing_training_workers), + ): + trainer = cfg.trainer.build() + trainer.fit() + finally: + if ray.is_initialized(): + ray.shutdown() + + return _load_trainer_samples(capture_dir) + + +@unittest.skipUnless( + BASELINE_PATH and STUDENT_MODEL_PATH, + "XTUNER_OPD_BASELINE and XTUNER_OPD_STUDENT_MODEL are required", +) +class TestPGOPDOldLogprobs(unittest.TestCase): + def test_trainer_old_logprobs_and_advantage_error_propagation(self) -> None: + baseline_path = Path(str(BASELINE_PATH)).expanduser().resolve() + student_model_path = Path(str(STUDENT_MODEL_PATH)).expanduser().resolve() + baseline = torch.load(baseline_path, map_location="cpu", weights_only=True) + baseline_samples = baseline["samples"] + + work_root = Path( + os.getenv( + "XTUNER_OPD_OLD_LOGPROB_WORK_DIR", + str(REPO_ROOT / "work_dirs/test_pg_opd_old_logprobs"), + ) + ).expanduser() + run_dir = work_root / f"run_{time.strftime('%Y%m%d_%H%M%S')}_{os.getpid()}" + debug_rollout_dir = run_dir / "debug_rollout" + capture_dir = run_dir / "trainer_capture" + debug_rollout_dir.mkdir(parents=True) + capture_dir.mkdir() + + train_batch = _build_debug_rollout_batch(baseline_samples) + torch.save(train_batch, debug_rollout_dir / "debug_rollout_1.pt") + trainer_samples_by_rollout_id = _run_trainer_once( + baseline_samples, + baseline_path=baseline_path, + student_model_path=student_model_path, + debug_rollout_dir=debug_rollout_dir, + capture_dir=capture_dir, + run_dir=run_dir, + ) + + result_path = run_dir / "result.pt" + result_samples = [] + + for baseline_sample in baseline_samples: + rollout_id = int(baseline_sample["sample_index"]) + trainer_sample = trainer_samples_by_rollout_id[rollout_id] + loss_mask = torch.as_tensor(baseline_sample["loss_mask"], dtype=torch.bool) + response_token_ids = torch.as_tensor(baseline_sample["response_token_ids"], dtype=torch.long)[loss_mask] + torch.testing.assert_close( + trainer_sample["response_token_ids"], + response_token_ids, + rtol=0, + atol=0, + ) + baseline_old_log_probs = torch.as_tensor( + baseline_sample["old_log_probs"], + dtype=torch.float32, + )[loss_mask] + trainer_old_log_probs = trainer_sample["old_log_probs"] + baseline_advantages = torch.as_tensor( + baseline_sample["advantages"], + dtype=torch.float32, + )[loss_mask] + trainer_advantages = trainer_sample["advantages"] + + old_logprob_error = trainer_old_log_probs - baseline_old_log_probs + advantage_error = trainer_advantages - baseline_advantages + propagation_residual = advantage_error + old_logprob_error + sample_num_tokens = old_logprob_error.numel() + sample_summary = { + "num_tokens": sample_num_tokens, + "old_logprobs_mean_abs_error": old_logprob_error.abs().mean().item(), + "old_logprobs_max_abs_error": old_logprob_error.abs().max().item(), + "advantages_mean_abs_error": advantage_error.abs().mean().item(), + "advantages_max_abs_error": advantage_error.abs().max().item(), + "propagation_mean_abs_error": propagation_residual.abs().mean().item(), + "propagation_max_abs_error": propagation_residual.abs().max().item(), + } + result_samples.append( + { + "sample_index": rollout_id, + "response_token_ids": response_token_ids, + "baseline_old_log_probs": baseline_old_log_probs, + "trainer_old_log_probs": trainer_old_log_probs, + "old_logprob_error": old_logprob_error, + "baseline_advantages": baseline_advantages, + "trainer_advantages": trainer_advantages, + "advantage_error": advantage_error, + "propagation_residual": propagation_residual, + "summary": sample_summary, + } + ) + + with self.subTest(sample_index=rollout_id): + torch.testing.assert_close( + advantage_error, + -old_logprob_error, + rtol=1e-5, + atol=1e-5, + msg=f"Result: {result_path}", + ) + + global_result = { + "response_token_ids": torch.cat([sample["response_token_ids"] for sample in result_samples]), + "baseline_old_log_probs": torch.cat([sample["baseline_old_log_probs"] for sample in result_samples]), + "trainer_old_log_probs": torch.cat([sample["trainer_old_log_probs"] for sample in result_samples]), + "old_logprob_error": torch.cat([sample["old_logprob_error"] for sample in result_samples]), + "baseline_advantages": torch.cat([sample["baseline_advantages"] for sample in result_samples]), + "trainer_advantages": torch.cat([sample["trainer_advantages"] for sample in result_samples]), + "advantage_error": torch.cat([sample["advantage_error"] for sample in result_samples]), + "propagation_residual": torch.cat([sample["propagation_residual"] for sample in result_samples]), + } + summary = { + "num_samples": len(result_samples), + "num_tokens": global_result["old_logprob_error"].numel(), + "old_logprobs_mean_abs_error": global_result["old_logprob_error"].abs().mean().item(), + "old_logprobs_max_abs_error": global_result["old_logprob_error"].abs().max().item(), + "advantages_mean_abs_error": global_result["advantage_error"].abs().mean().item(), + "advantages_max_abs_error": global_result["advantage_error"].abs().max().item(), + "propagation_mean_abs_error": global_result["propagation_residual"].abs().mean().item(), + "propagation_max_abs_error": global_result["propagation_residual"].abs().max().item(), + } + torch.save( + { + "summary": summary, + "global": global_result, + "samples": result_samples, + }, + result_path, + ) + + with self.subTest(scope="global"): + torch.testing.assert_close( + global_result["advantage_error"], + -global_result["old_logprob_error"], + rtol=1e-5, + atol=1e-5, + msg=f"Result: {result_path}", + ) + + print( + f"old_logprobs: mean_abs={summary['old_logprobs_mean_abs_error']}, " + f"max_abs={summary['old_logprobs_max_abs_error']}\n" + f"advantages: mean_abs={summary['advantages_mean_abs_error']}, " + f"max_abs={summary['advantages_max_abs_error']}\n" + f"propagation: mean_abs={summary['propagation_mean_abs_error']}, " + f"max_abs={summary['propagation_max_abs_error']}\n" + f"result: {result_path}" + ) + + +@unittest.skipUnless( + BASELINE_PATH and TEACHER_MODEL_PATH, + "XTUNER_OPD_BASELINE and XTUNER_OPD_TEACHER_MODEL are required", +) +class TestTeacherLogprobClient(unittest.IsolatedAsyncioTestCase): + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + baseline = torch.load(BASELINE_PATH, map_location="cpu", weights_only=False) + cls.samples = baseline["samples"] + port = find_free_ports()[0] + cls.teacher_endpoint = f"http://127.0.0.1:{port}" + cls.teacher_backend = TeacherLogprobClient._resolve_backend_from_env() + teacher_env = os.environ.copy() + teacher_command = build_teacher_server_command( + OPDTeacherConfig( + name="teacher", + endpoint=cls.teacher_endpoint, + launch_config=OPDTeacherLaunchConfig( + model_path=str(TEACHER_MODEL_PATH), + cuda_visible_devices=teacher_env.get("CUDA_VISIBLE_DEVICES") + or "7", + ), + ), + cls.teacher_backend, + ) + if not teacher_command: + raise ValueError("Teacher must define launch_config for local startup") + cls.teacher_process = subprocess.Popen( + teacher_command, + cwd=REPO_ROOT, + env=teacher_env, + start_new_session=True, + ) + cls.addClassCleanup(cls._stop_teacher) + _wait_for_teacher(cls.teacher_process, cls.teacher_endpoint, cls.teacher_backend) + + @classmethod + def _stop_teacher(cls) -> None: + if cls.teacher_process.poll() is not None: + return + os.killpg(cls.teacher_process.pid, signal.SIGTERM) + try: + cls.teacher_process.wait(timeout=30) + except subprocess.TimeoutExpired: + os.killpg(cls.teacher_process.pid, signal.SIGKILL) + cls.teacher_process.wait() + + @unittest.skipUnless(os.getenv("XTUNER_USE_SGLANG", "0") == "1", "XTUNER_USE_SGLANG=1 is required") + async def test_compute_logprobs_with_sglang_matches_baseline(self) -> None: + self.assertEqual(self.teacher_backend, "sglang") + await self._assert_compute_logprobs_matches_baseline() + + @unittest.skipUnless(os.getenv("XTUNER_USE_LMDEPLOY", "0") == "1", "XTUNER_USE_LMDEPLOY=1 is required") + async def test_compute_logprobs_with_lmdeploy_matches_baseline(self) -> None: + self.assertEqual(self.teacher_backend, "lmdeploy") + await self._assert_compute_logprobs_matches_baseline() + + async def _assert_compute_logprobs_matches_baseline(self) -> None: + client = TeacherLogprobClient(OPDTeacherConfig(name="teacher", endpoint=self.teacher_endpoint)) + self.addAsyncCleanup(client._client.aclose) + + for sample in self.samples: + with self.subTest(sample_index=sample["sample_index"]): + prompt_ids = torch.as_tensor(sample["prompt_token_ids"], dtype=torch.long).tolist() + response_ids = torch.as_tensor(sample["response_token_ids"], dtype=torch.long).tolist() + expected_logprobs = torch.as_tensor(sample["teacher_log_probs"], dtype=torch.float32) + state = RolloutState( + rollout_id=int(sample["sample_index"]), + group_id=int(sample["group_index"]), + message=[], + prompt_ids=prompt_ids, + tokens=prompt_ids, + response="", + response_ids=response_ids, + status=Status.COMPLETED, + ) + + result = await client.compute_logprobs(state) + + self.assertEqual(result.status, Status.COMPLETED, result.error_msg) + self.assertEqual(result.teacher_tokens, response_ids) + actual_logprobs = torch.tensor(result.teacher_logprobs, dtype=torch.float32) + self.assertEqual(actual_logprobs.shape, expected_logprobs.shape) + absolute_errors = torch.abs(actual_logprobs - expected_logprobs) + mae = absolute_errors.mean().item() + self.assertLessEqual( + mae, + 0.05, + ( + f"Teacher logprob MAE {mae:.8f} exceeds threshold 0.05; " + f"max_abs_error={absolute_errors.max().item():.8f}" + ), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/rl/test_prepare_train_data.py b/tests/rl/test_prepare_train_data.py index 3c178b4769..0a209df5bb 100644 --- a/tests/rl/test_prepare_train_data.py +++ b/tests/rl/test_prepare_train_data.py @@ -7,9 +7,6 @@ - VLM 样本使用 train_prompt_ids，并保留 multimodal 训练字段。 - 无效 rollout group 会被跳过。 - 缺失 reward、logprob/mask 长度不一致、pack_max_length 过小时 fail fast。 - -注意：当前训练 contract 中 data_dict["advantage"] 比 shifted_labels 多 1 个元素； -metric 统计使用 actual_advantages[:-1]，测试会显式固定这个行为。 """ import unittest @@ -36,6 +33,7 @@ class TestPrepareTrainData(unittest.TestCase): def _build_trainer(self, advantages: list[float]): trainer = BaseRLTrainer.__new__(BaseRLTrainer) trainer._advantage_estimator = _FakeAdvantageEstimator(advantages) + trainer._opd_config = None trainer.tokenizer = MagicMock(return_value={"input_ids": torch.tensor([[999]])}) trainer.logger = MagicMock() return trainer @@ -102,8 +100,8 @@ def test_text_path_builds_shifted_training_tensors(self): batch["rollout_logprobs"], torch.tensor([[0.0, 0.0, 0.1, 0.2, 0.3]], dtype=torch.float32), ) - self.assertEqual(batch["advantage"], [1.5, 1.5, 1.5, 1.5, 0.0, 1.5]) - self.assertEqual(len(batch["advantage"]), batch["shifted_labels"].numel() + 1) + self.assertEqual(batch["advantage"], [0.0, 0.0, 1.5, 0.0, 1.5]) + self.assertEqual(len(batch["advantage"]), batch["shifted_labels"].numel()) self.assertIs(batch["seq_ctx"].rollout_routed_experts, routed_experts) self.assertEqual(info["training_samples"], 1) self.assertEqual(info["training_tokens"], 5) @@ -136,8 +134,8 @@ def test_multi_sample_group_uses_each_sample_reward_and_advantage(self): data_batches, info = self._prepare(trainer, [[first, second]]) self.assertEqual(len(data_batches), 2) - self.assertEqual(data_batches[0]["advantage"], [1.5, 1.5, 1.5, 1.5, 1.5]) - self.assertEqual(data_batches[1]["advantage"], [-2.0, -2.0, -2.0, -2.0, -2.0]) + self.assertEqual(data_batches[0]["advantage"], [0.0, 0.0, 1.5, 1.5]) + self.assertEqual(data_batches[1]["advantage"], [0.0, 0.0, -2.0, -2.0]) self.assertEqual(info["batch_size"], 2) self.assertEqual(info["rewards/min"], -1.0) self.assertEqual(info["rewards/max"], 3.0) diff --git a/tests/rl/test_producer.py b/tests/rl/test_producer.py index 364db9f09c..44f8e67c67 100644 --- a/tests/rl/test_producer.py +++ b/tests/rl/test_producer.py @@ -119,6 +119,8 @@ async def mock_gen(rs, **kwargs): mock_agent_loop.generate_group = mock_gen mock_agent_loop.is_valid_sample_fn = is_valid_sample_fn + mock_agent_loop.teacher_clients = {} + mock_agent_loop.collect_rollout_group = AgentLoop.collect_rollout_group.__get__(mock_agent_loop) return mock_agent_loop def _build_context( @@ -132,6 +134,7 @@ def _build_context( train_step: int = 0, model_step: int = 0, progress: ProduceProgress | None = None, + is_valid_sample_fn=None, ) -> ProduceContext: # 测试只走新的 ProduceContext 入口，不再覆盖旧散装参数兼容逻辑。 if progress is None: @@ -349,6 +352,7 @@ def is_valid_sample_fn(samples): self._build_agent_loop(is_valid_sample_fn=is_valid_sample_fn), self._build_sampler(), batch_size=1, + is_valid_sample_fn=is_valid_sample_fn, ) completed_group = [make_rollout_state(1, status=Status.COMPLETED)] @@ -433,6 +437,7 @@ def is_valid_sample_fn(samples): self._build_agent_loop(is_valid_sample_fn=is_valid_sample_fn), self._build_sampler(), batch_size=1, + is_valid_sample_fn=is_valid_sample_fn, ) completed_group = [ @@ -524,6 +529,7 @@ async def mock_gen(rs, **kwargs): train_step=4, model_step=3, progress=self._build_progress(task_name, target=2), + is_valid_sample_fn=is_valid_sample_fn, ) await strategy.produce_batch(ctx) @@ -554,7 +560,7 @@ async def mock_gen(rs, **kwargs): r.status = Status.COMPLETED return rs - mock_agent_loop.generate_group = mock_gen + mock_agent_loop.collect_rollout_group = mock_gen sampler_cfg = SamplerConfig.model_construct(dataloader_cfg=self.mock_dataloader_cfg) produce_strategy_cfg = AsyncProduceStrategyConfig(over_sample_threshold=1) diff --git a/tests/rl/test_rl_colocate_trainer.py b/tests/rl/test_rl_colocate_trainer.py index b5ef8e96e8..1b6b3481ed 100644 --- a/tests/rl/test_rl_colocate_trainer.py +++ b/tests/rl/test_rl_colocate_trainer.py @@ -100,7 +100,7 @@ async def generate_group(rollout_states, **kwargs): state.response_model_steps = [model_step] return rollout_states - agent_loop.generate_group = generate_group + agent_loop.collect_rollout_group = generate_group return agent_loop diff --git a/xtuner/v1/data_proto/rl_data.py b/xtuner/v1/data_proto/rl_data.py index ec00526144..95e3f23831 100644 --- a/xtuner/v1/data_proto/rl_data.py +++ b/xtuner/v1/data_proto/rl_data.py @@ -109,6 +109,8 @@ class RolloutState(BaseModel): tool_calls: list[RolloutToolCall] | None = None response_ids: list[int] | None = None logprobs: list[float] | None = None + teacher_tokens: list[int] | None = None + teacher_logprobs: list[float] | None = None routed_experts: np.ndarray | RayObjectRef | list[RayObjectRef] | None = None finish_reason: str | None = None # response_mask: 记录response_ids中哪个token算loss, 与response_ids长度相同，每轮rollout在 agent_loop.generate 中覆盖写 @@ -244,6 +246,8 @@ def reset_rollout_response(rollout_state: RolloutState) -> RolloutState: rollout_state.response = "" rollout_state.response_ids = [] rollout_state.logprobs = [] + rollout_state.teacher_tokens = None + rollout_state.teacher_logprobs = None rollout_state.routed_experts = None rollout_state.finish_reason = None rollout_state.response_mask = None diff --git a/xtuner/v1/loss/__init__.py b/xtuner/v1/loss/__init__.py index d2f20b3a16..51c7596daf 100644 --- a/xtuner/v1/loss/__init__.py +++ b/xtuner/v1/loss/__init__.py @@ -11,7 +11,7 @@ ZLossKwargs, ) from .mtp_loss import MTPLossContext -from .rl_loss import LogProbConfig, LogProbContext +from .rl_loss import LogProbConfig, LogProbContext, TopKLogProbConfig, TopKLogProbContext __all__ = [ @@ -33,6 +33,8 @@ "MTPLossContext", "LogProbConfig", "LogProbContext", + "TopKLogProbConfig", + "TopKLogProbContext", ] import torch diff --git a/xtuner/v1/loss/rl_loss.py b/xtuner/v1/loss/rl_loss.py index b6d15291d1..fd22b2c629 100644 --- a/xtuner/v1/loss/rl_loss.py +++ b/xtuner/v1/loss/rl_loss.py @@ -2,6 +2,7 @@ import torch import torch.nn.functional as F +from pydantic import Field from torch.distributed.device_mesh import DeviceMesh from xtuner.v1.rl.utils.misc import gather_logprobs @@ -91,3 +92,75 @@ def forward( else: logprobs, _ = self.loss_fn(hidden_states, head_weight, head_bias, self.loss_kwargs) return logprobs, (None, {}) + + +class TopKLogProbConfig(BaseLossConfig): + """Select model Top-K tokens and compute their exact full-softmax log + probabilities.""" + + top_k: int = Field(gt=0) + + @property + def loss_ctx_cls(self) -> type["TopKLogProbContext"]: + return TopKLogProbContext + + @property + def _loss_kwargs_cls(self) -> type[BaseLossKwargs]: + return BaseLossKwargs + + def build(self, data: dict, sp_mesh: DeviceMesh | None = None) -> "TopKLogProbContext": + del data, sp_mesh + return self.loss_ctx_cls(self, self._loss_kwargs_cls()) + + +class TopKLogProbContext(BaseLossContext): + """Select model Top-K IDs without storing full log-softmax.""" + + loss_cfg: TopKLogProbConfig + loss_kwargs: BaseLossKwargs + + def loss_fn( + self, + hidden_states: torch.Tensor, + head_weight: torch.Tensor, + head_bias: torch.Tensor | None, + loss_kwargs: BaseLossKwargs, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, dict[str, Any]]]: + del loss_kwargs + logits = F.linear(hidden_states, head_weight, head_bias).float() + selected_logits, token_ids = torch.topk(logits, k=self.loss_cfg.top_k, dim=-1) + selected_logprobs = selected_logits - torch.logsumexp(logits, dim=-1, keepdim=True) + return selected_logprobs, (token_ids, {}) + + def chunk_mode( + self, + hidden_states: torch.Tensor, + head_weight: torch.Tensor, + head_bias: torch.Tensor | None, + loss_kwargs: BaseLossKwargs, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, dict[str, Any]]]: + assert self.loss_cfg.chunk_size is not None, "chunk_size must be set in chunk mode" + + logprob_chunks = [] + token_id_chunks = [] + for hidden_states_chunk in torch.split(hidden_states, self.loss_cfg.chunk_size, dim=1): + logprobs, (token_ids, _) = self.loss_fn( + hidden_states_chunk, + head_weight, + head_bias, + loss_kwargs, + ) + logprob_chunks.append(logprobs) + token_id_chunks.append(token_ids) + return torch.cat(logprob_chunks, dim=1), (torch.cat(token_id_chunks, dim=1), {}) + + def forward( + self, + hidden_states: torch.Tensor, + head_weight: torch.Tensor, + head_bias: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, dict[str, Any]]]: + if self.loss_cfg.mode == "chunk": + return self.chunk_mode(hidden_states, head_weight, head_bias, self.loss_kwargs) + else: + return self.loss_fn(hidden_states, head_weight, head_bias, self.loss_kwargs) diff --git a/xtuner/v1/model/utils/misc.py b/xtuner/v1/model/utils/misc.py index 70fbf2d2d2..67d970e5e6 100644 --- a/xtuner/v1/model/utils/misc.py +++ b/xtuner/v1/model/utils/misc.py @@ -108,6 +108,17 @@ def get(self): "reduced_train_policy_kl1_sum", "reduced_train_policy_kl3_sum", "reduced_train_policy_valid_count", + "reduced_distillation_kl_sum", + "reduced_distillation_abs_loss_sum", + "reduced_distillation_valid_count", + "reduced_opd_reverse_kl_sum", + "reduced_opd_abs_logprob_loss_sum", + "reduced_topk_opd_kl_sum", + "reduced_topk_opd_loss_sum", + "reduced_topk_opd_student_selected_mass_sum", + "reduced_topk_opd_teacher_selected_mass_sum", + "reduced_topk_opd_overlap_fraction_sum", + "reduced_topk_opd_valid_count", ) max_keys = ("max_ratio", "reduced_train_policy_ratio_max") min_keys = ("reduced_train_policy_ratio_min",) diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index 63b21b18f7..844b5ed87c 100644 --- a/xtuner/v1/rl/agent_loop/agent_loop.py +++ b/xtuner/v1/rl/agent_loop/agent_loop.py @@ -11,6 +11,12 @@ from ray.util.placement_group import PlacementGroup from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status +from xtuner.v1.rl.distillation import ( + DistillationConfig, + RolloutTeacherClient, + route_rollout_teacher_client, + validate_opd_sample_params, +) from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController from xtuner.v1.rl.rollout.constants import AGENT_LOOP_RAY_GENERATE_MAX_CONCURRENCY @@ -78,6 +84,7 @@ def build( logger=None, *, is_valid_sample_fn: IsValidSampleFn | None = None, + distillation_config: DistillationConfig | None = None, ) -> AgentLoopSpec: if self.cpu_resources is None: agent_loop = self.build_local( @@ -86,6 +93,7 @@ def build( logger=logger, ) agent_loop.is_valid_sample_fn = is_valid_sample_fn + agent_loop.configure_distillation(distillation_config) return agent_loop concurrency = AGENT_LOOP_RAY_GENERATE_MAX_CONCURRENCY @@ -103,6 +111,7 @@ def build( judger=judger, logger=logger, is_valid_sample_fn=is_valid_sample_fn, + distillation_config=distillation_config, ) return self._build_ray_actor( rollout_controller=rollout_controller, @@ -111,6 +120,7 @@ def build( judger=judger, logger=logger, is_valid_sample_fn=is_valid_sample_fn, + distillation_config=distillation_config, ) @abstractmethod @@ -130,6 +140,7 @@ def _build_ray_actor( judger: Judger | None = None, logger=None, is_valid_sample_fn: IsValidSampleFn | None = None, + distillation_config: DistillationConfig | None = None, ) -> RayAgentLoopProxy: ray_agent_loop = ray.remote( concurrency_groups={ @@ -149,6 +160,7 @@ def _build_ray_actor( actor_memory=cpu_resources.cpu_memory_per_worker, capture_child_tasks=True, is_valid_sample_fn=is_valid_sample_fn, + distillation_config=distillation_config, ), ) @@ -162,6 +174,7 @@ def _build_ray_actors( logger=None, start_bundle_idx: int = 0, is_valid_sample_fn: IsValidSampleFn | None = None, + distillation_config: DistillationConfig | None = None, ) -> list[RayAgentLoopProxy]: ray_agent_loop = ray.remote( concurrency_groups={ @@ -182,6 +195,7 @@ def _build_ray_actors( actor_memory_per_worker=cpu_resources.cpu_memory_per_worker, capture_child_tasks=True, is_valid_sample_fn=is_valid_sample_fn, + distillation_config=distillation_config, ), ) @@ -195,6 +209,7 @@ def _build_router( logger=None, start_bundle_idx: int = 0, is_valid_sample_fn: IsValidSampleFn | None = None, + distillation_config: DistillationConfig | None = None, ) -> RouterAgentLoop: return RouterAgentLoop( workers=self._build_ray_actors( @@ -206,6 +221,7 @@ def _build_router( logger=logger, start_bundle_idx=start_bundle_idx, is_valid_sample_fn=is_valid_sample_fn, + distillation_config=distillation_config, ), rollout_ctl=rollout_controller, ) @@ -234,6 +250,18 @@ def __init__( else: self.logger = logger self._judger_pause_event = asyncio.Event() + self.teacher_clients: dict[str, RolloutTeacherClient] = {} + self.data_source_teacher_map: dict[str, str] = {} + + def configure_distillation(self, distillation_config: DistillationConfig | None) -> None: + if distillation_config is None: + return + + validate_opd_sample_params(self.sample_params) + self.teacher_clients = { + teacher.name: RolloutTeacherClient(teacher) for teacher in distillation_config.rollout_teachers + } + self.data_source_teacher_map = dict(distillation_config.data_source_teacher_map) @abstractmethod async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: ... @@ -259,6 +287,48 @@ async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> l # Keep sample validation as the final group-generation step. return maybe_filter_invalid_sample(group_samples, self.is_valid_sample_fn, self.logger) + async def collect_rollout_group( + self, + rollout_state: list[RolloutState], + *, + is_valid_sample_func: Callable[[list[RolloutState]], bool] | None = None, + **kwargs, + ) -> list[RolloutState]: + if is_valid_sample_func is None and self.teacher_clients: + teacher = route_rollout_teacher_client( + rollout_state[0], + data_source_teacher_map=self.data_source_teacher_map, + teacher_clients=self.teacher_clients, + ) + + async def generate_and_score(state: RolloutState) -> RolloutState: + state.sample_params = self.sample_params + state = await self.generate_sample(state, **kwargs) + if state.status == Status.COMPLETED: + state = await teacher.compute_logprobs(state) + return state + + group = list(await asyncio.gather(*(create_task(generate_and_score(state)) for state in rollout_state))) + if self.judger is not None and self.enable_batch_judge and get_group_status(group) == Status.COMPLETED: + group = await self.run_judger(group) + return group + + group = await self.generate_group(rollout_state, **kwargs) + if get_group_status(group) != Status.COMPLETED: + return group + if is_valid_sample_func is not None and not is_valid_sample_func(group): + for state in group: + state.status = Status.FILTERED + return group + if self.teacher_clients: + teacher = route_rollout_teacher_client( + group[0], + data_source_teacher_map=self.data_source_teacher_map, + teacher_clients=self.teacher_clients, + ) + group = list(await asyncio.gather(*(create_task(teacher.compute_logprobs(state)) for state in group))) + return group + @overload async def run_judger(self, rollout_state: RolloutState) -> RolloutState: ... @@ -345,6 +415,13 @@ async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> l finally: await self._release_worker(worker) + async def collect_rollout_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: + worker = await self._pick_worker() + try: + return await worker.collect_rollout_group.remote(rollout_state, **kwargs) + finally: + await self._release_worker(worker) + def get_worker_status(self) -> dict[str, int]: return {str(worker): load for worker, load in self._worker_loads.items()} @@ -373,6 +450,8 @@ def __init__( judger: Judger | None = None, logger=None, is_valid_sample_fn: IsValidSampleFn | None = None, + *, + distillation_config: DistillationConfig | None = None, ): self.agent_loop = agent_loop_config.build_local( rollout_controller=rollout_controller, @@ -380,6 +459,7 @@ def __init__( logger=logger, ) self.agent_loop.is_valid_sample_fn = is_valid_sample_fn + self.agent_loop.configure_distillation(distillation_config) @ray_method(concurrency_group=AGENT_LOOP_CONCURRENCY_GROUP_GENERATE) async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: @@ -389,6 +469,10 @@ async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> Rollou async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: return await self.agent_loop.generate_group(rollout_state, **kwargs) + @ray_method(concurrency_group=AGENT_LOOP_CONCURRENCY_GROUP_GENERATE) + async def collect_rollout_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: + return await self.agent_loop.collect_rollout_group(rollout_state, **kwargs) + @ray_method async def get_rollout_ctl(self): return self.agent_loop.rollout_ctl diff --git a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py index 8f2edb2cd6..e988b793f6 100644 --- a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py @@ -9,6 +9,7 @@ from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast from xtuner.v1.data_proto.rl_data import Status from xtuner.v1.rl.agent_loop import AgentLoopConfig, IsValidSampleFn +from xtuner.v1.rl.distillation import DistillationConfig from xtuner.v1.rl.judger import ComposedJudgerConfig, JudgerConfig, build_judger from xtuner.v1.rl.replay_buffer import ReplayBuffer from xtuner.v1.rl.rollout import RolloutController @@ -18,6 +19,7 @@ _MANAGER_STATE_PATH, _STATUS_POLL_INTERVAL_S, _TASK_CHECKPOINT_DIR, + IsValidSampleFn, ProduceBatchResult, _TaskRunner, _TaskSamplerView, @@ -129,6 +131,7 @@ def build( replay_buffer: ReplayBuffer, logger=None, sync_weights_interval: int = 1, + distillation_config: DistillationConfig | None = None, ) -> "AgentLoopManager": tasks = self.tasks if isinstance(self.tasks, list) else [self.tasks] if not tasks: @@ -146,6 +149,7 @@ def build( judger=build_judger(task_cfg.judger_config) if task_cfg.judger_config is not None else None, logger=logger, is_valid_sample_fn=task_cfg.is_valid_sample_fn, + distillation_config=distillation_config, ) produce_strategy = task_cfg.produce_strategy_config.build( sync_weights_interval=sync_weights_interval, @@ -158,6 +162,7 @@ def build( agent_loop=agent_loop, produce_strategy=produce_strategy, sampler=sampler, + is_valid_sample_fn=task_cfg.filter_func, weight=task_cfg.weight, order=order, ) diff --git a/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py index 232223caed..dd245a9695 100644 --- a/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py @@ -27,6 +27,7 @@ _MANAGER_STATE_PATH, _STATUS_POLL_INTERVAL_S, _TASK_CHECKPOINT_DIR, + IsValidSampleFn, ProduceBatchResult, ProduceBatchStatus, _TaskRunner, @@ -69,6 +70,7 @@ def build( replay_buffer: ReplayBuffer, logger=None, sync_weights_interval: int = 1, + distillation_config: DistillationConfig | None = None, ) -> "DisaggAgentLoopManager": tasks = self.tasks if isinstance(self.tasks, list) else [self.tasks] if not tasks: @@ -86,6 +88,7 @@ def build( judger=build_judger(task_cfg.judger_config) if task_cfg.judger_config is not None else None, logger=logger, is_valid_sample_fn=task_cfg.is_valid_sample_fn, + distillation_config=distillation_config, ) produce_strategy = task_cfg.produce_strategy_config.build( sync_weights_interval=sync_weights_interval, @@ -98,6 +101,7 @@ def build( agent_loop=agent_loop, produce_strategy=produce_strategy, sampler=sampler, + is_valid_sample_fn=task_cfg.filter_func, weight=task_cfg.weight, order=order, ) diff --git a/xtuner/v1/rl/agent_loop_manager/disagg_producer.py b/xtuner/v1/rl/agent_loop_manager/disagg_producer.py index 76503ebe17..2fdde63ab1 100644 --- a/xtuner/v1/rl/agent_loop_manager/disagg_producer.py +++ b/xtuner/v1/rl/agent_loop_manager/disagg_producer.py @@ -424,7 +424,7 @@ async def produce_batch(self, ctx: DisaggProduceContext) -> ProduceBatchStatus: async def spawn_one() -> asyncio.Task: rollout_state = await ctx.sample_group(from_expired_pool=sample_expired) return create_task( - ctx.generate_group( + ctx.collect_rollout_group( rollout_state, enable_partial_rollout=self.enable_partial_rollout, ) diff --git a/xtuner/v1/rl/agent_loop_manager/produce_utils.py b/xtuner/v1/rl/agent_loop_manager/produce_utils.py index 9135fb395e..fef9a0f746 100644 --- a/xtuner/v1/rl/agent_loop_manager/produce_utils.py +++ b/xtuner/v1/rl/agent_loop_manager/produce_utils.py @@ -127,7 +127,7 @@ async def sample_group(self, *, from_expired_pool: bool) -> list[RolloutState]: group_status = [Status.EXPIRED, Status.ABORTED] if from_expired_pool else [Status.ABORTED] return await self.sampler.sample(task_name=self.task_name, group_status=group_status) - async def generate_group( + async def collect_rollout_group( self, rollout_state: list[RolloutState], *, @@ -141,13 +141,15 @@ async def generate_group( start = time.perf_counter() if isinstance(self.agent_loop, ray.actor.ActorHandle): - result = await self.agent_loop.generate_group.remote( + result = await self.agent_loop.collect_rollout_group.remote( rollout_state, + is_valid_sample_func=self.is_valid_sample_fn, enable_partial_rollout=enable_partial_rollout, ) else: - result = await self.agent_loop.generate_group( + result = await self.agent_loop.collect_rollout_group( rollout_state, + is_valid_sample_func=self.is_valid_sample_fn, enable_partial_rollout=enable_partial_rollout, ) elapsed = time.perf_counter() - start @@ -168,9 +170,6 @@ async def put_generated_group(self, group: list[RolloutState]) -> bool: rewards_count = 0 for item in group: if item.reward is None or "score" not in item.reward: - logger.warning( - f"Missing reward score in item (rollout_id: {item.rollout_id}) of generated group for task {self.task_name}. This item will be skipped in reward statistics." - ) continue # TODO: 在 agent 存在一拆多的情况下，这个 raw reward 统计会不准，但是考虑到在这区分有点 hard code，应该暂时不处理 rewards_sum += float(item.reward["score"]) @@ -257,6 +256,7 @@ class _TaskRunner: agent_loop: AgentLoopSpec produce_strategy: Any sampler: Sampler + is_valid_sample_fn: IsValidSampleFn | None = None weight: float = 1.0 order: int = 0 diff --git a/xtuner/v1/rl/agent_loop_manager/producer.py b/xtuner/v1/rl/agent_loop_manager/producer.py index 2f7f660e70..75c89bf1d1 100644 --- a/xtuner/v1/rl/agent_loop_manager/producer.py +++ b/xtuner/v1/rl/agent_loop_manager/producer.py @@ -283,7 +283,7 @@ async def produce_batch(self, ctx: ProduceContext) -> None: for _ in range(ctx.task_batch_size): rollout_state = await ctx.sampler.sample(task_name=ctx.task_name) - task = create_task(ctx.generate_group(rollout_state)) + task = create_task(ctx.collect_rollout_group(rollout_state)) pending_tasks.add(task) logger.info(f"[SyncProduceStrategy] Started {len(pending_tasks)} initial tasks.") @@ -316,7 +316,7 @@ async def produce_batch(self, ctx: ProduceContext) -> None: completed_sample_count, ctx.task_batch_size ): rollout_state = await ctx.sampler.sample(task_name=ctx.task_name) - task = create_task(ctx.generate_group(rollout_state)) + task = create_task(ctx.collect_rollout_group(rollout_state)) pending_tasks.add(task) progress_displayer.close() @@ -405,7 +405,7 @@ async def produce_batch(self, ctx: ProduceContext) -> None: async def spawn_one() -> asyncio.Task: rollout_state = await ctx.sample_group(from_expired_pool=sample_expired) return create_task( - ctx.generate_group( + ctx.collect_rollout_group( rollout_state, enable_partial_rollout=self.enable_partial_rollout, ) diff --git a/xtuner/v1/rl/distillation/__init__.py b/xtuner/v1/rl/distillation/__init__.py new file mode 100644 index 0000000000..85c604ac17 --- /dev/null +++ b/xtuner/v1/rl/distillation/__init__.py @@ -0,0 +1,30 @@ +from .config import ( + DistillationConfig, + RolloutTeacherConfig, + RolloutTeacherLaunchConfig, + TeacherConfig, + TrainTeacherConfig, +) +from .rollout_teacher_client import ( + RolloutTeacherClient, + RolloutTeacherReplicaRouter, + route_rollout_teacher_client, + validate_opd_sample_params, +) +from .train_teacher_manager import TrainTeacherManager, TrainTeacherOutputs, TrainTeacherTimings + + +__all__ = [ + "DistillationConfig", + "RolloutTeacherConfig", + "RolloutTeacherLaunchConfig", + "TeacherConfig", + "TrainTeacherConfig", + "RolloutTeacherClient", + "RolloutTeacherReplicaRouter", + "route_rollout_teacher_client", + "TrainTeacherManager", + "TrainTeacherOutputs", + "TrainTeacherTimings", + "validate_opd_sample_params", +] diff --git a/xtuner/v1/rl/distillation/config.py b/xtuner/v1/rl/distillation/config.py new file mode 100644 index 0000000000..8630e32726 --- /dev/null +++ b/xtuner/v1/rl/distillation/config.py @@ -0,0 +1,147 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Literal, cast + +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator + +from xtuner.v1.config.fsdp import FSDPConfig +from xtuner.v1.model.base import TransformerConfig +from xtuner.v1.model.compose.base import BaseComposeConfig +from xtuner.v1.rl.loss.distillation_loss import DistillationLossConfig + + +class RolloutTeacherLaunchConfig(BaseModel): + model_config = ConfigDict(extra="forbid") + + model_path: str | Path + num_workers: int = Field(default=1, gt=0) + server_port: int = Field(gt=0, le=65535) + dtype: Literal["auto", "float16", "bfloat16"] = "bfloat16" + tensor_parallel_size: int = Field(default=1, gt=0) + expert_parallel_size: int = Field(default=1, gt=0) + context_length: int | None = Field(default=None, gt=0) + max_batch_size: int | None = Field(default=None, gt=0) + log_level: Literal["critical", "error", "warning", "info", "debug"] | None = None + chunked_prefill_size: int | None = Field(default=4096, gt=0) + max_prefill_token_num: int | None = Field(default=4096, gt=0) + gpu_memory_utilization: float = Field(default=0.6, gt=0.0, le=1.0) + + +class RolloutTeacherConfig(BaseModel): + """Teacher served by an external rollout/inference engine.""" + + model_config = ConfigDict(extra="forbid") + + name: str + num_replicas: int = Field(default=1, gt=0) + endpoints: list[str] = Field(default_factory=list) + api_key: str | None = None + request_timeout_s: float = Field(default=1200.0, gt=0.0) + max_retry_per_sample: int = Field(default=2, ge=0) + max_concurrency: int = Field(default=128, gt=0) + enable_prefix_caching: bool = False + launch_config: RolloutTeacherLaunchConfig | None = None + + +class TrainTeacherConfig(BaseModel): + """Frozen Teacher loaded and executed by each training worker.""" + + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) + + name: str + model_path: str | Path + model_cfg: TransformerConfig | BaseComposeConfig + fsdp_cfg: FSDPConfig | None = None + + +TeacherConfig = RolloutTeacherConfig | TrainTeacherConfig + + +class DistillationConfig(BaseModel): + """Distillation objective, Teacher runtimes, and data-source routing.""" + + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) + + loss_config: DistillationLossConfig + teachers: list[TeacherConfig] = Field(min_length=1) + data_source_teacher_map: dict[str, str] = Field(min_length=1) + _teacher_index_by_data_source: dict[str, int] = PrivateAttr(default_factory=dict) + + @model_validator(mode="after") + def validate_config(self) -> DistillationConfig: + teacher_names = [teacher.name for teacher in self.teachers] + if len(teacher_names) != len(set(teacher_names)): + raise ValueError("Distillation teacher names must be unique") + unknown_teachers = set(self.data_source_teacher_map.values()) - set(teacher_names) + if unknown_teachers: + raise ValueError(f"data_source_teacher_map references unknown teachers: {sorted(unknown_teachers)}") + + if self.loss_config.uses_sampled_token_targets: + teacher_types = {type(teacher) for teacher in self.teachers} + if len(teacher_types) != 1: + raise ValueError( + "Sampled-token distillation requires all teachers to use the same runtime type; " + "mixing RolloutTeacherConfig and TrainTeacherConfig is not supported" + ) + else: + incompatible_teachers = [ + teacher.name for teacher in self.teachers if not isinstance(teacher, TrainTeacherConfig) + ] + if incompatible_teachers: + raise ValueError( + f"loss_mode={self.loss_config.loss_mode!r} requires " + f"TrainTeacherConfig, got incompatible teachers: {incompatible_teachers}" + ) + teacher_index_by_name = {teacher.name: index for index, teacher in enumerate(self.teachers)} + self._teacher_index_by_data_source = { + data_source: teacher_index_by_name[teacher_name] + for data_source, teacher_name in self.data_source_teacher_map.items() + } + return self + + @property + def rollout_teachers(self) -> list[RolloutTeacherConfig]: + return [teacher for teacher in self.teachers if isinstance(teacher, RolloutTeacherConfig)] + + @property + def train_teachers(self) -> list[TrainTeacherConfig]: + return [teacher for teacher in self.teachers if isinstance(teacher, TrainTeacherConfig)] + + @property + def teacher_index_by_data_source(self) -> dict[str, int]: + return self._teacher_index_by_data_source + + def validate_student_model(self, student_model_cfg: TransformerConfig | BaseComposeConfig) -> None: + """Validate contracts that require both Student and TrainTeacher + configs.""" + if not self.train_teachers: + return + student_lm_cfg = ( + student_model_cfg.text_config if isinstance(student_model_cfg, BaseComposeConfig) else student_model_cfg + ) + student_vocab_size = cast(TransformerConfig, student_lm_cfg).vocab_size + for teacher in self.train_teachers: + teacher_lm_cfg = ( + teacher.model_cfg.text_config + if isinstance(teacher.model_cfg, BaseComposeConfig) + else teacher.model_cfg + ) + teacher_vocab_size = cast(TransformerConfig, teacher_lm_cfg).vocab_size + if teacher_vocab_size != student_vocab_size: + raise ValueError( + f"distillation teacher {teacher.name!r} vocab_size={teacher_vocab_size} does not match " + f"student vocab_size={student_vocab_size}" + ) + + def resolve_teacher_endpoints( + self, + endpoint_map: dict[str, list[str]], + ) -> DistillationConfig: + teachers = [ + teacher + if not isinstance(teacher, RolloutTeacherConfig) or teacher.launch_config is None + else teacher.model_copy(update={"endpoints": endpoint_map[teacher.name]}) + for teacher in self.teachers + ] + return self.model_copy(update={"teachers": teachers}) diff --git a/xtuner/v1/rl/distillation/rollout_teacher_client.py b/xtuner/v1/rl/distillation/rollout_teacher_client.py new file mode 100644 index 0000000000..b16b9fef19 --- /dev/null +++ b/xtuner/v1/rl/distillation/rollout_teacher_client.py @@ -0,0 +1,339 @@ +from __future__ import annotations + +import asyncio +import math +import os +import time +from hashlib import blake2b +from typing import Any, Literal, cast + +import httpx + +from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status +from xtuner.v1.utils import get_logger + +from .config import RolloutTeacherConfig + + +logger = get_logger() + + +def validate_opd_sample_params(sample_params: SampleParams) -> None: + identity_sampling_params: dict[str, Any] = { + "temperature": 1.0, + "top_p": 1.0, + "top_k": 0, + "repetition_penalty": 1.0, + "presence_penalty": 0.0, + "frequency_penalty": 0.0, + "min_tokens": 0, + } + non_identity_params = { + name: getattr(sample_params, name) + for name, expected in identity_sampling_params.items() + if getattr(sample_params, name) != expected + } + if non_identity_params: + raise ValueError(f"PG-OPD requires identity student sampling, got {non_identity_params}") + if not sample_params.return_logprob or not sample_params.return_token_ids: + raise ValueError("PG-OPD requires return_logprob=True and return_token_ids=True") + + +class RolloutTeacherReplicaRouter: + """Resolve a physical teacher replica for one scoring request.""" + + def __init__(self, num_replicas: int) -> None: + if num_replicas <= 0: + raise ValueError(f"num_replicas must be positive, got {num_replicas}") + self._num_replicas = num_replicas + + def resolve_replica_idx( + self, + *, + teacher_name: str, + data_source: str, + group_id: int | None, + ) -> int: + if self._num_replicas == 1: + return 0 + routing_key = f"{teacher_name}\0{data_source}\0{group_id}".encode() + digest = blake2b(routing_key, digest_size=8).digest() + return int.from_bytes(digest, byteorder="big") % self._num_replicas + + +class RolloutTeacherClient: + """Asynchronous teacher client scoped to one AgentLoop.""" + + def __init__(self, config: RolloutTeacherConfig) -> None: + self.config = config + self.name = config.name + self.backend = self._resolve_backend_from_env() + self.urls = [f"{endpoint.rstrip('/')}/generate" for endpoint in config.endpoints] + self._semaphores = [asyncio.Semaphore(config.max_concurrency) for _ in self.urls] + self._replica_router = RolloutTeacherReplicaRouter(len(self.urls)) + + headers = {"Content-Type": "application/json"} + if config.api_key is not None: + headers["Authorization"] = f"Bearer {config.api_key}" + self._client = httpx.AsyncClient(headers=headers, timeout=config.request_timeout_s) + + async def compute_logprobs(self, state: RolloutState) -> RolloutState: + start = time.perf_counter() + try: + prompt_ids = cast(list[int], state.prompt_ids) + response_ids = cast(list[int], state.response_ids) + if not prompt_ids or not response_ids: + state.status = Status.FAILED + state.error_msg = f"Teacher {self.name!r} scoring requires non-empty prompt_ids and response_ids" + return state + routed_replica_idx = self._replica_router.resolve_replica_idx( + teacher_name=self.name, + data_source=str(state.extra_fields.get("origin_data_source", "")), + group_id=state.group_id, + ) + image_data = state.extra_fields.get("image_data") + expanded_prompt_len = None + # Recompute the final prompt token so the first response token's + # logprob remains available while the earlier prompt can be reused. + logprob_start_len = 0 + if self.config.enable_prefix_caching: + if not image_data: + logprob_start_len = len(prompt_ids) - 1 + else: + expanded_prompt_len = len(state.extra_fields["train_prompt_ids"]) + logprob_start_len = expanded_prompt_len - 1 + payload = self._construct_payload( + prompt_ids, + response_ids, + logprob_start_len=logprob_start_len, + image_data=image_data, + ) + + attempt_idx = 0 + while True: + replica_idx = (routed_replica_idx + attempt_idx) % len(self.urls) + url = self.urls[replica_idx] + try: + async with self._semaphores[replica_idx]: + response = await self._client.post(url, json=payload) + response.raise_for_status() + teacher_tokens, teacher_logprobs = self._parse_response( + response, + response_ids, + logprob_start_len=logprob_start_len, + expanded_prompt_len=expanded_prompt_len, + ) + state.teacher_tokens = teacher_tokens + state.teacher_logprobs = teacher_logprobs + return state + except (httpx.HTTPStatusError, httpx.RequestError, ValueError) as exc: + if attempt_idx >= self.config.max_retry_per_sample: + state.status = Status.FAILED + state.error_msg = ( + f"Teacher {self.name!r} logprobs calculation failed after {attempt_idx + 1} attempts; " + f"group_id={state.group_id}; routed_replica={routed_replica_idx}; " + f"replica={replica_idx}; endpoint={url}; last_error={exc}" + ) + logger.warning(state.error_msg) + return state + attempt_idx += 1 + await asyncio.sleep(0.1) + finally: + state.extra_fields["teacher_score_time_s"] = time.perf_counter() - start + + @staticmethod + def _resolve_backend_from_env() -> Literal["sglang", "lmdeploy"]: + use_sglang = os.environ.get("XTUNER_USE_SGLANG", "0") == "1" + use_lmdeploy = os.environ.get("XTUNER_USE_LMDEPLOY", "0") == "1" + use_vllm = os.environ.get("XTUNER_USE_VLLM", "0") == "1" + + if use_vllm: + raise RuntimeError("RolloutTeacherClient supports only SGLang or LMDeploy, not vLLM") + if use_sglang == use_lmdeploy: + raise RuntimeError("Exactly one of XTUNER_USE_SGLANG and XTUNER_USE_LMDEPLOY must be set to 1") + return "sglang" if use_sglang else "lmdeploy" + + def _construct_payload( + self, + prompt_ids: list[int], + response_ids: list[int], + *, + logprob_start_len: int, + image_data: Any | None = None, + ) -> dict[str, Any]: + if self.backend == "sglang": + payload = self._construct_sglang_payload( + prompt_ids, + response_ids, + logprob_start_len=logprob_start_len, + ) + elif self.backend == "lmdeploy": + payload = self._construct_lmdeploy_payload( + prompt_ids, + response_ids, + logprob_start_len=logprob_start_len, + ) + else: + raise RuntimeError(f"Unsupported teacher backend: {self.backend}") + if image_data: + payload["image_data"] = image_data + return payload + + @staticmethod + def _construct_sglang_payload( + prompt_ids: list[int], + response_ids: list[int], + *, + logprob_start_len: int, + ) -> dict[str, Any]: + return { + "input_ids": prompt_ids + response_ids, + "sampling_params": { + "max_new_tokens": 0, + "temperature": 0, + "skip_special_tokens": False, + }, + "return_logprob": True, + "logprob_start_len": logprob_start_len, + "top_logprobs_num": 0, + "stream": False, + } + + @staticmethod + def _construct_lmdeploy_payload( + prompt_ids: list[int], + response_ids: list[int], + *, + logprob_start_len: int, + ) -> dict[str, Any]: + return { + "input_ids": prompt_ids + response_ids, + "return_logprob": True, + "logprob_start_len": logprob_start_len, + "max_tokens": 0, + "stream": False, + } + + def _parse_response( + self, + response: httpx.Response, + response_ids: list[int], + *, + logprob_start_len: int, + expanded_prompt_len: int | None = None, + ) -> tuple[list[int], list[float]]: + if self.backend == "sglang": + response_logprobs = self._parse_sglang_response( + response, + response_ids, + logprob_start_len=logprob_start_len, + expanded_prompt_len=expanded_prompt_len, + ) + else: + response_logprobs = self._parse_lmdeploy_response( + response, + response_ids, + logprob_start_len=logprob_start_len, + expanded_prompt_len=expanded_prompt_len, + ) + return self._validate_response_logprobs(response_logprobs, response_ids) + + @staticmethod + def _parse_sglang_response( + response: httpx.Response, + response_ids: list[int], + *, + logprob_start_len: int, + expanded_prompt_len: int | None = None, + ) -> list[Any]: + raw_logprobs = RolloutTeacherClient._get_input_token_logprobs(response) + prompt_token_count = response.json()["meta_info"]["prompt_tokens"] + if expanded_prompt_len is not None: + expected_prompt_token_count = expanded_prompt_len + len(response_ids) + if prompt_token_count != expected_prompt_token_count: + raise ValueError( + "SGLang expanded prompt length mismatch: " + f"expected {expected_prompt_token_count} total input tokens " + f"({expanded_prompt_len} prompt + {len(response_ids)} response), " + f"got {prompt_token_count}" + ) + # SGLang includes an unscorable placeholder at the requested boundary, + # so N processed tokens with boundary S produce N-S rows. + expected_rows = prompt_token_count - logprob_start_len + if len(raw_logprobs) != expected_rows: + raise ValueError( + "SGLang teacher logprob length mismatch: " + f"expected {expected_rows} rows for logprob_start_len={logprob_start_len}, " + f"got {len(raw_logprobs)}" + ) + return raw_logprobs[-len(response_ids) :] + + @staticmethod + def _parse_lmdeploy_response( + response: httpx.Response, + response_ids: list[int], + *, + logprob_start_len: int, + expanded_prompt_len: int | None = None, + ) -> list[Any]: + raw_logprobs = RolloutTeacherClient._get_input_token_logprobs(response) + prompt_token_count = response.json()["meta_info"]["prompt_tokens"] + if expanded_prompt_len is not None: + expected_prompt_token_count = expanded_prompt_len + len(response_ids) + if prompt_token_count != expected_prompt_token_count: + raise ValueError( + "LMDeploy expanded prompt length mismatch: " + f"expected {expected_prompt_token_count} total input tokens " + f"({expanded_prompt_len} prompt + {len(response_ids)} response), " + f"got {prompt_token_count}" + ) + # LMDeploy omits the unscorable boundary row, hence one fewer row than + # SGLang for the same processed input and logprob boundary. + expected_rows = prompt_token_count - logprob_start_len - 1 + if len(raw_logprobs) != expected_rows: + raise ValueError( + "LMDeploy teacher logprob length mismatch: " + f"expected {expected_rows} rows for logprob_start_len={logprob_start_len}, " + f"got {len(raw_logprobs)}" + ) + return raw_logprobs[-len(response_ids) :] + + @staticmethod + def _get_input_token_logprobs(response: httpx.Response) -> list[Any]: + try: + raw_logprobs = response.json()["meta_info"]["input_token_logprobs"] + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("Invalid teacher response") from exc + if not isinstance(raw_logprobs, list): + raise ValueError("Invalid teacher response") + return raw_logprobs + + @staticmethod + def _validate_response_logprobs( + response_logprobs: list[Any], + response_ids: list[int], + ) -> tuple[list[int], list[float]]: + try: + teacher_tokens = [item[1] for item in response_logprobs] + teacher_logprobs = [float(item[0]) for item in response_logprobs] + except (TypeError, IndexError, ValueError) as exc: + raise ValueError("Invalid teacher response") from exc + + if len(teacher_logprobs) != len(response_ids): + raise ValueError("Teacher logprob length mismatch") + if teacher_tokens != response_ids: + raise ValueError("Teacher token ids mismatch") + if not all(math.isfinite(logprob) for logprob in teacher_logprobs): + raise ValueError("Teacher logprobs contain NaN or Inf") + return teacher_tokens, teacher_logprobs + + +def route_rollout_teacher_client( + state: RolloutState, + *, + data_source_teacher_map: dict[str, str], + teacher_clients: dict[str, RolloutTeacherClient], +) -> RolloutTeacherClient: + data_source = state.extra_fields["origin_data_source"] + teacher_name = data_source_teacher_map[data_source] + return teacher_clients[teacher_name] diff --git a/xtuner/v1/rl/distillation/train_teacher_manager.py b/xtuner/v1/rl/distillation/train_teacher_manager.py new file mode 100644 index 0000000000..dc9b2099f4 --- /dev/null +++ b/xtuner/v1/rl/distillation/train_teacher_manager.py @@ -0,0 +1,239 @@ +from __future__ import annotations + +import time +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Iterator, cast + +import torch + +from xtuner.v1.data_proto.sequence_context import DSATopKCacheState, SequenceContext +from xtuner.v1.loss import LogProbConfig, LogProbContext, TopKLogProbConfig +from xtuner.v1.model.compose.base import BaseComposeConfig +from xtuner.v1.rl.loss import DistillationLossConfig +from xtuner.v1.rl.model_utils import FrozenModel, build_frozen_model +from xtuner.v1.utils import get_device, get_torch_device_module + +from .config import DistillationConfig + + +DEVICE = get_device() +DEVICE_MODULE = get_torch_device_module() + + +@dataclass +class TrainTeacherTimings: + """Wall-clock seconds spent in the frozen Teacher lifecycle.""" + + compute: float = 0.0 + onload: float = 0.0 + offload: float = 0.0 + + +@dataclass +class TrainTeacherOutputs: + teacher_logprobs: list[torch.Tensor] + target_token_ids: list[torch.Tensor] | None = None + timings: TrainTeacherTimings = field(default_factory=TrainTeacherTimings) + + +class TrainTeacherManager: + """Execute training-side Teachers in one TrainingWorker process. + + The manager owns Teacher model construction, deterministic Teacher-major scheduling, CPU/device residency, and + sampled-token or top-k output calculation. The caller remains responsible for preparing distributed inputs and + swapping the Actor and optimizer around the Teacher phase. + """ + + def __init__(self, distillation_config: DistillationConfig, *, chunk_size: int | None) -> None: + self.loss_config = cast(DistillationLossConfig, distillation_config.loss_config) + mode = "chunk" if chunk_size is not None else "eager" + self.logprob_config = LogProbConfig(chunk_size=chunk_size, mode=mode) + self.topk_logprob_config: TopKLogProbConfig | None = None + if self.loss_config.uses_topk_targets: + self.topk_logprob_config = TopKLogProbConfig( + top_k=cast(int, self.loss_config.top_k), + chunk_size=chunk_size, + mode=mode, + ) + + self._teachers: list[FrozenModel] = [ + build_frozen_model(teacher.model_cfg, teacher.model_path, teacher.fsdp_cfg) + for teacher in distillation_config.train_teachers + ] + self._teacher_is_composed = [ + isinstance(teacher.model_cfg, BaseComposeConfig) for teacher in distillation_config.train_teachers + ] + self._teachers_by_name = { + teacher_config.name: teacher + for teacher_config, teacher in zip(distillation_config.train_teachers, self._teachers) + } + + def compute_logprobs( + self, + *, + seq_ctx_list: list[SequenceContext], + shifted_labels_list: list[torch.Tensor], + teacher_indices_list: list[torch.Tensor], + ) -> TrainTeacherOutputs: + timings = TrainTeacherTimings() + if self.loss_config.uses_sampled_token_targets: + return TrainTeacherOutputs( + teacher_logprobs=self._compute_sampled_logprobs( + seq_ctx_list, + shifted_labels_list, + teacher_indices_list, + timings, + ), + timings=timings, + ) + + target_token_ids, teacher_logprobs = self._compute_topk_targets( + seq_ctx_list, + teacher_indices_list, + timings, + ) + return TrainTeacherOutputs( + teacher_logprobs=teacher_logprobs, + target_token_ids=target_token_ids, + timings=timings, + ) + + def offload_all_to_cpu(self) -> None: + for teacher in self._teachers: + self._offload_to_cpu(teacher) + + def offload_to_disk(self, teacher_name: str) -> None: + """Reserve the disk-offload lifecycle boundary for a later backend.""" + if teacher_name not in self._teachers_by_name: + raise KeyError(f"Unknown training Teacher: {teacher_name!r}") + raise NotImplementedError("Train Teacher disk offload is not implemented") + + @staticmethod + def _offload_to_cpu(teacher: FrozenModel) -> None: + teacher.to_device("cpu") + if hasattr(DEVICE_MODULE, "empty_cache"): + DEVICE_MODULE.empty_cache() + + @staticmethod + def _synchronize_device() -> None: + if str(DEVICE) != "cpu" and hasattr(DEVICE_MODULE, "synchronize"): + DEVICE_MODULE.synchronize() + + @contextmanager + def _teacher_on_device( + self, + teacher: FrozenModel, + timings: TrainTeacherTimings, + ) -> Iterator[None]: + onload_begin = time.perf_counter() + teacher.to_device(DEVICE) + timings.onload += time.perf_counter() - onload_begin + + compute_begin = time.perf_counter() + try: + yield + self._synchronize_device() + finally: + timings.compute += time.perf_counter() - compute_begin + offload_begin = time.perf_counter() + self._offload_to_cpu(teacher) + timings.offload += time.perf_counter() - offload_begin + + @staticmethod + def _teacher_seq_ctx(seq_ctx: SequenceContext, *, is_composed: bool) -> SequenceContext: + overrides = { + "rollout_routed_experts": None, + "offload_rollout_routed_experts": False, + "dsa_topk_cache": DSATopKCacheState(), + } + if not is_composed: + # A VLM tokenizer supplies 3D M-RoPE position ids for every sample, + # including text-only samples. Plain language-model Teachers need + # standard 2D packed positions. Passing ``None`` makes + # SequenceContext rebuild them from the packed sequence lengths. + # Visual fields are irrelevant to the text Teacher and may refer to + # a pack routed to another Teacher, so do not retain them here. + overrides.update( + position_ids=None, + image_grid_thw=None, + deepstack_visual_embeds=None, + visual_pos_masks=None, + pixel_values=None, + inputs_embeds=None, + num_img_tokens=None, + ) + return seq_ctx.copy(**overrides) + + def _compute_topk_targets( + self, + seq_ctx_list: list[SequenceContext], + teacher_indices_list: list[torch.Tensor], + timings: TrainTeacherTimings, + ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: + top_k = cast(int, self.loss_config.top_k) + topk_logprob_config = cast(TopKLogProbConfig, self.topk_logprob_config) + target_ids = [ + torch.zeros((*teacher_indices.shape, top_k), dtype=torch.long, device=DEVICE) + for teacher_indices in teacher_indices_list + ] + target_logprobs = [ + torch.zeros((*teacher_indices.shape, top_k), dtype=torch.float32, device=DEVICE) + for teacher_indices in teacher_indices_list + ] + + # All ranks iterate Teachers and local packs in the same order so every + # FSDP rank enters the same collective sequence, even when a rank has no + # tokens routed to a particular Teacher. + for teacher_index, teacher in enumerate(self._teachers): + with self._teacher_on_device(teacher, timings): + # Every rank forwards every local pack before selecting routed + # tokens so all ranks enter the same FSDP collective sequence. + for batch_index, (seq_ctx, teacher_indices) in enumerate(zip(seq_ctx_list, teacher_indices_list)): + loss_ctx = topk_logprob_config.build(data={}) + assert loss_ctx is not None + with torch.no_grad(): + output = teacher( + seq_ctx=self._teacher_seq_ctx( + seq_ctx, + is_composed=self._teacher_is_composed[teacher_index], + ), + loss_ctx={"lm": loss_ctx}, + ) + selected = teacher_indices == teacher_index + target_ids[batch_index][selected] = cast(torch.Tensor, output.logits)[selected] + target_logprobs[batch_index][selected] = cast(torch.Tensor, output.loss)[selected] + return target_ids, target_logprobs + + def _compute_sampled_logprobs( + self, + seq_ctx_list: list[SequenceContext], + shifted_labels_list: list[torch.Tensor], + teacher_indices_list: list[torch.Tensor], + timings: TrainTeacherTimings, + ) -> list[torch.Tensor]: + target_logprobs = [ + torch.zeros_like(shifted_labels, dtype=torch.float32) for shifted_labels in shifted_labels_list + ] + + # Keep the Teacher-major schedule identical across ranks for FSDP. + for teacher_index, teacher in enumerate(self._teachers): + with self._teacher_on_device(teacher, timings): + for batch_index, (seq_ctx, shifted_labels, teacher_indices) in enumerate( + zip(seq_ctx_list, shifted_labels_list, teacher_indices_list) + ): + loss_ctx = cast( + LogProbContext, + self.logprob_config.build(data={"shifted_labels": shifted_labels}), + ) + with torch.no_grad(): + output = teacher( + seq_ctx=self._teacher_seq_ctx( + seq_ctx, + is_composed=self._teacher_is_composed[teacher_index], + ), + loss_ctx={"lm": loss_ctx}, + ) + selected = teacher_indices == teacher_index + target_logprobs[batch_index][selected] = cast(torch.Tensor, output.loss)[selected] + return target_logprobs diff --git a/xtuner/v1/rl/loss/__init__.py b/xtuner/v1/rl/loss/__init__.py index 371ea27bff..45d3e1d069 100644 --- a/xtuner/v1/rl/loss/__init__.py +++ b/xtuner/v1/rl/loss/__init__.py @@ -5,6 +5,13 @@ compute_kl_loss_weight, finalize_train_policy_metrics, ) +from .distillation_loss import ( + DistillationLossConfig, + DistillationLossContext, + DistillationLossKwargs, + compute_topk_distillation_kl, + finalize_distillation_metrics, +) from .grpo_loss import GRPOLossConfig, GRPOLossContext, GRPOLossKwargs from .loss_fn import check_config, get_policy_loss_fn, kl_penalty, pg_loss_fn, register_policy_loss, sft_loss_fn from .oreal_loss import OrealLossConfig, OrealLossContext, OrealLossKwargs diff --git a/xtuner/v1/rl/loss/base_loss.py b/xtuner/v1/rl/loss/base_loss.py index 108af41ed8..e2eccc9676 100644 --- a/xtuner/v1/rl/loss/base_loss.py +++ b/xtuner/v1/rl/loss/base_loss.py @@ -304,4 +304,5 @@ def reduce_values(keys, op): extra_info_dict["reduced_train_policy_ratio_min"] = ratio_min # legacy metric，keep here extra_info_dict["max_ratio"] = ratio_max + return extra_info_dict diff --git a/xtuner/v1/rl/loss/distillation_loss.py b/xtuner/v1/rl/loss/distillation_loss.py new file mode 100644 index 0000000000..e655c06ec7 --- /dev/null +++ b/xtuner/v1/rl/loss/distillation_loss.py @@ -0,0 +1,447 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from typing import Any, Literal, cast + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from pydantic import Field, model_validator +from torch.distributed.device_mesh import DeviceMesh +from typing_extensions import Self + +from xtuner.v1.loss.utils import sp_split +from xtuner.v1.utils import get_logger +from xtuner.v1.utils.device import get_device + +from ..utils import gather_logprobs +from .base_loss import ( + BaseRLLossConfig, + BaseRLLossContext, + BaseRLLossKwargs, + compute_kl_loss_weight, +) +from .loss_fn import get_policy_loss_fn, kl_penalty + + +DEVICE = get_device() +logger = get_logger() + +TopKDistillationMode = Literal["forward", "reverse", "forward_kl_topk"] +DistillationLossMode = Literal["k1", "k3", "forward", "reverse", "forward_kl_topk"] + + +def compute_topk_distillation_kl( + student_logprobs: torch.Tensor, + teacher_logprobs: torch.Tensor, + loss_mode: TopKDistillationMode, + log_prob_min_clamp: float | None = None, + loss_max_clamp: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute a KL loss over teacher-selected Top-K tokens.""" + student_probs = student_logprobs.exp() + teacher_probs = teacher_logprobs.exp() + student_selected_mass = student_probs.sum(dim=-1) + teacher_selected_mass = teacher_probs.sum(dim=-1) + + if loss_mode == "forward_kl_topk": + if log_prob_min_clamp is not None: + student_logprobs = student_logprobs.clamp_min(log_prob_min_clamp) + teacher_logprobs = teacher_logprobs.clamp_min(log_prob_min_clamp) + teacher_probs = teacher_logprobs.exp() + loss = (teacher_probs * (teacher_logprobs - student_logprobs)).sum(dim=-1) + if loss_max_clamp is not None: + loss = loss.clamp(min=-loss_max_clamp, max=loss_max_clamp) + return loss, student_selected_mass, teacher_selected_mass + + eps = torch.finfo(student_logprobs.dtype).tiny + student_tail = (1.0 - student_selected_mass).clamp_min(eps) + teacher_tail = (1.0 - teacher_selected_mass).clamp_min(eps) + if loss_mode == "forward": + selected_kl = teacher_probs * (teacher_logprobs - student_logprobs) + tail_kl = teacher_tail * (teacher_tail.log() - student_tail.log()) + else: + selected_kl = student_probs * (student_logprobs - teacher_logprobs) + tail_kl = student_tail * (student_tail.log() - teacher_tail.log()) + loss = selected_kl.sum(dim=-1) + tail_kl + if loss_max_clamp is not None: + loss = loss.clamp(max=loss_max_clamp) + return loss, student_selected_mass, teacher_selected_mass + + +class DistillationLossConfig(BaseRLLossConfig): + """Configuration shared by sampled-token PG-OPD and direct GKD losses.""" + + loss_mode: DistillationLossMode = "k1" + use_policy_gradient: bool = True + task_adv_weight: float = Field(default=0.0, ge=0.0) + distillation_loss_weight: float = Field(default=1.0, ge=0.0) + top_k: int | None = Field(default=None, gt=0) + log_prob_min_clamp: float | None = None + loss_max_clamp: float | None = Field(default=None, gt=0.0) + + @model_validator(mode="after") + def validate_loss_options(self) -> Self: + if self.loss_mode in ("forward", "reverse", "forward_kl_topk"): + if self.use_policy_gradient: + raise ValueError(f"{self.loss_mode} only supports direct backpropagation") + if self.top_k is None: + raise ValueError(f"{self.loss_mode} requires top_k") + elif self.top_k is not None: + raise ValueError(f"top_k is not used by loss_mode={self.loss_mode}") + + if self.loss_mode == "k1" and not self.use_policy_gradient: + raise ValueError("k1 only supports the policy-gradient path") + return self + + @property + def uses_sampled_token_targets(self) -> bool: + return self.loss_mode in ("k1", "k3") + + @property + def uses_topk_targets(self) -> bool: + return not self.uses_sampled_token_targets + + @property + def loss_ctx_cls(self) -> type["DistillationLossContext"]: + return DistillationLossContext + + @property + def _loss_kwargs_cls(self) -> type["DistillationLossKwargs"]: + return DistillationLossKwargs + + def build( + self, + data: dict, + sp_mesh: DeviceMesh | None = None, + ) -> "DistillationLossContext | None": + if "shifted_labels" not in data or "advantages" not in data: + return None + + loss_kwargs = DistillationLossKwargs( + shifted_labels=data["shifted_labels"], + advantages=data["advantages"], + rollout_logprobs=data.get("rollout_logprobs"), + old_logprobs=data.get("old_logprobs"), + ref_logprobs=data.get("ref_logprobs"), + is_weights=data.get("rollout_is_weights"), + teacher_logprobs=data.get("teacher_logprobs"), + target_token_ids=data.get("target_token_ids"), + ).to(DEVICE) + if sp_mesh is not None and sp_mesh.size() > 1: + loss_kwargs = loss_kwargs.sp_split(sp_mesh) + return self.loss_ctx_cls(self, loss_kwargs) + + +class DistillationLossKwargs(BaseRLLossKwargs): + teacher_logprobs: torch.Tensor | None = None + target_token_ids: torch.Tensor | None = None + distillation_loss_weight: torch.Tensor | None = None + + def sp_split(self, sp_mesh: DeviceMesh) -> Self: + super().sp_split(sp_mesh) + if self.teacher_logprobs is not None: + self.teacher_logprobs = sp_split( + self.teacher_logprobs, + sp_mesh=sp_mesh, + split_dim=1, + padding_value=0.0, + ) + if self.target_token_ids is not None: + self.target_token_ids = sp_split( + self.target_token_ids, + sp_mesh=sp_mesh, + split_dim=1, + padding_value=0, + ) + return self + + def to(self, device: torch.device | str) -> Self: + super().to(device) + if self.teacher_logprobs is not None: + self.teacher_logprobs = self.teacher_logprobs.to(device) + if self.target_token_ids is not None: + self.target_token_ids = self.target_token_ids.to(device) + if self.distillation_loss_weight is not None: + self.distillation_loss_weight = self.distillation_loss_weight.to(device) + return self + + +class DistillationLossContext(BaseRLLossContext): + loss_cfg: DistillationLossConfig + loss_kwargs: DistillationLossKwargs + + def __init__(self, loss_cfg: DistillationLossConfig, loss_kwargs: DistillationLossKwargs): + super().__init__(loss_cfg, loss_kwargs) + self.policy_loss_fn = get_policy_loss_fn(self.loss_cfg.policy_loss_cfg.get("loss_type", "vanilla")) + + @staticmethod + def build_batches( # type: ignore[override] + loss_ctx_list: list["DistillationLossContext"], + ) -> list["DistillationLossContext"]: + assert loss_ctx_list, "loss_ctx_list can not be empty" + loss_cfg = loss_ctx_list[0].loss_cfg + shifted_labels_list = [loss_ctx.loss_kwargs.shifted_labels for loss_ctx in loss_ctx_list] + rank_grad_tokens = sum((labels != loss_cfg.ignore_idx).sum() for labels in shifted_labels_list) + global_grad_tokens = cast(torch.Tensor, rank_grad_tokens) + if dist.is_initialized(): + dist.all_reduce(global_grad_tokens, op=dist.ReduceOp.SUM) + if global_grad_tokens == 0: + logger.warning("Global gradient tokens is 0; using one as the loss denominator") + global_grad_tokens.add_(1) + + for loss_ctx in loss_ctx_list: + loss_kwargs = loss_ctx.loss_kwargs + shifted_labels = loss_kwargs.shifted_labels + assert loss_kwargs.old_logprobs is not None, "old_logprobs can not be None" + + policy_loss_weight = torch.ones_like(shifted_labels, dtype=torch.float32) / global_grad_tokens + policy_loss_weight[shifted_labels == loss_cfg.ignore_idx] = 0.0 + if loss_kwargs.is_weights is not None: + policy_loss_weight = policy_loss_weight * loss_kwargs.is_weights + + if loss_cfg.use_kl_loss: + assert loss_kwargs.ref_logprobs is not None, "ref_logprobs can not be None" + kl_loss_weight = compute_kl_loss_weight( + shifted_labels, + global_grad_tokens, + loss_cfg.kl_loss_coef, + loss_cfg.ignore_idx, + ) + else: + kl_loss_weight = None + + distillation_loss_weight = None + if not loss_cfg.use_policy_gradient: + distillation_loss_weight = ( + torch.ones_like(shifted_labels, dtype=torch.float32) + / global_grad_tokens + * loss_cfg.distillation_loss_weight + ) + distillation_loss_weight[shifted_labels == loss_cfg.ignore_idx] = 0.0 + + loss_kwargs.policy_loss_weight = policy_loss_weight + loss_kwargs.kl_loss_weight = kl_loss_weight + loss_kwargs.distillation_loss_weight = distillation_loss_weight + loss_kwargs.global_grad_tokens = global_grad_tokens + return loss_ctx_list + + def _compute_sampled_distillation_loss( + self, + student_logprobs: torch.Tensor, + teacher_logprobs: torch.Tensor, + ) -> torch.Tensor: + if self.loss_cfg.loss_mode == "k1": + loss = student_logprobs - teacher_logprobs + else: + log_ratio = (teacher_logprobs - student_logprobs).clamp(min=-20.0, max=20.0) + loss = torch.exp(log_ratio) - log_ratio - 1.0 + if self.loss_cfg.loss_max_clamp is not None: + loss = loss.clamp( + min=-self.loss_cfg.loss_max_clamp, + max=self.loss_cfg.loss_max_clamp, + ) + return loss + + def _compute_topk_distillation_loss( + self, + logits: torch.Tensor, + loss_kwargs: DistillationLossKwargs, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + target_token_ids = cast(torch.Tensor, loss_kwargs.target_token_ids) + teacher_logprobs = cast(torch.Tensor, loss_kwargs.teacher_logprobs) + student_topk_logprobs = torch.gather(logits, dim=-1, index=target_token_ids) + student_topk_logprobs = student_topk_logprobs - torch.logsumexp(logits, dim=-1, keepdim=True) + + loss_mode = cast(TopKDistillationMode, self.loss_cfg.loss_mode) + + loss, student_mass, teacher_mass = compute_topk_distillation_kl( + student_topk_logprobs, + teacher_logprobs, + loss_mode, + self.loss_cfg.log_prob_min_clamp, + self.loss_cfg.loss_max_clamp, + ) + student_topk_ids = torch.topk(logits.detach(), k=target_token_ids.size(-1), dim=-1).indices + overlap = (student_topk_ids.unsqueeze(-1) == target_token_ids.detach().unsqueeze(-2)).any(dim=-1) + return loss, student_mass, teacher_mass, overlap.float().mean(dim=-1) + + def loss_fn( + self, + hidden_states: torch.Tensor, + head_weight: torch.Tensor, + head_bias: torch.Tensor | None, + loss_kwargs: DistillationLossKwargs, + ) -> tuple[torch.Tensor, tuple[torch.Tensor | None, dict[str, Any]]]: + logits = F.linear(hidden_states, head_weight, head_bias).float() + shifted_labels = loss_kwargs.shifted_labels + old_logprobs = cast(torch.Tensor, loss_kwargs.old_logprobs) + policy_loss_weight = cast(torch.Tensor, loss_kwargs.policy_loss_weight) + current_logprobs = gather_logprobs(logits, shifted_labels) + + topk_metrics: dict[str, torch.Tensor] = {} + if self.loss_cfg.uses_sampled_token_targets: + teacher_logprobs = cast(torch.Tensor, loss_kwargs.teacher_logprobs) + distillation_student_logprobs = old_logprobs if self.loss_cfg.use_policy_gradient else current_logprobs + per_token_distillation_loss = self._compute_sampled_distillation_loss( + distillation_student_logprobs, + teacher_logprobs, + ) + else: + ( + per_token_distillation_loss, + student_selected_mass, + teacher_selected_mass, + overlap_fraction, + ) = self._compute_topk_distillation_loss(logits, loss_kwargs) + topk_metrics = { + "reduced_topk_opd_student_selected_mass_sum": student_selected_mass.detach(), + "reduced_topk_opd_teacher_selected_mass_sum": teacher_selected_mass.detach(), + "reduced_topk_opd_overlap_fraction_sum": overlap_fraction.detach(), + } + + if self.loss_cfg.use_policy_gradient: + combined_advantages = ( + self.loss_cfg.task_adv_weight * loss_kwargs.advantages + - self.loss_cfg.distillation_loss_weight * per_token_distillation_loss.detach() + ) + loss = self.policy_loss_fn( + current_logprobs, + old_logprobs, + combined_advantages, + policy_loss_weight, + self.loss_cfg.policy_loss_cfg, + ) + else: + distillation_loss_weight = cast(torch.Tensor, loss_kwargs.distillation_loss_weight) + task_advantages = self.loss_cfg.task_adv_weight * loss_kwargs.advantages + task_loss = self.policy_loss_fn( + current_logprobs, + old_logprobs, + task_advantages, + policy_loss_weight, + self.loss_cfg.policy_loss_cfg, + ) + distillation_loss = (per_token_distillation_loss * distillation_loss_weight).sum() + loss = task_loss + distillation_loss + if self.loss_cfg.uses_topk_targets: + topk_metrics["reduced_topk_opd_loss_sum"] = distillation_loss.detach() + + valid_mask = shifted_labels != self.loss_cfg.ignore_idx + valid_float = valid_mask.float() + log_ratio = current_logprobs.detach() - old_logprobs.detach() + log_ratio_safe = torch.clamp(log_ratio, min=-20.0, max=20.0) + ratio = torch.exp(log_ratio_safe) + ratio_max = ratio.masked_fill(~valid_mask, 0.0).max() + ratio_min = ratio.masked_fill(~valid_mask, float("inf")).min() + extra_info = { + "max_ratio": ratio_max, + "reduced_train_policy_ratio_abs_dev_sum": ((ratio - 1.0).abs() * valid_float).sum(), + "reduced_train_policy_kl1_sum": (-log_ratio * valid_float).sum(), + "reduced_train_policy_kl3_sum": ((ratio - 1.0 - log_ratio_safe) * valid_float).sum(), + "reduced_train_policy_valid_count": valid_float.sum(), + "reduced_train_policy_ratio_max": ratio_max, + "reduced_train_policy_ratio_min": ratio_min, + "reduced_distillation_kl_sum": (per_token_distillation_loss.detach() * valid_float).sum(), + "reduced_distillation_abs_loss_sum": (per_token_distillation_loss.detach().abs() * valid_float).sum(), + "reduced_distillation_valid_count": valid_float.sum(), + **{ + key: (value * valid_float).sum() if key != "reduced_topk_opd_loss_sum" else value + for key, value in topk_metrics.items() + }, + } + if self.loss_cfg.loss_mode == "k1": + extra_info["reduced_opd_reverse_kl_sum"] = (per_token_distillation_loss.detach() * valid_float).sum() + extra_info["reduced_opd_abs_logprob_loss_sum"] = ( + per_token_distillation_loss.detach().abs() * valid_float + ).sum() + if self.loss_cfg.uses_topk_targets: + extra_info["reduced_topk_opd_kl_sum"] = (per_token_distillation_loss.detach() * valid_float).sum() + extra_info["reduced_topk_opd_valid_count"] = valid_float.sum() + + cliprange_low = self.loss_cfg.policy_loss_cfg.get("cliprange_low") + cliprange_high = self.loss_cfg.policy_loss_cfg.get("cliprange_high") + if cliprange_low is not None and cliprange_high is not None: + extra_info["reduced_train_policy_clip_low_count"] = ( + ((ratio < 1 - cliprange_low) & valid_mask).float().sum() + ) + extra_info["reduced_train_policy_clip_high_count"] = ( + ((ratio > 1 + cliprange_high) & valid_mask).float().sum() + ) + + if self.loss_cfg.use_kl_loss: + ref_logprobs = loss_kwargs.ref_logprobs + kl_loss_weight = loss_kwargs.kl_loss_weight + assert ref_logprobs is not None and kl_loss_weight is not None + loss = loss + kl_penalty( + current_logprobs, + ref_logprobs, + kl_loss_weight, + self.loss_cfg.kl_loss_type, + ) + + return loss, (logits, extra_info) + + +def finalize_distillation_metrics( + extra_info_dict: dict[str, Any], + device: str | torch.device, +) -> dict[str, Any]: + def reduce_values(keys: tuple[str, ...]) -> dict[str, float]: + values = torch.tensor( + [extra_info_dict.pop(key, 0.0) for key in keys], + dtype=torch.float32, + device=device, + ) + if dist.is_initialized(): + dist.all_reduce(values, op=dist.ReduceOp.SUM) + return dict(zip(keys, values.tolist())) + + if "reduced_distillation_valid_count" in extra_info_dict: + distillation_keys = [ + "reduced_distillation_kl_sum", + "reduced_distillation_abs_loss_sum", + "reduced_distillation_valid_count", + ] + has_opd_metrics = "reduced_opd_reverse_kl_sum" in extra_info_dict + if has_opd_metrics: + distillation_keys.extend( + [ + "reduced_opd_reverse_kl_sum", + "reduced_opd_abs_logprob_loss_sum", + ] + ) + values = reduce_values(tuple(distillation_keys)) + valid_count = values["reduced_distillation_valid_count"] + extra_info_dict["reduced_distillation_kl"] = ( + values["reduced_distillation_kl_sum"] / valid_count if valid_count > 0 else 0.0 + ) + extra_info_dict["reduced_distillation_abs_loss"] = ( + values["reduced_distillation_abs_loss_sum"] / valid_count if valid_count > 0 else 0.0 + ) + if has_opd_metrics: + extra_info_dict["opd_reverse_kl"] = ( + values["reduced_opd_reverse_kl_sum"] / valid_count if valid_count > 0 else 0.0 + ) + extra_info_dict["opd_abs_logprob_loss"] = ( + values["reduced_opd_abs_logprob_loss_sum"] / valid_count if valid_count > 0 else 0.0 + ) + + if "reduced_topk_opd_valid_count" in extra_info_dict: + topk_keys = ( + "reduced_topk_opd_kl_sum", + "reduced_topk_opd_loss_sum", + "reduced_topk_opd_student_selected_mass_sum", + "reduced_topk_opd_teacher_selected_mass_sum", + "reduced_topk_opd_overlap_fraction_sum", + "reduced_topk_opd_valid_count", + ) + values = reduce_values(topk_keys) + valid_count = values["reduced_topk_opd_valid_count"] + for output_key, sum_key in { + "reduced_topk_opd_kl": "reduced_topk_opd_kl_sum", + "reduced_topk_opd_student_selected_mass": "reduced_topk_opd_student_selected_mass_sum", + "reduced_topk_opd_teacher_selected_mass": "reduced_topk_opd_teacher_selected_mass_sum", + "reduced_topk_opd_overlap_fraction": "reduced_topk_opd_overlap_fraction_sum", + }.items(): + extra_info_dict[output_key] = values[sum_key] / valid_count if valid_count > 0 else 0.0 + extra_info_dict["reduced_topk_opd_loss"] = values["reduced_topk_opd_loss_sum"] + return extra_info_dict diff --git a/xtuner/v1/rl/model_utils.py b/xtuner/v1/rl/model_utils.py new file mode 100644 index 0000000000..404569d956 --- /dev/null +++ b/xtuner/v1/rl/model_utils.py @@ -0,0 +1,54 @@ +from pathlib import Path +from typing import cast + +import torch + +from xtuner.v1.config.fsdp import FSDPConfig +from xtuner.v1.float8.float8_handler import Float8Handler +from xtuner.v1.model.base import BaseModel as XtunerBaseModel +from xtuner.v1.model.base import TransformerConfig +from xtuner.v1.model.compose.base import BaseComposeConfig, BaseComposeModel +from xtuner.v1.utils import get_torch_device_module + + +DEVICE_MODULE = get_torch_device_module() + +FrozenModel = BaseComposeModel | XtunerBaseModel + + +def build_frozen_model( + model_cfg: TransformerConfig | BaseComposeConfig, + load_from: str | Path, + fsdp_cfg: FSDPConfig | None = None, +) -> FrozenModel: + """Build a frozen FSDP model and leave it resident on CPU.""" + with torch.device("meta"): + model = model_cfg.build() + + if isinstance(model_cfg, BaseComposeConfig): + assert model_cfg.text_config.float8_cfg is None, "BaseComposeConfig does not support float8" + if fsdp_cfg is None: + fsdp_cfg = FSDPConfig(recompute_ratio=0, cpu_offload=False, requires_grad=False) + model = model.fully_shard(fsdp_cfg) + model.from_hf(hf_path=load_from) + model.eval() # type: ignore + else: + model_cfg = cast(TransformerConfig, model_cfg) + if model_cfg.float8_cfg is not None and model_cfg.float8_cfg.enable_float8: + float8_handler = Float8Handler( + scaling_granularity_gemm=model_cfg.float8_cfg.scaling_granularity_gemm, + scaling_granularity_grouped_gemm=model_cfg.float8_cfg.scaling_granularity_grouped_gemm, + ) + else: + float8_handler = None + if fsdp_cfg is None: + fsdp_cfg = FSDPConfig(recompute_ratio=0, cpu_offload=False, requires_grad=False) + model = model.fully_shard(fsdp_cfg) # type: ignore + model.from_hf(hf_path=load_from) + model.eval() # type: ignore + if float8_handler is not None: + float8_handler.precompute_float8_dynamic_scale_for_fsdp(model) # type: ignore + + model.to_device("cpu") # type: ignore + DEVICE_MODULE.empty_cache() + return model diff --git a/xtuner/v1/rl/trainer/controller.py b/xtuner/v1/rl/trainer/controller.py index 0cc6b9e322..7484a542c4 100644 --- a/xtuner/v1/rl/trainer/controller.py +++ b/xtuner/v1/rl/trainer/controller.py @@ -4,6 +4,7 @@ import ray import torch +from typing_extensions import NotRequired from xtuner.v1.data_proto.sequence_context import SequenceContext from xtuner.v1.model.compose.base import BaseComposeConfig @@ -22,6 +23,8 @@ class ColateItem(TypedDict): shifted_labels: torch.Tensor advantage: float rollout_logprobs: torch.Tensor | None + teacher_logprobs: NotRequired[torch.Tensor | None] + teacher_indices: NotRequired[torch.Tensor] def _summarize_process_group_results(results: list[dict[str, Any]]) -> str: @@ -121,6 +124,14 @@ def _packing(self, data_batches, pack_max_length, language_cfg): if "rollout_logprobs" in data_batches[0] and data_batches[0]["rollout_logprobs"] is not None: rollout_logprobs_list = [data_batches[i]["rollout_logprobs"] for i in indices] + teacher_logprobs_list = None + if "teacher_logprobs" in data_batches[0] and data_batches[0]["teacher_logprobs"] is not None: + teacher_logprobs_list = [data_batches[i]["teacher_logprobs"] for i in indices] + + teacher_indices_list = None + if "teacher_indices" in data_batches[0] and data_batches[0]["teacher_indices"] is not None: + teacher_indices_list = [data_batches[i]["teacher_indices"] for i in indices] + if pad_len > 0: # Reduce the attn calculation time by using multiple short sequence packs pad_tokens = tuple( @@ -162,6 +173,22 @@ def _packing(self, data_batches, pack_max_length, language_cfg): device=data_batches[0]["shifted_labels"].device, ) rollout_logprobs_list.append(pad_rollout_logprobs) + if teacher_logprobs_list is not None: + pad_teacher_logprobs = torch.zeros( + 1, + pad_len, + dtype=data_batches[0]["teacher_logprobs"].dtype, + device=data_batches[0]["shifted_labels"].device, + ) + teacher_logprobs_list.append(pad_teacher_logprobs) + if teacher_indices_list is not None: + pad_teacher_indices = torch.full( + (1, pad_len), + -1, + dtype=data_batches[0]["teacher_indices"].dtype, + device=data_batches[0]["teacher_indices"].device, + ) + teacher_indices_list.append(pad_teacher_indices) seq_ctx = SequenceContext.cat(seq_ctx_list) shifted_labels = torch.cat(label_list, dim=1) # (1, max_len) @@ -172,12 +199,22 @@ def _packing(self, data_batches, pack_max_length, language_cfg): if rollout_logprobs_list is not None: rollout_logprobs = torch.cat(rollout_logprobs_list, dim=1) # (1, max_len) + teacher_logprobs = None + if teacher_logprobs_list is not None: + teacher_logprobs = torch.cat(teacher_logprobs_list, dim=1) # (1, max_len) + + teacher_indices = None + if teacher_indices_list is not None: + teacher_indices = torch.cat(teacher_indices_list, dim=1) # (1, max_len) + packed_data_batches.append( { "seq_ctx": seq_ctx, "shifted_labels": shifted_labels, "advantages": advantages, "rollout_logprobs": rollout_logprobs, + "teacher_logprobs": teacher_logprobs, + "teacher_indices": teacher_indices, } ) return packed_data_batches @@ -260,11 +297,26 @@ def fit(self, data_batches: list[ColateItem], pack_max_length: int, rollout_idx: pad_rollout_logprobs = torch.zeros( 1, pack_max_length, dtype=packed_data_batches[0]["rollout_logprobs"].dtype, device="cpu" ) + pad_teacher_logprobs = None + if "teacher_logprobs" in packed_data_batches[0] and packed_data_batches[0]["teacher_logprobs"] is not None: + pad_teacher_logprobs = torch.zeros( + 1, pack_max_length, dtype=packed_data_batches[0]["teacher_logprobs"].dtype, device="cpu" + ) + pad_teacher_indices = None + if "teacher_indices" in packed_data_batches[0] and packed_data_batches[0]["teacher_indices"] is not None: + pad_teacher_indices = torch.full( + (1, pack_max_length), + -1, + dtype=packed_data_batches[0]["teacher_indices"].dtype, + device="cpu", + ) pad_data = { "seq_ctx": pad_seq_ctx, "shifted_labels": pad_shifted_labels, "advantages": pad_advantages, "rollout_logprobs": pad_rollout_logprobs, + "teacher_logprobs": pad_teacher_logprobs, + "teacher_indices": pad_teacher_indices, } pad_data_samples = [pad_data for _ in range(pad_num)] packed_data_batches = packed_data_batches + pad_data_samples diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index 1cb2cbe65d..a358893329 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -35,16 +35,24 @@ from xtuner.v1.datasets.config import DataloaderConfig from xtuner.v1.datasets.dataloader import Dataloader from xtuner.v1.engine.train_engine import TrainEngine, TrainStepInfo -from xtuner.v1.float8.float8_handler import Float8Handler from xtuner.v1.loss import BaseLossContext, CELossConfig, LogProbConfig from xtuner.v1.loss.ce_loss import CELossContext, LMHeadLossContext from xtuner.v1.loss.mtp_loss import MTPLossConfig, MTPLossContext -from xtuner.v1.model.base import BaseModel as XtunerBaseModel +from xtuner.v1.loss.utils import sp_split from xtuner.v1.model.base import ModelItem, TransformerConfig -from xtuner.v1.model.compose.base import BaseComposeConfig, BaseComposeModel +from xtuner.v1.model.compose.base import BaseComposeConfig from xtuner.v1.model.utils.misc import ModelForwardExtraLogInfo from xtuner.v1.profiler import profiling_memory, profiling_time -from xtuner.v1.rl.loss import BaseRLLossConfig, BaseRLLossContext, finalize_train_policy_metrics, kl_penalty +from xtuner.v1.rl.distillation import DistillationConfig, TrainTeacherManager, TrainTeacherTimings +from xtuner.v1.rl.loss import ( + BaseRLLossConfig, + BaseRLLossContext, + DistillationLossConfig, + finalize_distillation_metrics, + finalize_train_policy_metrics, + kl_penalty, +) +from xtuner.v1.rl.model_utils import build_frozen_model from xtuner.v1.rl.utils import SingleAcceleratorWorker from xtuner.v1.rl.weight_update import WeightUpdater from xtuner.v1.train.trainer import LoadCheckpointConfig @@ -157,6 +165,7 @@ class WorkerConfig(BaseModel): free_rollout_routed_experts_in_worker: bool = True # 默认不需要用户配置 offload_rollout_routed_experts: bool = False offload_old_logprobs: bool = False + distillation_config: DistillationConfig | None = None # sft config sft_dataloader_cfg: DataloaderConfig | None = None @@ -190,6 +199,8 @@ class WorkerInputItem(TypedDict): shifted_labels: torch.LongTensor advantages: torch.Tensor rollout_logprobs: torch.Tensor | None + teacher_logprobs: NotRequired[torch.Tensor | None] + teacher_indices: NotRequired[torch.Tensor] class WorkerTrainLogItem(TypedDict, total=False): @@ -203,6 +214,9 @@ class WorkerTrainLogItem(TypedDict, total=False): class WorkerLogItem(TypedDict): train_entropy: float + teacher_compute_time: NotRequired[float] + teacher_onload_time: NotRequired[float] + teacher_offload_time: NotRequired[float] rollout_entropy: NotRequired[float] mismatch_metrics: NotRequired[dict[str, float]] rollout_is_metrics: NotRequired[dict[str, float]] @@ -255,10 +269,18 @@ def __init__( self._has_ref = True if worker_cfg.ref_load_from is None: worker_cfg.ref_load_from = worker_cfg.load_from - self._ref_model = self._build_ref_model( + self._ref_model = build_frozen_model( worker_cfg.model_cfg, worker_cfg.ref_load_from, worker_cfg.ref_model_fsdp_cfg ) + self._train_teacher_manager: TrainTeacherManager | None = None + distillation_config = worker_cfg.distillation_config + if distillation_config is not None and distillation_config.train_teachers: + self._train_teacher_manager = TrainTeacherManager( + distillation_config, + chunk_size=worker_cfg.loss_cfg.chunk_size, + ) + self._optimizer_steps = worker_cfg.optimizer_steps profile_step = worker_cfg.profile_step if isinstance(profile_step, int): @@ -394,46 +416,6 @@ def _build_engine(self, worker_cfg: WorkerConfig) -> TrainEngine: self.logger.info(f"The `compile_cfg` of model is {json.dumps(engine.model.compile_cfg, indent=4)}") return engine - def _build_ref_model( - self, - ref_model_cfg: TransformerConfig | BaseComposeConfig, - load_from: str | Path, - ref_model_fsdp_cfg: FSDPConfig | None = None, - ): - # TODO: 需要重构，使得能更优雅的兼容 mllm - model: BaseComposeModel | XtunerBaseModel - with torch.device("meta"): - model = ref_model_cfg.build() - - if isinstance(ref_model_cfg, BaseComposeConfig): - assert ref_model_cfg.text_config.float8_cfg is None, "BaseComposeConfig does not support float8" - if ref_model_fsdp_cfg is None: - ref_model_fsdp_cfg = FSDPConfig(recompute_ratio=0, cpu_offload=False, requires_grad=False) - model = model.fully_shard(ref_model_fsdp_cfg) - model.from_hf(hf_path=load_from) - model.eval() # type: ignore - else: - ref_model_cfg = cast(TransformerConfig, ref_model_cfg) - if ref_model_cfg.float8_cfg is not None and ref_model_cfg.float8_cfg.enable_float8: - float8_handler = Float8Handler( - scaling_granularity_gemm=ref_model_cfg.float8_cfg.scaling_granularity_gemm, - scaling_granularity_grouped_gemm=ref_model_cfg.float8_cfg.scaling_granularity_grouped_gemm, - ) - else: - float8_handler = None - if ref_model_fsdp_cfg is None: - ref_model_fsdp_cfg = FSDPConfig(recompute_ratio=0, cpu_offload=False, requires_grad=False) - model = model.fully_shard(ref_model_fsdp_cfg) # type: ignore - - model.from_hf(hf_path=load_from) - model.eval() # type: ignore - if float8_handler is not None: - # As the ref model is not updated, we only compute params' scales once - float8_handler.precompute_float8_dynamic_scale_for_fsdp(model) # type: ignore - model.to_device("cpu") # type: ignore - DEVICE_MODULE.empty_cache() - return model - def _init_data_mesh( self, sp_size: int, @@ -514,6 +496,44 @@ def compute_ref_logprobs( self._ref_model.to_device("cpu") return ref_logprobs_list + def _compute_train_teacher_outputs( + self, + seq_ctx_list: list[SequenceContext], + teacher_indices_list: list[torch.Tensor], + loss_ctx_list: list[BaseRLLossContext], + ) -> TrainTeacherTimings: + self._engine.put_model_to_device("cpu") + self._engine.put_optimizer_to_device("cpu") + if hasattr(DEVICE_MODULE, "empty_cache"): + DEVICE_MODULE.empty_cache() + try: + assert self._train_teacher_manager is not None + outputs = self._train_teacher_manager.compute_logprobs( + seq_ctx_list=seq_ctx_list, + shifted_labels_list=[loss_ctx.loss_kwargs.shifted_labels for loss_ctx in loss_ctx_list], + teacher_indices_list=teacher_indices_list, + ) + if outputs.target_token_ids is None: + for loss_ctx, teacher_logprobs in zip(loss_ctx_list, outputs.teacher_logprobs): + loss_ctx.loss_kwargs.teacher_logprobs = teacher_logprobs + else: + for loss_ctx, token_ids, teacher_logprobs in zip( + loss_ctx_list, + outputs.target_token_ids, + outputs.teacher_logprobs, + ): + loss_ctx.loss_kwargs.target_token_ids = token_ids + loss_ctx.loss_kwargs.teacher_logprobs = teacher_logprobs + return outputs.timings + finally: + if self._train_teacher_manager is not None: + self._train_teacher_manager.offload_all_to_cpu() + self._onload_actor_and_optimizer() + + def _onload_actor_and_optimizer(self) -> None: + self._engine.put_model_to_device(DEVICE) + self._engine.put_optimizer_to_device(DEVICE) + def _add_rollout_routed_experts( self, seq_ctx: SequenceContext, rollout_routed_experts: torch.Tensor | list[torch.Tensor | ray.ObjectRef] ): @@ -623,6 +643,8 @@ def _maybe_profiling(self, global_train_step: int, phase: str): def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLogItem: # NOTE: sglang会清除logger handle, 重新创建 self.logger = get_logger(log_dir=self.log_dir, tag="TrainingWorker") + if self._train_teacher_manager is None: + self._onload_actor_and_optimizer() loss_cfg: BaseRLLossConfig = self.config.loss_cfg num_batches = len(data_batches) iters_per_step = math.ceil(num_batches / self._optimizer_steps) @@ -635,11 +657,13 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo # Init loss_ctx: shifted_labels, advantages, rollout_logprobs seq_ctx_list: list[SequenceContext] = [] loss_ctx_list: list[BaseRLLossContext] = [] + teacher_indices_list: list[torch.Tensor] = [] mtp_loss_ctx_list: list[list[MTPLossContext]] = [] prepare_inputs_begin = time.perf_counter() for data in data_batches: # update seq_ctx seq_ctx = data["seq_ctx"] + teacher_indices = data.get("teacher_indices") pixel_values = seq_ctx.pixel_values if pixel_values is not None: if not isinstance(pixel_values, np.ndarray): @@ -667,11 +691,22 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo advantages = data["advantages"].to(DEVICE) rollout_logprobs = data.get("rollout_logprobs", None) rollout_logprobs = rollout_logprobs.to(DEVICE) if rollout_logprobs is not None else None + if teacher_indices is not None: + teacher_indices = teacher_indices.to(DEVICE) + if self.sp_mesh.size() > 1: + teacher_indices = sp_split( + teacher_indices, + sp_mesh=self.sp_mesh, + split_dim=1, + padding_value=-1, + ) + teacher_indices_list.append(teacher_indices) loss_ctx = loss_cfg.build( data={ "shifted_labels": shifted_labels, "advantages": advantages, "rollout_logprobs": rollout_logprobs, + "teacher_logprobs": data.get("teacher_logprobs", None), }, sp_mesh=self.sp_mesh, ) @@ -709,12 +744,25 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo shifted_labels_list = [loss_ctx.loss_kwargs.shifted_labels for loss_ctx in loss_ctx_list] rollout_logprobs_list = [loss_ctx.loss_kwargs.rollout_logprobs for loss_ctx in loss_ctx_list] + worker_log_item: WorkerLogItem = {"train_entropy": 0.0, "train_metrics": [], "sft_train_metrics": {}} + if self._train_teacher_manager is not None: + # Training-side Teachers share the training workers' devices. Keep + # Actor and optimizer on CPU while each frozen Teacher produces its + # targets, then restore Actor state for old-logprob and train forward. + teacher_timings = self._compute_train_teacher_outputs( + seq_ctx_list, + teacher_indices_list, + loss_ctx_list, + ) + worker_log_item["teacher_compute_time"] = teacher_timings.compute + worker_log_item["teacher_onload_time"] = teacher_timings.onload + worker_log_item["teacher_offload_time"] = teacher_timings.offload + # compute old logprobs old_logprobs_list = self.compute_actor_logprobs(seq_ctx_list, shifted_labels_list) for old_logprobs, loss_ctx in zip(old_logprobs_list, loss_ctx_list): loss_ctx.loss_kwargs.old_logprobs = old_logprobs - worker_log_item: WorkerLogItem = {"train_entropy": 0.0, "train_metrics": [], "sft_train_metrics": {}} logger_msg = f"Rollout {rollout_idx}: " # compute entropy @@ -908,6 +956,8 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo if isinstance(v, (torch.Tensor, int, float)) } extra_info_dict = finalize_train_policy_metrics(extra_info_dict, DEVICE) + if isinstance(loss_cfg, DistillationLossConfig): + extra_info_dict = finalize_distillation_metrics(extra_info_dict, DEVICE) train_step_info.pop("total_loss") # type: ignore[misc] max_memory = DEVICE_MODULE.max_memory_allocated() / (1024**3) # type: ignore[attr-defined] reserved_memory = DEVICE_MODULE.max_memory_reserved() / (1024**3) # type: ignore[attr-defined] diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 49dee2b650..b8dcda7c4f 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -32,6 +32,7 @@ ProduceBatchStatus, ) from xtuner.v1.rl.agent_loop_manager.produce_utils import default_should_continue_fn +from xtuner.v1.rl.distillation import DistillationConfig from xtuner.v1.rl.evaluator import EvaluatorConfig from xtuner.v1.rl.health_manager import RLHealthManager, _NoOpRLHealthManager from xtuner.v1.rl.replay_buffer import ( @@ -67,6 +68,13 @@ DEVICE = get_device() DEVICE_MODULE = get_torch_device_module() +_DISTILLATION_METRIC_PREFIXES = ( + "reduced_distillation_", + "reduced_topk_opd_", + "opd_", + "topk_opd_", +) + def _to_cpu_tensor(value: np.ndarray | None, *, dtype: torch.dtype | None = None) -> torch.Tensor | None: if value is None: @@ -412,6 +420,7 @@ class BaseRLTrainerConfig(BaseModel): total_epochs: int | None = None train_batch_size: int advantage_estimator_config: BaseAdvantageConfig = Field(default_factory=GRPOAdvantageConfig) + distillation_config: DistillationConfig | None = None sync_weights_interval: int = 1 enable_evaluate: bool = True @@ -449,6 +458,10 @@ def _validate_sync_intervals(self): raise ValueError(f"total_train_steps must be positive, got {self.total_train_steps}.") if self.total_epochs is not None and self.total_epochs <= 0: raise ValueError(f"total_epochs must be positive, got {self.total_epochs}.") + if self.distillation_config is not None: + if self.train_worker_cfg.loss_cfg != self.distillation_config.loss_config: + raise ValueError("train_worker_cfg.loss_cfg must be distillation_config.loss_config") + self.distillation_config.validate_student_model(self.train_worker_cfg.model_cfg) _validate_sync_intervals( sync_weights_interval=self.sync_weights_interval, checkpoint_interval=self.checkpoint_interval, @@ -652,6 +665,25 @@ class BaseRLTrainer: _debug_train_files: dict[int, Path] def _init_common(self, cfg: BaseRLTrainerConfig, *, meta_path: str, logger_tag: str) -> None: + if cfg.distillation_config is not None and cfg.distillation_config.rollout_teachers: + endpoint_map = json.loads(os.environ.get("XTUNER_OPD_TEACHER_ENDPOINTS_JSON", "{}")) + cfg.distillation_config = cfg.distillation_config.resolve_teacher_endpoints(endpoint_map) + + self._distillation_config = cfg.distillation_config + self._distillation_loss_cfg = ( + self._distillation_config.loss_config if self._distillation_config is not None else None + ) + self._train_teacher_config = ( + self._distillation_config + if self._distillation_config is not None and self._distillation_config.train_teachers + else None + ) + self._rollout_teacher_config = ( + self._distillation_config + if self._distillation_config is not None and self._distillation_config.rollout_teachers + else None + ) + check_fa3() self._init_work_dir_and_meta(cfg, meta_path) self._init_load_source(cfg) @@ -755,6 +787,7 @@ def _init_train_worker_config(self, cfg: BaseRLTrainerConfig, log_dir: Path) -> cfg.train_worker_cfg.free_rollout_routed_experts_in_worker = False cfg.train_worker_cfg.load_from = cfg.load_from cfg.train_worker_cfg.log_dir = log_dir + cfg.train_worker_cfg.distillation_config = cfg.distillation_config self._train_worker_cfg = cfg.train_worker_cfg def _init_rollout_config(self, cfg: BaseRLTrainerConfig, log_dir: Path) -> None: @@ -782,12 +815,18 @@ def _init_runtime_flags(self, cfg: BaseRLTrainerConfig) -> None: def _build_agent_loop_components(self, cfg: BaseRLTrainerConfig, replay_buffer) -> None: self.tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer_path, trust_remote_code=True) + # Sampled-token objectives require identity Student sampling even when + # Teacher logprobs are produced later by the training workers. + agent_loop_distillation_config = self._rollout_teacher_config + if self._distillation_loss_cfg is not None and self._distillation_loss_cfg.uses_sampled_token_targets: + agent_loop_distillation_config = self._distillation_config agent_loop_manager = cfg.agent_loop_manager_cfg.build( rollout_controller=self.rollout_controller, tokenizer=self.tokenizer, replay_buffer=replay_buffer, logger=self.logger, sync_weights_interval=cfg.sync_weights_interval, + distillation_config=agent_loop_distillation_config, ) self.agent_loop_manager = cast(AgentLoopManager | DisaggAgentLoopManager, agent_loop_manager) @@ -802,6 +841,7 @@ def _build_agent_loop_components(self, cfg: BaseRLTrainerConfig, replay_buffer) replay_buffer=replay_buffer, logger=self.logger, sync_weights_interval=cfg.sync_weights_interval, + distillation_config=None, ), ) @@ -1007,7 +1047,7 @@ def _train_one_batch( step_timer_dict: dict, *, offload_rollout_before_train: bool = False, - onload_train_before_train: bool = False, + resume_train_before_train: bool = False, raw_rewards_sum: float = 0.0, raw_rewards_count: int = 0, ) -> TrainInfo: @@ -1020,17 +1060,14 @@ def _train_one_batch( self._save_trajectories(train_batch, train_trajectory_path) self.logger.info(f"Train step {train_step} train trajectories saved to {train_trajectory_path}") - # 共卡训练前切换资源：检查 rollout -> offload rollout -> onload train。 + # 共卡训练前切换资源：检查 rollout -> offload rollout -> resume train NCCL。 if offload_rollout_before_train: ray.get(self.rollout_controller.offload.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT) - if onload_train_before_train: + if resume_train_before_train: if getattr(self, "_train_nccl_suspended", False): with timer("resume_train_nccl", step_timer_dict): self.train_controller.resume_train_nccl_process_groups() self._train_nccl_suspended = False - with timer("onload", step_timer_dict): - self.train_controller.onload(target="all") - self.logger.info("Training controller loaded") with timer("prepare_data", step_timer_dict): data_batches, data_info = self._prepare_train_data( @@ -1118,24 +1155,32 @@ def _prepare_train_data( raw_rewards_sum: float = 0.0, raw_rewards_count: int = 0, ): - rewards_list = [] # Per-session rewards for distribution metrics. Agentic sessions may split into several # trainable segments that share one reward; counting that reward once per session keeps - # rewards/* from being weighted by segment count. rewards_list stays per-segment for counts. + # rewards/* from being weighted by segment count. cluster_rewards_list: list[float] = [] + teacher_rewards: dict[str, list[float]] = {} advantages_list = [] prompt_len_list = [] response_len_list = [] tool_turns_list: list[int] = [] training_tokens = 0 + training_samples = 0 data_batches = [] + teacher_index_by_data_source = ( + self._train_teacher_config.teacher_index_by_data_source if self._train_teacher_config is not None else None + ) for j, group in enumerate(data_groups): if not is_valid_for_training(group, self.logger): self.logger.error(f"Skip one data group {group} due to rollout failed or empty response.") continue + training_samples += len(group) + task_adv_weight = ( + self._distillation_loss_cfg.task_adv_weight if self._distillation_loss_cfg is not None else 1.0 + ) prompt_ids = None if any(data.input_ids is None for data in group): is_vlm_model = "train_prompt_ids" in group[0].extra_fields @@ -1147,19 +1192,22 @@ def _prepare_train_data( assert prompt_ids is not None and len(prompt_ids) > 0, ( f"Prompt ids cannot be None or empty in data: {group[0]}" ) - rewards = [] - # Agentic rollouts may split one model session into multiple trainable segments. - # Compute the group advantage once per session, then broadcast it back to each segment. + for data in group: + # 有可能有重复，但是没有其他更好办法 + turns = data.extra_fields.get("agent_tool_turns") + if isinstance(turns, int): + tool_turns_list.append(turns) + # Collect rewards independently from task-advantage computation. Pure OPD may omit + # rewards entirely; when rewards are present they remain useful observability signals. cluster_index_by_key: dict[Any, int] = {} cluster_rewards: list[float] = [] cluster_representatives: list[RolloutState] = [] sample_cluster_indices: list[int] = [] for data in group: - assert data.reward is not None and "score" in data.reward, ( - f"Reward is missing or does not contain 'score' key in data: {data}" - ) + if data.reward is None or "score" not in data.reward: + assert task_adv_weight == 0, f"Reward is missing or does not contain 'score' key in data: {data}" + continue reward = float(data.reward["score"]) - rewards.append(reward) # session_id is only set by agentic loops / XTUNER_DETERMINISTIC; plain RL falls back # to rollout_id, which the sampler always assigns. Segments of one session share a key. cluster_key = data.session_id if data.session_id is not None else data.rollout_id @@ -1170,19 +1218,34 @@ def _prepare_train_data( cluster_rewards.append(reward) cluster_representatives.append(data) sample_cluster_indices.append(cluster_index) - # 有可能有重复，但是没有其他更好办法 - turns = data.extra_fields.get("agent_tool_turns") - if isinstance(turns, int): - tool_turns_list.append(turns) - rewards_list.extend(rewards) cluster_rewards_list.extend(cluster_rewards) - rewards_tensor = torch.tensor(cluster_rewards, dtype=torch.float32) - cluster_advantages = self._advantage_estimator.compute(rewards_tensor, cluster_representatives) - sample_advantages = [cluster_advantages[cluster_index].item() for cluster_index in sample_cluster_indices] + if self._distillation_config is not None: + for reward, representative in zip(cluster_rewards, cluster_representatives): + data_source = representative.extra_fields.get("origin_data_source") + if not isinstance(data_source, str): + continue + teacher_name = self._distillation_config.data_source_teacher_map.get(data_source) + if teacher_name is not None: + teacher_rewards.setdefault(teacher_name, []).append(reward) + + if task_adv_weight == 0: + sample_advantages = [0.0] * len(group) + else: + # Agentic rollouts may split one model session into multiple trainable segments. + # Compute the group advantage once per session, then broadcast it back to each segment. + rewards_tensor = torch.tensor(cluster_rewards, dtype=torch.float32) + cluster_advantages = self._advantage_estimator.compute(rewards_tensor, cluster_representatives) + sample_advantages = [ + cluster_advantages[cluster_index].item() for cluster_index in sample_cluster_indices + ] prompt_repeat_k = len(group) for i in range(prompt_repeat_k): + teacher_index = None + if teacher_index_by_data_source is not None: + data_source = group[i].extra_fields["origin_data_source"] + teacher_index = teacher_index_by_data_source[data_source] if group[i].input_ids is not None: raw_input_ids = cast(list[int], group[i].input_ids) labels = cast(list[int] | None, group[i].labels) @@ -1234,6 +1297,11 @@ def _prepare_train_data( "advantage": actual_advantages, "rollout_logprobs": rollout_logprobs, } + if teacher_index is not None: + data_dict["teacher_indices"] = torch.full_like( + shifted_labels_t, + teacher_index, + ) seq_ctx.rollout_routed_experts = group[i].routed_experts data_batches.append(data_dict) @@ -1284,12 +1352,9 @@ def _prepare_train_data( shifted_labels = [-100] * (len(prompt_ids) - 1) + response_labels shifted_labels_t = torch.tensor(shifted_labels, dtype=torch.int64).unsqueeze(0) - # 根据 response_mask 计算新的 advantages - advatnages_val = sample_advantages[i] - actual_advantages = [advatnages_val] * len(prompt_ids) + [ - 0.0 if mask == 0 else advatnages_val for mask in response_mask - ] - advantages_list.extend(actual_advantages[:-1]) + base_advantage = sample_advantages[i] + actual_advantages = [0.0 if label == -100 else base_advantage for label in shifted_labels] + advantages_list.extend(actual_advantages) assert len(input_ids) <= pack_max_length, f"{len(input_ids)} vs {pack_max_length}" training_tokens += len(input_ids) @@ -1314,6 +1379,21 @@ def _prepare_train_data( "advantage": actual_advantages, "rollout_logprobs": rollout_logprobs, } + if ( + self._distillation_config is not None + and self._distillation_loss_cfg is not None + and self._distillation_config.rollout_teachers + ): + response_teacher_logprobs = cast(list[float], group[i].teacher_logprobs) + data_dict["teacher_logprobs"] = torch.tensor( + [0.0] * (len(prompt_ids) - 1) + response_teacher_logprobs, + dtype=torch.float32, + ).unsqueeze(0) + if teacher_index is not None: + data_dict["teacher_indices"] = torch.full_like( + shifted_labels_t, + teacher_index, + ) seq_ctx.rollout_routed_experts = group[i].routed_experts # n,layer*expert @@ -1321,8 +1401,8 @@ def _prepare_train_data( if not XTUNER_DETERMINISTIC: random.shuffle(data_batches) - # rewards/* report the per-session reward distribution; batch_size/training_samples below - # still use rewards_list (per-segment) so counts reflect the actual training samples. + # rewards/* report the per-session reward distribution; batch_size/training_samples + # count the valid rollout segments included in training. rewards_t = torch.tensor(cluster_rewards_list).float() if cluster_rewards_list else torch.tensor([0.0]).float() advantages_t = torch.tensor(advantages_list).float() if advantages_list else torch.tensor([0.0]).float() prompt_len_t = torch.tensor(prompt_len_list).float() if prompt_len_list else torch.tensor([0.0]).float() @@ -1330,8 +1410,8 @@ def _prepare_train_data( raw_rewards_mean = raw_rewards_sum / raw_rewards_count if raw_rewards_count > 0 else rewards_t.mean().item() info_dict = { - "batch_size": len(rewards_list), - "training_samples": len(rewards_list), + "batch_size": training_samples, + "training_samples": training_samples, "training_tokens": training_tokens, "rewards/mean": rewards_t.mean().item(), "rewards/min": rewards_t.min().item(), @@ -1348,6 +1428,9 @@ def _prepare_train_data( "prompt_len/min": prompt_len_t.min().item(), "prompt_len/max": prompt_len_t.max().item(), } + for teacher_name, rewards in teacher_rewards.items(): + if rewards: + info_dict[f"rewards/{teacher_name}/mean"] = torch.tensor(rewards, dtype=torch.float32).mean().item() if tool_turns_list: tool_turns_t = torch.tensor(tool_turns_list, dtype=torch.float32) info_dict["tool_turns/mean"] = tool_turns_t.mean().item() @@ -1470,6 +1553,12 @@ def _log_step( all_scalars.update({f"{k}": v for k, v in rank0_mismatch_metrics.items()}) all_scalars.update({"entropy/rollout": rank0_rollout_entropy}) all_scalars.update({"entropy/train": rank0_log_item["train_entropy"]}) + if "teacher_compute_time" in rank0_log_item: + all_scalars["time/train_teacher_compute"] = rank0_log_item["teacher_compute_time"] + if "teacher_onload_time" in rank0_log_item: + all_scalars["time/train_teacher_onload"] = rank0_log_item["teacher_onload_time"] + if "teacher_offload_time" in rank0_log_item: + all_scalars["time/train_teacher_offload"] = rank0_log_item["teacher_offload_time"] for worker_idx, log_item in enumerate(train_info["workers_log_item"]): if not self._display_all_workers_log and worker_idx > 0: break @@ -1481,6 +1570,10 @@ def _log_step( for key, value in mini_batch_metrics.items(): avg_value = sum(value) / len(value) all_scalars.update({f"train_metrics/worker_{worker_idx}/step_avg_{key}": avg_value}) + if worker_idx == 0 and key.startswith(_DISTILLATION_METRIC_PREFIXES): + all_scalars[f"distillation/{key}"] = avg_value + if key in ("opd_reverse_kl", "opd_abs_logprob_loss"): + all_scalars[key] = avg_value rank_sft_log = log_item["sft_train_metrics"] for k, v in rank_sft_log.items(): @@ -1509,8 +1602,9 @@ def _save_trajectories(self, data_groups: list[list[RolloutState]], save_path: P continue valid_groups.append(group) for data in group: - assert data.reward is not None - rewards.append(data.reward["score"]) + reward = data.reward.get("score") if data.reward is not None else None + if reward is not None: + rewards.append(reward) response_ids = self._get_trajectory_response_ids(data) response_len_list.append(len(response_ids)) @@ -1520,20 +1614,20 @@ def _save_trajectories(self, data_groups: list[list[RolloutState]], save_path: P with open(save_path, "w", encoding="utf-8") as f: summary = { "reward_mean": rewards_tensor.mean().item(), - "reward_std": rewards_tensor.std().item(), + "reward_std": rewards_tensor.std(unbiased=False).item(), "reward_max": rewards_tensor.max().item(), "reward_min": rewards_tensor.min().item(), "response_len_mean": response_lens.mean().item(), - "response_len_std": response_lens.std().item(), + "response_len_std": response_lens.std(unbiased=False).item(), "response_len_max": response_lens.max().item(), "response_len_min": response_lens.min().item(), - "total_len": len(rewards), + "total_len": len(response_len_list), } json.dump(summary, f, ensure_ascii=False, separators=(",", ":")) f.write("\n") for group in valid_groups: for data in group: - assert data.reward is not None + reward = data.reward.get("score") if data.reward is not None else None response_ids = self._get_trajectory_response_ids(data) response = data.response if response is None and response_ids: @@ -1553,7 +1647,7 @@ def _save_trajectories(self, data_groups: list[list[RolloutState]], save_path: P "prompt": data.message, "label": ground_truth, "response": response, - "reward": data.reward["score"], + "reward": reward, "prompt_len": data.num_tokens, "response_len": len(response_ids), "reward_payload": data.reward, @@ -1839,7 +1933,7 @@ def _fit(self): train_step, step_timer_dict, offload_rollout_before_train=True, - onload_train_before_train=True, + resume_train_before_train=True, raw_rewards_sum=produce_result.raw_rewards_sum, raw_rewards_count=produce_result.raw_rewards_count, ) @@ -1879,7 +1973,7 @@ def _fit_debug_train(self) -> None: train_step, step_timer_dict, offload_rollout_before_train=False, - onload_train_before_train=False, + resume_train_before_train=False, ) eval_log_info: dict[str, float] = {} produce_result = ProduceBatchResult(rollout_states=train_batch) From d40d0e7fb5220f1991d9279dd9507f63f02fb37a Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Fri, 28 Aug 2026 04:13:40 +0000 Subject: [PATCH 02/20] support compute top-k logprobs in lmdeploy --- xtuner/v1/data_proto/rl_data.py | 6 +- xtuner/v1/rl/agent_loop/agent_loop.py | 3 +- xtuner/v1/rl/distillation/config.py | 22 ++------ .../rl/distillation/rollout_teacher_client.py | 55 ++++++++++++++++--- xtuner/v1/rl/trainer/controller.py | 36 +++++++++++- xtuner/v1/rl/trainer/worker.py | 2 + xtuner/v1/train/rl_trainer.py | 31 ++++++++--- 7 files changed, 119 insertions(+), 36 deletions(-) diff --git a/xtuner/v1/data_proto/rl_data.py b/xtuner/v1/data_proto/rl_data.py index 95e3f23831..6c6a96f65e 100644 --- a/xtuner/v1/data_proto/rl_data.py +++ b/xtuner/v1/data_proto/rl_data.py @@ -109,8 +109,10 @@ class RolloutState(BaseModel): tool_calls: list[RolloutToolCall] | None = None response_ids: list[int] | None = None logprobs: list[float] | None = None - teacher_tokens: list[int] | None = None - teacher_logprobs: list[float] | None = None + # Sampled-token teachers store one value per response token. Top-K + # teachers store one K-wide row per response token. + teacher_tokens: list[int] | list[list[int]] | None = None + teacher_logprobs: list[float] | list[list[float]] | None = None routed_experts: np.ndarray | RayObjectRef | list[RayObjectRef] | None = None finish_reason: str | None = None # response_mask: 记录response_ids中哪个token算loss, 与response_ids长度相同，每轮rollout在 agent_loop.generate 中覆盖写 diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index 844b5ed87c..1a0c2e34a6 100644 --- a/xtuner/v1/rl/agent_loop/agent_loop.py +++ b/xtuner/v1/rl/agent_loop/agent_loop.py @@ -259,7 +259,8 @@ def configure_distillation(self, distillation_config: DistillationConfig | None) validate_opd_sample_params(self.sample_params) self.teacher_clients = { - teacher.name: RolloutTeacherClient(teacher) for teacher in distillation_config.rollout_teachers + teacher.name: RolloutTeacherClient(teacher, distillation_config.loss_config) + for teacher in distillation_config.rollout_teachers } self.data_source_teacher_map = dict(distillation_config.data_source_teacher_map) diff --git a/xtuner/v1/rl/distillation/config.py b/xtuner/v1/rl/distillation/config.py index 8630e32726..f9557637db 100644 --- a/xtuner/v1/rl/distillation/config.py +++ b/xtuner/v1/rl/distillation/config.py @@ -77,22 +77,12 @@ def validate_config(self) -> DistillationConfig: if unknown_teachers: raise ValueError(f"data_source_teacher_map references unknown teachers: {sorted(unknown_teachers)}") - if self.loss_config.uses_sampled_token_targets: - teacher_types = {type(teacher) for teacher in self.teachers} - if len(teacher_types) != 1: - raise ValueError( - "Sampled-token distillation requires all teachers to use the same runtime type; " - "mixing RolloutTeacherConfig and TrainTeacherConfig is not supported" - ) - else: - incompatible_teachers = [ - teacher.name for teacher in self.teachers if not isinstance(teacher, TrainTeacherConfig) - ] - if incompatible_teachers: - raise ValueError( - f"loss_mode={self.loss_config.loss_mode!r} requires " - f"TrainTeacherConfig, got incompatible teachers: {incompatible_teachers}" - ) + teacher_types = {type(teacher) for teacher in self.teachers} + if len(teacher_types) != 1: + raise ValueError( + "Distillation requires all teachers to use the same runtime type; " + "mixing RolloutTeacherConfig and TrainTeacherConfig is not supported" + ) teacher_index_by_name = {teacher.name: index for index, teacher in enumerate(self.teachers)} self._teacher_index_by_data_source = { data_source: teacher_index_by_name[teacher_name] diff --git a/xtuner/v1/rl/distillation/rollout_teacher_client.py b/xtuner/v1/rl/distillation/rollout_teacher_client.py index b16b9fef19..c02049635d 100644 --- a/xtuner/v1/rl/distillation/rollout_teacher_client.py +++ b/xtuner/v1/rl/distillation/rollout_teacher_client.py @@ -10,6 +10,7 @@ import httpx from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status +from xtuner.v1.rl.loss import DistillationLossConfig from xtuner.v1.utils import get_logger from .config import RolloutTeacherConfig @@ -64,10 +65,13 @@ def resolve_replica_idx( class RolloutTeacherClient: """Asynchronous teacher client scoped to one AgentLoop.""" - def __init__(self, config: RolloutTeacherConfig) -> None: + def __init__(self, config: RolloutTeacherConfig, loss_config: DistillationLossConfig) -> None: self.config = config + self.loss_config = loss_config self.name = config.name self.backend = self._resolve_backend_from_env() + if self.loss_config.uses_topk_targets and self.backend != "lmdeploy": + raise RuntimeError("Rollout Teacher Top-K targets currently require LMDeploy") self.urls = [f"{endpoint.rstrip('/')}/generate" for endpoint in config.endpoints] self._semaphores = [asyncio.Semaphore(config.max_concurrency) for _ in self.urls] self._replica_router = RolloutTeacherReplicaRouter(len(self.urls)) @@ -117,12 +121,21 @@ async def compute_logprobs(self, state: RolloutState) -> RolloutState: async with self._semaphores[replica_idx]: response = await self._client.post(url, json=payload) response.raise_for_status() - teacher_tokens, teacher_logprobs = self._parse_response( - response, - response_ids, - logprob_start_len=logprob_start_len, - expanded_prompt_len=expanded_prompt_len, - ) + teacher_tokens: list[int] | list[list[int]] + teacher_logprobs: list[float] | list[list[float]] + if self.loss_config.uses_topk_targets: + teacher_tokens, teacher_logprobs = self._parse_topk_response( + response, + response_ids, + expanded_prompt_len=expanded_prompt_len, + ) + else: + teacher_tokens, teacher_logprobs = self._parse_response( + response, + response_ids, + logprob_start_len=logprob_start_len, + expanded_prompt_len=expanded_prompt_len, + ) state.teacher_tokens = teacher_tokens state.teacher_logprobs = teacher_logprobs return state @@ -177,6 +190,8 @@ def _construct_payload( raise RuntimeError(f"Unsupported teacher backend: {self.backend}") if image_data: payload["image_data"] = image_data + if self.loss_config.uses_topk_targets: + payload["top_logprobs_num"] = cast(int, self.loss_config.top_k) return payload @staticmethod @@ -238,6 +253,32 @@ def _parse_response( ) return self._validate_response_logprobs(response_logprobs, response_ids) + def _parse_topk_response( + self, + response: httpx.Response, + response_ids: list[int], + *, + expanded_prompt_len: int | None = None, + ) -> tuple[list[list[int]], list[list[float]]]: + meta_info = response.json()["meta_info"] + prompt_token_count = meta_info["prompt_tokens"] + if expanded_prompt_len is not None: + expected_prompt_token_count = expanded_prompt_len + len(response_ids) + if prompt_token_count != expected_prompt_token_count: + raise ValueError( + "LMDeploy expanded prompt length mismatch: " + f"expected {expected_prompt_token_count} total input tokens " + f"({expanded_prompt_len} prompt + {len(response_ids)} response), " + f"got {prompt_token_count}" + ) + + response_topk = meta_info["input_top_logprobs"][-len(response_ids) :] + teacher_tokens = [[int(token_id) for _, token_id in row] for row in response_topk] + teacher_logprobs = [[float(logprob) for logprob, _ in row] for row in response_topk] + if not all(math.isfinite(logprob) for row in teacher_logprobs for logprob in row): + raise ValueError("Teacher Top-K logprobs contain NaN or Inf") + return teacher_tokens, teacher_logprobs + @staticmethod def _parse_sglang_response( response: httpx.Response, diff --git a/xtuner/v1/rl/trainer/controller.py b/xtuner/v1/rl/trainer/controller.py index 7484a542c4..ceb756088d 100644 --- a/xtuner/v1/rl/trainer/controller.py +++ b/xtuner/v1/rl/trainer/controller.py @@ -24,6 +24,7 @@ class ColateItem(TypedDict): advantage: float rollout_logprobs: torch.Tensor | None teacher_logprobs: NotRequired[torch.Tensor | None] + target_token_ids: NotRequired[torch.Tensor | None] teacher_indices: NotRequired[torch.Tensor] @@ -128,6 +129,10 @@ def _packing(self, data_batches, pack_max_length, language_cfg): if "teacher_logprobs" in data_batches[0] and data_batches[0]["teacher_logprobs"] is not None: teacher_logprobs_list = [data_batches[i]["teacher_logprobs"] for i in indices] + target_token_ids_list = None + if "target_token_ids" in data_batches[0] and data_batches[0]["target_token_ids"] is not None: + target_token_ids_list = [data_batches[i]["target_token_ids"] for i in indices] + teacher_indices_list = None if "teacher_indices" in data_batches[0] and data_batches[0]["teacher_indices"] is not None: teacher_indices_list = [data_batches[i]["teacher_indices"] for i in indices] @@ -174,13 +179,21 @@ def _packing(self, data_batches, pack_max_length, language_cfg): ) rollout_logprobs_list.append(pad_rollout_logprobs) if teacher_logprobs_list is not None: + teacher_logprobs_shape = data_batches[0]["teacher_logprobs"].shape[2:] pad_teacher_logprobs = torch.zeros( - 1, - pad_len, + (1, pad_len, *teacher_logprobs_shape), dtype=data_batches[0]["teacher_logprobs"].dtype, device=data_batches[0]["shifted_labels"].device, ) teacher_logprobs_list.append(pad_teacher_logprobs) + if target_token_ids_list is not None: + target_token_ids_shape = data_batches[0]["target_token_ids"].shape[2:] + pad_target_token_ids = torch.zeros( + (1, pad_len, *target_token_ids_shape), + dtype=data_batches[0]["target_token_ids"].dtype, + device=data_batches[0]["shifted_labels"].device, + ) + target_token_ids_list.append(pad_target_token_ids) if teacher_indices_list is not None: pad_teacher_indices = torch.full( (1, pad_len), @@ -203,6 +216,10 @@ def _packing(self, data_batches, pack_max_length, language_cfg): if teacher_logprobs_list is not None: teacher_logprobs = torch.cat(teacher_logprobs_list, dim=1) # (1, max_len) + target_token_ids = None + if target_token_ids_list is not None: + target_token_ids = torch.cat(target_token_ids_list, dim=1) + teacher_indices = None if teacher_indices_list is not None: teacher_indices = torch.cat(teacher_indices_list, dim=1) # (1, max_len) @@ -214,6 +231,7 @@ def _packing(self, data_batches, pack_max_length, language_cfg): "advantages": advantages, "rollout_logprobs": rollout_logprobs, "teacher_logprobs": teacher_logprobs, + "target_token_ids": target_token_ids, "teacher_indices": teacher_indices, } ) @@ -299,8 +317,19 @@ def fit(self, data_batches: list[ColateItem], pack_max_length: int, rollout_idx: ) pad_teacher_logprobs = None if "teacher_logprobs" in packed_data_batches[0] and packed_data_batches[0]["teacher_logprobs"] is not None: + teacher_logprobs_shape = packed_data_batches[0]["teacher_logprobs"].shape[2:] pad_teacher_logprobs = torch.zeros( - 1, pack_max_length, dtype=packed_data_batches[0]["teacher_logprobs"].dtype, device="cpu" + (1, pack_max_length, *teacher_logprobs_shape), + dtype=packed_data_batches[0]["teacher_logprobs"].dtype, + device="cpu", + ) + pad_target_token_ids = None + if "target_token_ids" in packed_data_batches[0] and packed_data_batches[0]["target_token_ids"] is not None: + target_token_ids_shape = packed_data_batches[0]["target_token_ids"].shape[2:] + pad_target_token_ids = torch.zeros( + (1, pack_max_length, *target_token_ids_shape), + dtype=packed_data_batches[0]["target_token_ids"].dtype, + device="cpu", ) pad_teacher_indices = None if "teacher_indices" in packed_data_batches[0] and packed_data_batches[0]["teacher_indices"] is not None: @@ -316,6 +345,7 @@ def fit(self, data_batches: list[ColateItem], pack_max_length: int, rollout_idx: "advantages": pad_advantages, "rollout_logprobs": pad_rollout_logprobs, "teacher_logprobs": pad_teacher_logprobs, + "target_token_ids": pad_target_token_ids, "teacher_indices": pad_teacher_indices, } pad_data_samples = [pad_data for _ in range(pad_num)] diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index a358893329..079036da84 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -200,6 +200,7 @@ class WorkerInputItem(TypedDict): advantages: torch.Tensor rollout_logprobs: torch.Tensor | None teacher_logprobs: NotRequired[torch.Tensor | None] + target_token_ids: NotRequired[torch.Tensor | None] teacher_indices: NotRequired[torch.Tensor] @@ -707,6 +708,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo "advantages": advantages, "rollout_logprobs": rollout_logprobs, "teacher_logprobs": data.get("teacher_logprobs", None), + "target_token_ids": data.get("target_token_ids", None), }, sp_mesh=self.sp_mesh, ) diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index b8dcda7c4f..f6f5897440 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -1384,11 +1384,28 @@ def _prepare_train_data( and self._distillation_loss_cfg is not None and self._distillation_config.rollout_teachers ): - response_teacher_logprobs = cast(list[float], group[i].teacher_logprobs) - data_dict["teacher_logprobs"] = torch.tensor( - [0.0] * (len(prompt_ids) - 1) + response_teacher_logprobs, - dtype=torch.float32, - ).unsqueeze(0) + if self._distillation_loss_cfg.uses_topk_targets: + top_k = cast(int, self._distillation_loss_cfg.top_k) + response_teacher_tokens = cast(list[list[int]], group[i].teacher_tokens) + response_topk_logprobs = cast(list[list[float]], group[i].teacher_logprobs) + assert len(response_teacher_tokens) == len(response_ids) + assert len(response_topk_logprobs) == len(response_ids) + prompt_target_tokens = [[0] * top_k for _ in range(len(prompt_ids) - 1)] + prompt_teacher_logprobs = [[0.0] * top_k for _ in range(len(prompt_ids) - 1)] + data_dict["target_token_ids"] = torch.tensor( + prompt_target_tokens + response_teacher_tokens, + dtype=torch.int64, + ).unsqueeze(0) + data_dict["teacher_logprobs"] = torch.tensor( + prompt_teacher_logprobs + response_topk_logprobs, + dtype=torch.float32, + ).unsqueeze(0) + else: + response_sampled_logprobs = cast(list[float], group[i].teacher_logprobs) + data_dict["teacher_logprobs"] = torch.tensor( + [0.0] * (len(prompt_ids) - 1) + response_sampled_logprobs, + dtype=torch.float32, + ).unsqueeze(0) if teacher_index is not None: data_dict["teacher_indices"] = torch.full_like( shifted_labels_t, @@ -1875,7 +1892,7 @@ def fit(self): self._exp_tracker.close() close_trace() - def _fit(self): + def _fit(self) -> None: self.logger.info("Start RL training") if self._cur_step >= self._total_train_steps: self.logger.info(f"Train steps {self._total_train_steps} reached, stop training") @@ -1901,7 +1918,7 @@ def _fit(self): model_step = self._get_colocate_rollout_model_step(init_train_step) for train_step in range(init_train_step, self._total_train_steps + 1): self.logger.info(f"Train step {train_step}/{self._total_train_steps} start") - step_timer_dict = {} + step_timer_dict: dict[str, float] = {} with timer("step", step_timer_dict): # 共卡一次调用内完成生产和消费。 self.rl_health_manager.set_rollout_resources_available(True) From 08a702431d88d2739c4d33542413bb90925b3a28 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Fri, 28 Aug 2026 05:37:15 +0000 Subject: [PATCH 03/20] fix agentic rl running bug --- .../agent_in_sandbox_loop.py | 34 +++++++++++++++- .../rl/distillation/rollout_teacher_client.py | 40 ++++++++++++++++--- .../rl/distillation/train_teacher_manager.py | 24 +++++++---- xtuner/v1/rl/trainer/controller.py | 9 +++-- xtuner/v1/train/rl_trainer.py | 27 +++++++++++++ 5 files changed, 117 insertions(+), 17 deletions(-) diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py index 75e6e13f23..15d12da47d 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py @@ -6,11 +6,13 @@ import json import traceback import uuid +from collections.abc import Callable from typing import Any, Literal from lagent.utils import create_object -from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status +from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status +from xtuner.v1.rl.distillation import route_rollout_teacher_client from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController from xtuner.v1.rl.utils import create_task @@ -228,6 +230,34 @@ def __init__( self._sample_semaphore = asyncio.Semaphore(max_concurrent_samples) if max_concurrent_samples else None self.mode = mode + async def collect_rollout_group( + self, + rollout_state: list[RolloutState], + *, + is_valid_sample_func: Callable[[list[RolloutState]], bool] | None = None, + **kwargs, + ) -> list[RolloutState]: + """Generate agent traces and score every completed segment.""" + group = await self.generate_group(rollout_state, **kwargs) + if get_group_status(group) != Status.COMPLETED: + return group + if is_valid_sample_func is not None and not is_valid_sample_func(group): + for state in group: + state.status = Status.FILTERED + return group + if not self.teacher_clients: + return group + + async def score_trace(state: RolloutState) -> RolloutState: + teacher = route_rollout_teacher_client( + state, + data_source_teacher_map=self.data_source_teacher_map, + teacher_clients=self.teacher_clients, + ) + return await teacher.compute_logprobs(state) + + return list(await asyncio.gather(*(create_task(score_trace(state)) for state in group))) + async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: async def generate_one(state: RolloutState) -> list[RolloutState]: if self._sample_semaphore is None: @@ -289,6 +319,7 @@ async def _build_rollout_states(self, rollout_state: RolloutState, item: AgentRo response_message.get("finish_reason") or ("stop" if item.status == RolloutStatus.COMPLETED else "error") ) rollout_state.reward = _extract_reward_payload(item) + rollout_state.extra_fields["origin_data_source"] = item.data_source rollout_state.extra_fields["agent_status"] = item.status.value selected_agent = _selected_agent(item) if selected_agent is not None: @@ -368,6 +399,7 @@ def _fill_eval_rollout_state(self, rollout_state: RolloutState, item: AgentRollo rollout_state.routed_experts = None rollout_state.response_mask = None rollout_state.response_model_steps = None + rollout_state.extra_fields["origin_data_source"] = item.data_source rollout_state.extra_fields["agent_status"] = item.status.value selected_agent = _selected_agent(item) if selected_agent is not None: diff --git a/xtuner/v1/rl/distillation/rollout_teacher_client.py b/xtuner/v1/rl/distillation/rollout_teacher_client.py index c02049635d..3734c490f5 100644 --- a/xtuner/v1/rl/distillation/rollout_teacher_client.py +++ b/xtuner/v1/rl/distillation/rollout_teacher_client.py @@ -84,12 +84,10 @@ def __init__(self, config: RolloutTeacherConfig, loss_config: DistillationLossCo async def compute_logprobs(self, state: RolloutState) -> RolloutState: start = time.perf_counter() try: - prompt_ids = cast(list[int], state.prompt_ids) - response_ids = cast(list[int], state.response_ids) - if not prompt_ids or not response_ids: - state.status = Status.FAILED - state.error_msg = f"Teacher {self.name!r} scoring requires non-empty prompt_ids and response_ids" + scoring_input = self._prepare_scoring_input(state) + if scoring_input is None: return state + prompt_ids, response_ids = scoring_input routed_replica_idx = self._replica_router.resolve_replica_idx( teacher_name=self.name, data_source=str(state.extra_fields.get("origin_data_source", "")), @@ -154,6 +152,38 @@ async def compute_logprobs(self, state: RolloutState) -> RolloutState: finally: state.extra_fields["teacher_score_time_s"] = time.perf_counter() - start + def _prepare_scoring_input(self, state: RolloutState) -> tuple[list[int], list[int]] | None: + if state.input_ids is not None: + input_ids = state.input_ids + labels = state.labels + if len(input_ids) < 2: + state.status = Status.FAILED + state.error_msg = f"Teacher {self.name!r} trace scoring requires at least two input_ids" + return None + if labels is None or len(labels) != len(input_ids): + state.status = Status.FAILED + state.error_msg = ( + f"Teacher {self.name!r} trace scoring requires input_ids and labels with equal lengths; " + f"got {len(input_ids)} and {None if labels is None else len(labels)}" + ) + return None + scoring_start = next((index for index, label in enumerate(labels[1:], start=1) if label != -100), None) + if scoring_start is None: + state.status = Status.FAILED + state.error_msg = f"Teacher {self.name!r} trace scoring requires at least one trainable label" + return None + # Keep later masked turns in the scored suffix so subsequent + # assistant turns retain their complete causal context. + return input_ids[:scoring_start], input_ids[scoring_start:] + + prompt_ids = cast(list[int] | None, state.prompt_ids) + response_ids = cast(list[int] | None, state.response_ids) + if not prompt_ids or not response_ids: + state.status = Status.FAILED + state.error_msg = f"Teacher {self.name!r} scoring requires non-empty prompt_ids and response_ids" + return None + return prompt_ids, response_ids + @staticmethod def _resolve_backend_from_env() -> Literal["sglang", "lmdeploy"]: use_sglang = os.environ.get("XTUNER_USE_SGLANG", "0") == "1" diff --git a/xtuner/v1/rl/distillation/train_teacher_manager.py b/xtuner/v1/rl/distillation/train_teacher_manager.py index dc9b2099f4..b441d5fbb6 100644 --- a/xtuner/v1/rl/distillation/train_teacher_manager.py +++ b/xtuner/v1/rl/distillation/train_teacher_manager.py @@ -57,16 +57,24 @@ def __init__(self, distillation_config: DistillationConfig, *, chunk_size: int | mode=mode, ) - self._teachers: list[FrozenModel] = [ - build_frozen_model(teacher.model_cfg, teacher.model_path, teacher.fsdp_cfg) - for teacher in distillation_config.train_teachers - ] + # Build every frozen Teacher before the Actor so checkpoint/config + # errors fail during worker initialization. Each Teacher is offloaded + # immediately, preventing multiple full models from co-residing on GPU. + self._teachers: list[FrozenModel] = [] + for teacher_config in distillation_config.train_teachers: + teacher = build_frozen_model( + teacher_config.model_cfg, + teacher_config.model_path, + teacher_config.fsdp_cfg, + ) + self._teachers.append(teacher) + self._teacher_is_composed = [ isinstance(teacher.model_cfg, BaseComposeConfig) for teacher in distillation_config.train_teachers ] - self._teachers_by_name = { - teacher_config.name: teacher - for teacher_config, teacher in zip(distillation_config.train_teachers, self._teachers) + self._teacher_index_by_name = { + teacher_config.name: teacher_index + for teacher_index, teacher_config in enumerate(distillation_config.train_teachers) } def compute_logprobs( @@ -105,7 +113,7 @@ def offload_all_to_cpu(self) -> None: def offload_to_disk(self, teacher_name: str) -> None: """Reserve the disk-offload lifecycle boundary for a later backend.""" - if teacher_name not in self._teachers_by_name: + if teacher_name not in self._teacher_index_by_name: raise KeyError(f"Unknown training Teacher: {teacher_name!r}") raise NotImplementedError("Train Teacher disk offload is not implemented") diff --git a/xtuner/v1/rl/trainer/controller.py b/xtuner/v1/rl/trainer/controller.py index ceb756088d..3925631931 100644 --- a/xtuner/v1/rl/trainer/controller.py +++ b/xtuner/v1/rl/trainer/controller.py @@ -103,9 +103,12 @@ def _packing(self, data_batches, pack_max_length, language_cfg): ) packed_data_batches = [] - is_qwen3_vl = False - if len(data_batches[0]["seq_ctx"].position_ids.shape) == 3: - is_qwen3_vl = True + is_qwen3_vl = any(data["seq_ctx"].position_ids.ndim == 3 for data in data_batches) + if is_qwen3_vl: + for data in data_batches: + position_ids = data["seq_ctx"].position_ids + if position_ids.ndim == 2: + data["seq_ctx"].position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) has_rollout_routed_experts = False if data_batches[0]["seq_ctx"].rollout_routed_experts is not None: diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index f6f5897440..9661120fcd 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -1267,6 +1267,31 @@ def _prepare_train_data( input_ids = raw_input_ids[:-1] shifted_labels = labels[1:] + teacher_logprobs = None + if ( + self._distillation_config is not None + and self._distillation_loss_cfg is not None + and self._distillation_config.rollout_teachers + ): + raw_teacher_logprobs = group[i].teacher_logprobs + if raw_teacher_logprobs is None: + raise ValueError( + f"Teacher logprobs cannot be None when distillation is enabled: {group[i]}" + ) + teacher_response_start = next( + (index for index, label in enumerate(shifted_labels) if label != -100), + len(shifted_labels), + ) + teacher_response_length = len(shifted_labels) - teacher_response_start + if len(raw_teacher_logprobs) != teacher_response_length: + raise ValueError( + "Teacher logprobs must align with the trainable suffix of shifted agent labels: " + f"{len(raw_teacher_logprobs)} vs {teacher_response_length}, data: {group[i]}" + ) + teacher_logprobs = torch.tensor( + [0.0] * teacher_response_start + raw_teacher_logprobs, + dtype=torch.float32, + ).unsqueeze(0) prompt_len = sum(label == -100 for label in shifted_labels) response_len = len(shifted_labels) - prompt_len prompt_len_list.append(prompt_len) @@ -1297,6 +1322,8 @@ def _prepare_train_data( "advantage": actual_advantages, "rollout_logprobs": rollout_logprobs, } + if teacher_logprobs is not None: + data_dict["teacher_logprobs"] = teacher_logprobs if teacher_index is not None: data_dict["teacher_indices"] = torch.full_like( shifted_labels_t, From 15c1868831d52601fd176f7763bbdf349b677e7a Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Fri, 28 Aug 2026 07:11:26 +0000 Subject: [PATCH 04/20] Fix rollout teacher distillation integration --- .../build_teacher_server_commands.py | 14 +- .../config/rl_dapo_math_mopd.py | 32 +- .../config/rl_dapo_math_mopd_async.py | 34 +- .../config/rl_dapo_math_opd.py | 28 +- tests/rl/test_on_policy_distillation.py | 825 ++++++------------ tests/rl/test_prepare_train_data.py | 160 +++- .../rl/distillation/rollout_teacher_client.py | 103 ++- xtuner/v1/train/rl_trainer.py | 117 ++- 8 files changed, 668 insertions(+), 645 deletions(-) diff --git a/recipe/on_policy_distillation/build_teacher_server_commands.py b/recipe/on_policy_distillation/build_teacher_server_commands.py index 88b3185320..ba71b8e686 100644 --- a/recipe/on_policy_distillation/build_teacher_server_commands.py +++ b/recipe/on_policy_distillation/build_teacher_server_commands.py @@ -20,13 +20,13 @@ REPO_ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO_ROOT)) -from xtuner.v1.rl.on_policy_distillation import OPDTeacherConfig # noqa: E402 +from xtuner.v1.rl.distillation import RolloutTeacherConfig # noqa: E402 from xtuner.v1.utils.config import Config # noqa: E402 @dataclass(frozen=True) class TeacherReplicaRequest: - teacher: OPDTeacherConfig + teacher: RolloutTeacherConfig replica_index: int @property @@ -46,7 +46,7 @@ def build_teacher_launch_server_commands( Args: config_path (str): Path to the XTuner Python config containing - ``opd_config``. + ``distillation_config``. backend (Literal["sglang", "lmdeploy"]): Teacher serving backend. Returns: @@ -77,7 +77,7 @@ def build_teacher_launch_server_commands( assert len(node_addresses) == node_count assert all(node_addresses) - teachers = config.opd_config.teachers + teachers = config.distillation_config.rollout_teachers replica_requests = _expand_teacher_replicas(teachers) health_path, model_info_path = _get_teacher_server_paths(backend) teacher_placements, student_local_num_workers = _allocate_teacher_devices( @@ -138,7 +138,7 @@ def build_teacher_launch_server_commands( def _expand_teacher_replicas( - teachers: list[OPDTeacherConfig], + teachers: list[RolloutTeacherConfig], ) -> list[TeacherReplicaRequest]: return [ TeacherReplicaRequest( @@ -280,7 +280,7 @@ def _get_teacher_server_paths( def _build_sglang_command( - teacher: OPDTeacherConfig, + teacher: RolloutTeacherConfig, *, server_port: int, ) -> list[str]: @@ -321,7 +321,7 @@ def _build_sglang_command( def _build_lmdeploy_command( - teacher: OPDTeacherConfig, + teacher: RolloutTeacherConfig, *, server_port: int, ) -> list[str]: diff --git a/recipe/on_policy_distillation/config/rl_dapo_math_mopd.py b/recipe/on_policy_distillation/config/rl_dapo_math_mopd.py index bef113ab89..7f930090b7 100644 --- a/recipe/on_policy_distillation/config/rl_dapo_math_mopd.py +++ b/recipe/on_policy_distillation/config/rl_dapo_math_mopd.py @@ -16,12 +16,12 @@ ) from xtuner.v1.rl.evaluator import EvaluatorConfig from xtuner.v1.rl.judger import ComposedJudgerConfig, GEO3KJudgerConfig, GSM8KJudgerConfig -from xtuner.v1.rl.loss import GRPOLossConfig -from xtuner.v1.rl.on_policy_distillation import ( - OPDConfig, - OPDTeacherConfig, - OPDTeacherLaunchConfig, +from xtuner.v1.rl.distillation import ( + DistillationConfig, + RolloutTeacherConfig, + RolloutTeacherLaunchConfig, ) +from xtuner.v1.rl.loss import DistillationLossConfig from xtuner.v1.rl.replay_buffer import SyncReplayBufferConfig from xtuner.v1.rl.rollout.worker import RolloutConfig from xtuner.v1.rl.trainer import WorkerConfig @@ -135,7 +135,7 @@ def mopd_compute_metric(samples): if hasattr(model_cfg, "z_loss_cfg"): model_cfg.z_loss_cfg = None optim_cfg = AdamWConfig(lr=1e-6, foreach=False, weight_decay=0.1, betas=(0.9, 0.98)) -loss_cfg = GRPOLossConfig( +loss_cfg = DistillationLossConfig( policy_loss_cfg={ "cliprange_high": 0.2, "cliprange_low": 0.2, @@ -150,6 +150,10 @@ def mopd_compute_metric(samples): kl_loss_type="low_var_kl", mode="chunk", chunk_size=512, + loss_mode="k1", + use_policy_gradient=True, + task_adv_weight=0.0, + distillation_loss_weight=1.0, ) train_worker_cfg = WorkerConfig( model_cfg=model_cfg, @@ -331,16 +335,14 @@ def mopd_compute_metric(samples): # Every training record's data_source must have an entry in # data_source_teacher_map. Teacher model paths are read from # GSM8K_TEACHER_MODEL_PATH and GEO3K_TEACHER_MODEL_PATH. -opd_config = OPDConfig( - mode="pg-opd", - task_adv_weight=0.0, - opd_adv_weight=1.0, +distillation_config = DistillationConfig( + loss_config=loss_cfg, teachers=[ - OPDTeacherConfig( + RolloutTeacherConfig( name="gsm8k_teacher", num_replicas=teacher_num_replicas, enable_prefix_caching=enable_prefix_caching, - launch_config=OPDTeacherLaunchConfig( + launch_config=RolloutTeacherLaunchConfig( model_path=gsm8k_teacher_model_path, num_workers=teacher_replica_num_workers, server_port=13141, @@ -351,11 +353,11 @@ def mopd_compute_metric(samples): gpu_memory_utilization=0.8, ), ), - OPDTeacherConfig( + RolloutTeacherConfig( name="geo3k_teacher", num_replicas=teacher_num_replicas, enable_prefix_caching=enable_prefix_caching, - launch_config=OPDTeacherLaunchConfig( + launch_config=RolloutTeacherLaunchConfig( model_path=geo3k_teacher_model_path, num_workers=teacher_replica_num_workers, server_port=13142, @@ -385,7 +387,7 @@ def mopd_compute_metric(samples): load_from=model_path, train_batch_size=train_batch_size, advantage_estimator_config=GRPOAdvantageConfig(eps=1e-8), - opd_config=opd_config, + distillation_config=distillation_config, enable_evaluate=enable_evaluate, enable_initial_evaluate=enable_evaluate, evaluate_step=evaluate_step, diff --git a/recipe/on_policy_distillation/config/rl_dapo_math_mopd_async.py b/recipe/on_policy_distillation/config/rl_dapo_math_mopd_async.py index 5abf0a1681..974f34af94 100644 --- a/recipe/on_policy_distillation/config/rl_dapo_math_mopd_async.py +++ b/recipe/on_policy_distillation/config/rl_dapo_math_mopd_async.py @@ -16,12 +16,12 @@ ) from xtuner.v1.rl.evaluator import EvaluatorConfig from xtuner.v1.rl.judger import ComposedJudgerConfig, GEO3KJudgerConfig, GSM8KJudgerConfig -from xtuner.v1.rl.loss import GRPOLossConfig -from xtuner.v1.rl.on_policy_distillation import ( - OPDConfig, - OPDTeacherConfig, - OPDTeacherLaunchConfig, +from xtuner.v1.rl.distillation import ( + DistillationConfig, + RolloutTeacherConfig, + RolloutTeacherLaunchConfig, ) +from xtuner.v1.rl.loss import DistillationLossConfig from xtuner.v1.rl.replay_buffer import AsyncReplayBufferConfig from xtuner.v1.rl.rollout.worker import RolloutConfig from xtuner.v1.rl.trainer import RolloutImportanceSampling, WorkerConfig @@ -130,7 +130,7 @@ def mopd_compute_metric(samples): if hasattr(model_cfg, "z_loss_cfg"): model_cfg.z_loss_cfg = None optim_cfg = AdamWConfig(lr=1e-6, foreach=False, weight_decay=0.1, betas=(0.9, 0.98)) -loss_cfg = GRPOLossConfig( +loss_cfg = DistillationLossConfig( policy_loss_cfg={ "cliprange_high": 0.2, "cliprange_low": 0.2, @@ -145,6 +145,10 @@ def mopd_compute_metric(samples): kl_loss_type="low_var_kl", mode="chunk", chunk_size=512, + loss_mode="k1", + use_policy_gradient=True, + task_adv_weight=0.0, + distillation_loss_weight=1.0, rollout_is=RolloutImportanceSampling( rollout_is_level="token", rollout_is_mode="both", @@ -221,6 +225,8 @@ def mopd_compute_metric(samples): over_sample_threshold=1.0, enable_partial_rollout=True, max_staleness=2, + max_token_staleness=0, + tail_batch_trigger_size=0 ) agent_loop_manager_cfg = AgentLoopManagerConfig( tasks=TaskSpecConfig( @@ -337,14 +343,12 @@ def mopd_compute_metric(samples): # Every training record's data_source must have an entry in # data_source_teacher_map. Teacher model paths are read from # GSM8K_TEACHER_MODEL_PATH and GEO3K_TEACHER_MODEL_PATH. -opd_config = OPDConfig( - mode="pg-opd", - task_adv_weight=0.0, - opd_adv_weight=1.0, +distillation_config = DistillationConfig( + loss_config=loss_cfg, teachers=[ - OPDTeacherConfig( + RolloutTeacherConfig( name="gsm8k_teacher", - launch_config=OPDTeacherLaunchConfig( + launch_config=RolloutTeacherLaunchConfig( model_path=gsm8k_teacher_model_path, num_workers=1, server_port=13141, @@ -355,9 +359,9 @@ def mopd_compute_metric(samples): gpu_memory_utilization=0.8, ), ), - OPDTeacherConfig( + RolloutTeacherConfig( name="geo3k_teacher", - launch_config=OPDTeacherLaunchConfig( + launch_config=RolloutTeacherLaunchConfig( model_path=geo3k_teacher_model_path, num_workers=1, server_port=13142, @@ -387,7 +391,7 @@ def mopd_compute_metric(samples): load_from=model_path, train_batch_size=train_batch_size, advantage_estimator_config=GRPOAdvantageConfig(eps=1e-8), - opd_config=opd_config, + distillation_config=distillation_config, enable_evaluate=enable_evaluate, enable_initial_evaluate=enable_evaluate, evaluate_step=evaluate_step, diff --git a/recipe/on_policy_distillation/config/rl_dapo_math_opd.py b/recipe/on_policy_distillation/config/rl_dapo_math_opd.py index da55149dc6..05e1e026bb 100644 --- a/recipe/on_policy_distillation/config/rl_dapo_math_opd.py +++ b/recipe/on_policy_distillation/config/rl_dapo_math_opd.py @@ -17,12 +17,12 @@ ) from xtuner.v1.rl.evaluator import EvaluatorConfig from xtuner.v1.rl.judger import DapoMathJudgerConfig -from xtuner.v1.rl.loss import GRPOLossConfig -from xtuner.v1.rl.on_policy_distillation import ( - OPDConfig, - OPDTeacherConfig, - OPDTeacherLaunchConfig, +from xtuner.v1.rl.distillation import ( + DistillationConfig, + RolloutTeacherConfig, + RolloutTeacherLaunchConfig, ) +from xtuner.v1.rl.loss import DistillationLossConfig from xtuner.v1.rl.replay_buffer import SyncReplayBufferConfig from xtuner.v1.rl.rollout.worker import RolloutConfig from xtuner.v1.rl.trainer import WorkerConfig @@ -92,7 +92,7 @@ if hasattr(model_cfg, "z_loss_cfg"): model_cfg.z_loss_cfg = None optim_cfg = AdamWConfig(lr=1e-6, foreach=False, weight_decay=0.1, betas=(0.9, 0.98)) -loss_cfg = GRPOLossConfig( +loss_cfg = DistillationLossConfig( policy_loss_cfg={ "cliprange_high": 0.28, "cliprange_low": 0.2, @@ -107,6 +107,10 @@ kl_loss_type="low_var_kl", mode="chunk", chunk_size=512, + loss_mode="k1", + use_policy_gradient=True, + task_adv_weight=0.0, + distillation_loss_weight=1.0, ) train_worker_cfg = WorkerConfig( model_cfg=model_cfg, @@ -206,15 +210,13 @@ evaluator_config = EvaluatorConfig() # 6. pure on-policy distillation -opd_config = OPDConfig( - mode="pg-opd", - task_adv_weight=0.0, - opd_adv_weight=1.0, +distillation_config = DistillationConfig( + loss_config=loss_cfg, teachers=[ - OPDTeacherConfig( + RolloutTeacherConfig( name="teacher", enable_prefix_caching=True, - launch_config=OPDTeacherLaunchConfig( + launch_config=RolloutTeacherLaunchConfig( model_path=teacher_model_path, num_workers=1, server_port=13141, @@ -236,7 +238,7 @@ load_from=model_path, train_batch_size=train_batch_size, advantage_estimator_config=GRPOAdvantageConfig(eps=1e-8), - opd_config=opd_config, + distillation_config=distillation_config, enable_evaluate=enable_evaluate, enable_initial_evaluate=enable_evaluate, evaluate_step=evaluate_step, diff --git a/tests/rl/test_on_policy_distillation.py b/tests/rl/test_on_policy_distillation.py index fa627bbee4..b9d3a98dfd 100644 --- a/tests/rl/test_on_policy_distillation.py +++ b/tests/rl/test_on_policy_distillation.py @@ -1,579 +1,334 @@ -import math -import os -import signal -import subprocess -import time +import tempfile import unittest from pathlib import Path -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import httpx import torch from recipe.on_policy_distillation.build_teacher_server_commands import ( - build_teacher_server_command, + build_teacher_launch_server_commands, ) from xtuner.v1.data_proto.rl_data import RolloutState, Status -from xtuner.v1.rl.loss import GRPOLossConfig -from xtuner.v1.rl.on_policy_distillation import ( - OPDConfig, - OPDTeacherConfig, - OPDTeacherLaunchConfig, - TeacherLogprobClient, - apply_opd_kl_to_advantages, +from xtuner.v1.data_proto.sequence_context import SequenceContext +from xtuner.v1.rl.distillation import RolloutTeacherClient, RolloutTeacherConfig +from xtuner.v1.rl.loss import DistillationLossConfig +from xtuner.v1.rl.trainer.controller import TrainingController + + +class TestDistillationRecipeConfig(unittest.TestCase): + def test_teacher_launcher_reads_distillation_config(self) -> None: + config_source = """ +from xtuner.v1.rl.distillation import ( + DistillationConfig, + RolloutTeacherConfig, + RolloutTeacherLaunchConfig, ) -from xtuner.v1.rl.utils import find_free_ports - - -REPO_ROOT = Path(__file__).resolve().parents[2] -BASELINE_PATH = os.getenv("XTUNER_OPD_BASELINE") -TEACHER_MODEL_PATH = os.getenv("XTUNER_OPD_TEACHER_MODEL") -STUDENT_MODEL_PATH = os.getenv("XTUNER_OPD_STUDENT_MODEL") -TEACHER_STARTUP_TIMEOUT_S = float(os.getenv("XTUNER_OPD_TEACHER_STARTUP_TIMEOUT_S", "1200")) -TRAINER_CONFIG_PATH = ( - REPO_ROOT / "recipe/on_policy_distillation/config/rl_dapo_math_opd.py" +from xtuner.v1.rl.loss import DistillationLossConfig + +loss_cfg = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) +distillation_config = DistillationConfig( + loss_config=loss_cfg, + teachers=[ + RolloutTeacherConfig( + name="teacher", + launch_config=RolloutTeacherLaunchConfig( + model_path="/models/teacher", + num_workers=1, + server_port=13141, + ), + ) + ], + data_source_teacher_map={"math": "teacher"}, ) - - -def _wait_for_teacher(process: subprocess.Popen, endpoint: str, backend: str) -> None: - deadline = time.monotonic() + TEACHER_STARTUP_TIMEOUT_S - health_path = "health_generate" if backend == "sglang" else "health" - with httpx.Client(timeout=1.0, trust_env=False) as client: - while time.monotonic() < deadline: - return_code = process.poll() - if return_code is not None: - raise RuntimeError(f"Teacher process exited during startup with code {return_code}") - try: - if client.get(f"{endpoint}/{health_path}").status_code == 200: - return - except httpx.RequestError: - pass - time.sleep(1.0) - raise TimeoutError(f"Teacher did not become ready within {TEACHER_STARTUP_TIMEOUT_S} seconds") - - -def _build_debug_rollout_batch(samples: list[dict]) -> list[list[RolloutState]]: - train_batch = [] - for sample in samples: - sample_index = int(sample["sample_index"]) - group_index = sample.get("group_index") - prompt_ids = torch.as_tensor(sample["prompt_token_ids"], dtype=torch.long).tolist() - response_ids = torch.as_tensor(sample["response_token_ids"], dtype=torch.long).tolist() - rollout_logprobs = torch.as_tensor(sample["rollout_log_probs"], dtype=torch.float32).tolist() - teacher_logprobs = torch.as_tensor(sample["teacher_log_probs"], dtype=torch.float32).tolist() - response_mask = torch.as_tensor(sample["loss_mask"], dtype=torch.bool).int().tolist() - response = str(sample["response"]) - - train_batch.append( - [ - RolloutState( - rollout_id=sample_index, - group_id=sample_index if group_index is None else int(group_index), - message=[], - prompt_ids=prompt_ids, - tokens=prompt_ids, - response=response, - response_ids=response_ids, - logprobs=rollout_logprobs, - teacher_tokens=response_ids, - teacher_logprobs=teacher_logprobs, - response_mask=response_mask, - reward={"score": 0.0}, - finish_reason="stop", - status=Status.COMPLETED, - extra_fields={"origin_data_source": "baseline"}, +""" + with tempfile.TemporaryDirectory() as temp_dir: + config_path = Path(temp_dir) / "distillation_config.py" + config_path.write_text(config_source) + with patch.dict( + "os.environ", + { + "NODE_COUNT": "1", + "NODE_RANK": "0", + "PROC_PER_NODE": "2", + "WORKER_ALL_SOCKET_ADDRS": "127.0.0.1", + }, + ): + endpoint_map, student_num_workers, student_local_num_workers, records = ( + build_teacher_launch_server_commands(str(config_path), "lmdeploy") ) - ] - ) - return train_batch - -@unittest.skipUnless(BASELINE_PATH, "XTUNER_OPD_BASELINE is required") -class TestPGOPDAdvantage(unittest.TestCase): - @classmethod - def setUpClass(cls) -> None: - super().setUpClass() - baseline = torch.load(BASELINE_PATH, map_location="cpu", weights_only=False) - cls.samples = baseline["samples"] + self.assertEqual(endpoint_map, {"teacher": ["http://127.0.0.1:13141"]}) + self.assertEqual(student_num_workers, 1) + self.assertEqual(student_local_num_workers, 1) + self.assertEqual(len(records), 1) + self.assertEqual(records[0][:4], ["teacher[0]", "0", "1", "http://127.0.0.1:13141"]) + self.assertIn("lmdeploy", records[0][7:]) + + +class TestRolloutTeacherClient(unittest.IsolatedAsyncioTestCase): + @staticmethod + def _response(payload: dict) -> httpx.Response: + return httpx.Response( + 200, + request=httpx.Request("POST", "http://teacher/generate"), + json=payload, + ) - def test_compute_advantages_matches_baseline(self) -> None: - config = OPDConfig( - teachers=[OPDTeacherConfig(name="teacher", endpoint="http://unused")], - data_source_teacher_map={"baseline": "teacher"}, + def _build_topk_client(self, *, max_retry_per_sample: int = 0) -> RolloutTeacherClient: + loss_config = DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="forward_kl_topk", + use_policy_gradient=False, + top_k=2, ) - loss_cfg = GRPOLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) - - for sample in self.samples: - with self.subTest(sample_index=sample["sample_index"]): - old_logprobs = torch.as_tensor(sample["old_log_probs"], dtype=torch.float32) - teacher_logprobs = torch.as_tensor(sample["teacher_log_probs"], dtype=torch.float32) - loss_mask = torch.as_tensor(sample["loss_mask"], dtype=torch.float32) - shifted_labels = torch.where( - loss_mask.bool(), - torch.zeros_like(loss_mask, dtype=torch.long), - torch.full_like(loss_mask, -100, dtype=torch.long), - ) - loss_ctx = loss_cfg.build( - { - "shifted_labels": shifted_labels, - "advantages": torch.zeros_like(old_logprobs), - "old_logprobs": old_logprobs, - "teacher_logprobs": teacher_logprobs, - } - ) - assert loss_ctx is not None - apply_opd_kl_to_advantages(loss_ctx, config=config) - actual_advantages = loss_ctx.loss_kwargs.advantages.cpu() - expected_advantages = torch.as_tensor(sample["advantages"], dtype=torch.float32) * loss_mask - try: - torch.testing.assert_close( - actual_advantages, - expected_advantages, - rtol=1e-5, - atol=1e-5, - ) - except AssertionError as error: - mismatch_indices = torch.nonzero( - ~torch.isclose(actual_advantages, expected_advantages, rtol=1e-5, atol=1e-5) - ).flatten() - mismatch_values = "\n".join( - ( - f"index={index}: " - f"actual={actual_advantages[index].item()!r}, " - f"expected={expected_advantages[index].item()!r}, " - f"abs_diff={abs(actual_advantages[index] - expected_advantages[index]).item()!r}" - ) - for index in mismatch_indices.tolist() - ) - raise AssertionError(f"{error}\n\nMismatched advantages:\n{mismatch_values}") from None - - -def _load_trainer_samples(capture_dir: Path) -> dict[int, dict]: - trainer_samples = {} - for capture_file in sorted(capture_dir.glob("rank_*.pt")): - for batch in torch.load(capture_file, map_location="cpu", weights_only=True): - shifted_labels = torch.as_tensor(batch["shifted_labels"], dtype=torch.long).reshape(-1) - old_log_probs = torch.as_tensor(batch["old_log_probs"], dtype=torch.float32).reshape(-1) - advantages = torch.as_tensor(batch["advantages"], dtype=torch.float32).reshape(-1) - boundaries = torch.as_tensor(batch["cu_seq_lens_q"], dtype=torch.long).tolist() - num_padding = int(batch["num_padding"]) - padding_start = shifted_labels.numel() - num_padding - - assert boundaries[len(batch["rollout_ids"])] == padding_start - torch.testing.assert_close( - shifted_labels[padding_start:], - torch.full_like(shifted_labels[padding_start:], -100), - rtol=0, - atol=0, + with patch.dict( + "os.environ", + {"XTUNER_USE_LMDEPLOY": "1", "XTUNER_USE_SGLANG": "0", "XTUNER_USE_VLLM": "0"}, + ): + client = RolloutTeacherClient( + RolloutTeacherConfig( + name="teacher", + endpoints=["http://teacher"], + max_retry_per_sample=max_retry_per_sample, + ), + loss_config, ) + self.addAsyncCleanup(client._client.aclose) + return client + + @staticmethod + def _state() -> RolloutState: + return RolloutState( + group_id=1, + message=[], + prompt_ids=[10, 11, 12], + response_ids=[13, 14], + status=Status.COMPLETED, + extra_fields={"origin_data_source": "math"}, + ) - for rollout_id, start, end in zip(batch["rollout_ids"], boundaries[:-1], boundaries[1:]): - response_mask = shifted_labels[start:end] != -100 - trainer_samples[int(rollout_id)] = { - "response_token_ids": shifted_labels[start:end][response_mask], - "old_log_probs": old_log_probs[start:end][response_mask], - "advantages": advantages[start:end][response_mask], + async def test_compute_sampled_token_logprobs_uses_current_interface(self) -> None: + response = httpx.Response( + 200, + request=httpx.Request("POST", "http://teacher/generate"), + json={ + "meta_info": { + "prompt_tokens": 5, + "input_token_logprobs": [ + [-0.1, 11], + [-0.2, 12], + [-0.3, 13], + [-0.4, 14], + ], } - - return trainer_samples - - -def _run_trainer_once( - baseline_samples: list[dict], - *, - baseline_path: Path, - student_model_path: Path, - debug_rollout_dir: Path, - capture_dir: Path, - run_dir: Path, -) -> dict[int, dict]: - import ray - - from xtuner.v1.rl.trainer.controller import TrainingController - from xtuner.v1.rl.trainer.worker import TrainingWorker, WorkerConfig - from xtuner.v1.train.rl_trainer import RLColocateTrainer - from xtuner.v1.utils import Config - - max_prompt_length = max(len(sample["prompt_token_ids"]) for sample in baseline_samples) - max_response_length = max(len(sample["response_token_ids"]) for sample in baseline_samples) - max_input_length = max(len(sample["token_ids"]) - 1 for sample in baseline_samples) - pack_max_length = math.ceil(max_input_length / 512) * 512 - num_workers = 8 - - class CapturingRLColocateTrainer(RLColocateTrainer): - def _prepare_train_data( - self, - data_groups, - pack_max_length, - raw_rewards_sum=0.0, - raw_rewards_count=0, + }, + ) + loss_config = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + with patch.dict( + "os.environ", + {"XTUNER_USE_LMDEPLOY": "1", "XTUNER_USE_SGLANG": "0", "XTUNER_USE_VLLM": "0"}, ): - data_batches, data_info = super()._prepare_train_data( - data_groups, - pack_max_length, - raw_rewards_sum=raw_rewards_sum, - raw_rewards_count=raw_rewards_count, - ) - rollout_states = (state for group in data_groups for state in group) - for data_batch, rollout_state in zip(data_batches, rollout_states): - data_batch["rollout_id"] = rollout_state.rollout_id - return data_batches, data_info - - class CapturingTrainingController(TrainingController): - def _packing(self, data_batches, pack_max_length, language_cfg): - pack_infos = self._get_pack_infos( - data_batches, - [data["seq_ctx"].input_ids.numel() for data in data_batches], - pack_max_length, + client = RolloutTeacherClient( + RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"]), + loss_config, ) - packed_data_batches = super()._packing(data_batches, pack_max_length, language_cfg) - for packed_data, pack_info in zip(packed_data_batches, pack_infos): - packed_data["rollout_ids"] = [data_batches[index]["rollout_id"] for index in pack_info["indices"]] - return packed_data_batches - - class CapturingTrainingWorker(TrainingWorker): - trainer_capture_dir = capture_dir - - def fit(self, data_batches, rollout_idx): - from unittest.mock import patch as mock_patch - - from xtuner.v1.rl.trainer import worker as worker_module - - captured_batches = [ - { - "rollout_ids": data.get("rollout_ids", []), - "cu_seq_lens_q": data["seq_ctx"].cu_seq_lens_q.detach().cpu(), - "num_padding": data["seq_ctx"].num_padding, - "shifted_labels": data["shifted_labels"].detach().cpu().reshape(-1), - } - for data in data_batches - ] - captured_batch_index = 0 - apply_opd = worker_module.apply_opd_kl_to_advantages - - def capture_opd_result(loss_ctx, *, config): - nonlocal captured_batch_index - reverse_kl_sum = apply_opd(loss_ctx, config=config) - captured_batches[captured_batch_index]["old_log_probs"] = ( - loss_ctx.loss_kwargs.old_logprobs.detach().cpu().reshape(-1) - ) - captured_batches[captured_batch_index]["advantages"] = ( - loss_ctx.loss_kwargs.advantages.detach().cpu().reshape(-1) - ) - captured_batch_index += 1 - return reverse_kl_sum - - with mock_patch.object(worker_module, "apply_opd_kl_to_advantages", capture_opd_result): - worker_log_item = TrainingWorker.fit(self, data_batches, rollout_idx) + self.addAsyncCleanup(client._client.aclose) + client._client.post = AsyncMock(return_value=response) + state = RolloutState( + group_id=1, + message=[], + prompt_ids=[10, 11, 12], + response_ids=[13, 14], + status=Status.COMPLETED, + extra_fields={"origin_data_source": "math"}, + ) - torch.save(captured_batches, self.trainer_capture_dir / f"rank_{self.rank}.pt") - return worker_log_item + result = await client.compute_logprobs(state) - def build_capturing_training_workers(self, placement_group): - from xtuner.v1.rl.utils import AutoAcceleratorWorkers + self.assertEqual(result.status, Status.COMPLETED) + self.assertEqual(result.teacher_tokens, [13, 14]) + self.assertEqual(result.teacher_logprobs, [-0.3, -0.4]) + self.assertIn("teacher_score_time_s", result.extra_fields) - capturing_worker_cls = ray.remote( - runtime_env={ - "env_vars": { - "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", - "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1", - "HCCL_NPU_SOCKET_PORT_RANGE": "auto", + async def test_malformed_sampled_response_becomes_failed_state(self) -> None: + response = self._response( + { + "meta_info": { + "input_token_logprobs": [ + [-0.1, 11], + [-0.2, 12], + [-0.3, 13], + [-0.4, 14], + ] } } - )(CapturingTrainingWorker) - train_workers, _ = AutoAcceleratorWorkers.from_placement_group( - capturing_worker_cls, - self, - placement_group, ) - ray.wait([worker.ready.remote() for worker in train_workers]) - return CapturingTrainingController(workers=train_workers) - - trainer_environment = { - "WORK_DIR": str(run_dir / "trainer"), - "MODEL_PATH": str(student_model_path), - "TEACHER_MODEL_PATH": os.getenv("XTUNER_OPD_TEACHER_MODEL", str(student_model_path)), - "DATA_PATH": os.getenv("XTUNER_OPD_DATA_PATH", str(baseline_path)), - "WORLD_SIZE": "1", - "ONLY_CALC_MISMATCH_RATIO": "1", - "XTUNER_DETERMINISTIC": "true", - "XTUNER_USE_FA3": os.getenv("XTUNER_OPD_USE_FA3", "1"), - "TOKENIZERS_PARALLELISM": "false", - "PYTHONUNBUFFERED": "1", - } - - with patch.dict(os.environ, trainer_environment, clear=False): - try: - ray.init( - num_cpus=12 * num_workers + 10, - num_gpus=num_workers, - include_dashboard=False, - _temp_dir="/dev/shm/xtuner-opd-old-logprobs", - ) - cfg = Config.fromfile(TRAINER_CONFIG_PATH) - cfg.trainer.resources.num_workers = num_workers - cfg.trainer.total_train_steps = 1 - cfg.trainer.train_batch_size = len(baseline_samples) - cfg.trainer.train_worker_cfg.optimizer_steps = 1 - cfg.trainer.train_worker_cfg.pack_max_length = pack_max_length - cfg.trainer.debug_train = True - cfg.trainer.debug_rollout_dir = debug_rollout_dir - with ( - patch("xtuner.v1.train.rl_trainer.XTUNER_DETERMINISTIC", True), - patch( - "xtuner.v1.train.rl_trainer.RLColocateTrainer", - CapturingRLColocateTrainer, + loss_config = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + with patch.dict( + "os.environ", + {"XTUNER_USE_LMDEPLOY": "1", "XTUNER_USE_SGLANG": "0", "XTUNER_USE_VLLM": "0"}, + ): + client = RolloutTeacherClient( + RolloutTeacherConfig( + name="teacher", + endpoints=["http://teacher"], + max_retry_per_sample=0, ), - patch.object(WorkerConfig, "build", build_capturing_training_workers), - ): - trainer = cfg.trainer.build() - trainer.fit() - finally: - if ray.is_initialized(): - ray.shutdown() - - return _load_trainer_samples(capture_dir) - - -@unittest.skipUnless( - BASELINE_PATH and STUDENT_MODEL_PATH, - "XTUNER_OPD_BASELINE and XTUNER_OPD_STUDENT_MODEL are required", -) -class TestPGOPDOldLogprobs(unittest.TestCase): - def test_trainer_old_logprobs_and_advantage_error_propagation(self) -> None: - baseline_path = Path(str(BASELINE_PATH)).expanduser().resolve() - student_model_path = Path(str(STUDENT_MODEL_PATH)).expanduser().resolve() - baseline = torch.load(baseline_path, map_location="cpu", weights_only=True) - baseline_samples = baseline["samples"] - - work_root = Path( - os.getenv( - "XTUNER_OPD_OLD_LOGPROB_WORK_DIR", - str(REPO_ROOT / "work_dirs/test_pg_opd_old_logprobs"), + loss_config, ) - ).expanduser() - run_dir = work_root / f"run_{time.strftime('%Y%m%d_%H%M%S')}_{os.getpid()}" - debug_rollout_dir = run_dir / "debug_rollout" - capture_dir = run_dir / "trainer_capture" - debug_rollout_dir.mkdir(parents=True) - capture_dir.mkdir() - - train_batch = _build_debug_rollout_batch(baseline_samples) - torch.save(train_batch, debug_rollout_dir / "debug_rollout_1.pt") - trainer_samples_by_rollout_id = _run_trainer_once( - baseline_samples, - baseline_path=baseline_path, - student_model_path=student_model_path, - debug_rollout_dir=debug_rollout_dir, - capture_dir=capture_dir, - run_dir=run_dir, - ) + self.addAsyncCleanup(client._client.aclose) + client._client.post = AsyncMock(return_value=response) - result_path = run_dir / "result.pt" - result_samples = [] - - for baseline_sample in baseline_samples: - rollout_id = int(baseline_sample["sample_index"]) - trainer_sample = trainer_samples_by_rollout_id[rollout_id] - loss_mask = torch.as_tensor(baseline_sample["loss_mask"], dtype=torch.bool) - response_token_ids = torch.as_tensor(baseline_sample["response_token_ids"], dtype=torch.long)[loss_mask] - torch.testing.assert_close( - trainer_sample["response_token_ids"], - response_token_ids, - rtol=0, - atol=0, - ) - baseline_old_log_probs = torch.as_tensor( - baseline_sample["old_log_probs"], - dtype=torch.float32, - )[loss_mask] - trainer_old_log_probs = trainer_sample["old_log_probs"] - baseline_advantages = torch.as_tensor( - baseline_sample["advantages"], - dtype=torch.float32, - )[loss_mask] - trainer_advantages = trainer_sample["advantages"] - - old_logprob_error = trainer_old_log_probs - baseline_old_log_probs - advantage_error = trainer_advantages - baseline_advantages - propagation_residual = advantage_error + old_logprob_error - sample_num_tokens = old_logprob_error.numel() - sample_summary = { - "num_tokens": sample_num_tokens, - "old_logprobs_mean_abs_error": old_logprob_error.abs().mean().item(), - "old_logprobs_max_abs_error": old_logprob_error.abs().max().item(), - "advantages_mean_abs_error": advantage_error.abs().mean().item(), - "advantages_max_abs_error": advantage_error.abs().max().item(), - "propagation_mean_abs_error": propagation_residual.abs().mean().item(), - "propagation_max_abs_error": propagation_residual.abs().max().item(), - } - result_samples.append( - { - "sample_index": rollout_id, - "response_token_ids": response_token_ids, - "baseline_old_log_probs": baseline_old_log_probs, - "trainer_old_log_probs": trainer_old_log_probs, - "old_logprob_error": old_logprob_error, - "baseline_advantages": baseline_advantages, - "trainer_advantages": trainer_advantages, - "advantage_error": advantage_error, - "propagation_residual": propagation_residual, - "summary": sample_summary, - } - ) + result = await client.compute_logprobs(self._state()) - with self.subTest(sample_index=rollout_id): - torch.testing.assert_close( - advantage_error, - -old_logprob_error, - rtol=1e-5, - atol=1e-5, - msg=f"Result: {result_path}", - ) + self.assertEqual(result.status, Status.FAILED) + self.assertIn("prompt_tokens", result.error_msg or "") - global_result = { - "response_token_ids": torch.cat([sample["response_token_ids"] for sample in result_samples]), - "baseline_old_log_probs": torch.cat([sample["baseline_old_log_probs"] for sample in result_samples]), - "trainer_old_log_probs": torch.cat([sample["trainer_old_log_probs"] for sample in result_samples]), - "old_logprob_error": torch.cat([sample["old_logprob_error"] for sample in result_samples]), - "baseline_advantages": torch.cat([sample["baseline_advantages"] for sample in result_samples]), - "trainer_advantages": torch.cat([sample["trainer_advantages"] for sample in result_samples]), - "advantage_error": torch.cat([sample["advantage_error"] for sample in result_samples]), - "propagation_residual": torch.cat([sample["propagation_residual"] for sample in result_samples]), - } - summary = { - "num_samples": len(result_samples), - "num_tokens": global_result["old_logprob_error"].numel(), - "old_logprobs_mean_abs_error": global_result["old_logprob_error"].abs().mean().item(), - "old_logprobs_max_abs_error": global_result["old_logprob_error"].abs().max().item(), - "advantages_mean_abs_error": global_result["advantage_error"].abs().mean().item(), - "advantages_max_abs_error": global_result["advantage_error"].abs().max().item(), - "propagation_mean_abs_error": global_result["propagation_residual"].abs().mean().item(), - "propagation_max_abs_error": global_result["propagation_residual"].abs().max().item(), - } - torch.save( + async def test_sampled_response_rejects_non_numeric_logprob(self) -> None: + response = self._response( { - "summary": summary, - "global": global_result, - "samples": result_samples, - }, - result_path, + "meta_info": { + "prompt_tokens": 5, + "input_token_logprobs": [ + [-0.1, 11], + [-0.2, 12], + [True, 13], + [-0.4, 14], + ], + } + } ) - - with self.subTest(scope="global"): - torch.testing.assert_close( - global_result["advantage_error"], - -global_result["old_logprob_error"], - rtol=1e-5, - atol=1e-5, - msg=f"Result: {result_path}", + loss_config = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + with patch.dict( + "os.environ", + {"XTUNER_USE_LMDEPLOY": "1", "XTUNER_USE_SGLANG": "0", "XTUNER_USE_VLLM": "0"}, + ): + client = RolloutTeacherClient( + RolloutTeacherConfig( + name="teacher", + endpoints=["http://teacher"], + max_retry_per_sample=0, + ), + loss_config, ) + self.addAsyncCleanup(client._client.aclose) + client._client.post = AsyncMock(return_value=response) + + result = await client.compute_logprobs(self._state()) + + self.assertEqual(result.status, Status.FAILED) + self.assertIn("non-numeric logprob", result.error_msg or "") + + async def test_malformed_topk_responses_become_failed_states(self) -> None: + valid_rows = [ + [[-0.1, 1], [-0.2, 2]], + [[-0.3, 3], [-0.4, 4]], + [[-0.5, 5], [-0.6, 6]], + [[-0.7, 7], [-0.8, 8]], + ] + malformed_payloads = { + "missing_meta_info": {}, + "missing_topk_field": {"meta_info": {"prompt_tokens": 5}}, + "wrong_topk_type": {"meta_info": {"prompt_tokens": 5, "input_top_logprobs": "invalid"}}, + "wrong_row_count": {"meta_info": {"prompt_tokens": 5, "input_top_logprobs": valid_rows[:-1]}}, + "ragged_k": {"meta_info": {"prompt_tokens": 5, "input_top_logprobs": [*valid_rows[:-1], [[-0.7, 7]]]}}, + "invalid_token_id": { + "meta_info": { + "prompt_tokens": 5, + "input_top_logprobs": [*valid_rows[:-1], [[-0.7, "7"], [-0.8, 8]]], + } + }, + } - print( - f"old_logprobs: mean_abs={summary['old_logprobs_mean_abs_error']}, " - f"max_abs={summary['old_logprobs_max_abs_error']}\n" - f"advantages: mean_abs={summary['advantages_mean_abs_error']}, " - f"max_abs={summary['advantages_max_abs_error']}\n" - f"propagation: mean_abs={summary['propagation_mean_abs_error']}, " - f"max_abs={summary['propagation_max_abs_error']}\n" - f"result: {result_path}" + for case_name, payload in malformed_payloads.items(): + with self.subTest(case=case_name): + client = self._build_topk_client() + client._client.post = AsyncMock(return_value=self._response(payload)) + + result = await client.compute_logprobs(self._state()) + + self.assertEqual(result.status, Status.FAILED) + self.assertIsNone(result.teacher_tokens) + self.assertIsNone(result.teacher_logprobs) + self.assertIn("last_error=", result.error_msg or "") + client._client.post.assert_awaited_once() + + client = self._build_topk_client() + non_finite_response = httpx.Response( + 200, + request=httpx.Request("POST", "http://teacher/generate"), + headers={"content-type": "application/json"}, + content=( + b'{"meta_info":{"prompt_tokens":5,"input_top_logprobs":' + b"[[[-0.1,1],[-0.2,2]],[[-0.3,3],[-0.4,4]]," + b"[[0.5,5],[-0.6,6]],[[NaN,7],[-0.8,8]]]}}" + ), ) + client._client.post = AsyncMock(return_value=non_finite_response) + result = await client.compute_logprobs(self._state()) -@unittest.skipUnless( - BASELINE_PATH and TEACHER_MODEL_PATH, - "XTUNER_OPD_BASELINE and XTUNER_OPD_TEACHER_MODEL are required", -) -class TestTeacherLogprobClient(unittest.IsolatedAsyncioTestCase): - @classmethod - def setUpClass(cls) -> None: - super().setUpClass() - baseline = torch.load(BASELINE_PATH, map_location="cpu", weights_only=False) - cls.samples = baseline["samples"] - port = find_free_ports()[0] - cls.teacher_endpoint = f"http://127.0.0.1:{port}" - cls.teacher_backend = TeacherLogprobClient._resolve_backend_from_env() - teacher_env = os.environ.copy() - teacher_command = build_teacher_server_command( - OPDTeacherConfig( - name="teacher", - endpoint=cls.teacher_endpoint, - launch_config=OPDTeacherLaunchConfig( - model_path=str(TEACHER_MODEL_PATH), - cuda_visible_devices=teacher_env.get("CUDA_VISIBLE_DEVICES") - or "7", - ), - ), - cls.teacher_backend, - ) - if not teacher_command: - raise ValueError("Teacher must define launch_config for local startup") - cls.teacher_process = subprocess.Popen( - teacher_command, - cwd=REPO_ROOT, - env=teacher_env, - start_new_session=True, + self.assertEqual(result.status, Status.FAILED) + self.assertIn("NaN or Inf", result.error_msg or "") + + async def test_invalid_topk_response_is_retried_before_success(self) -> None: + invalid_response = self._response({"meta_info": {"prompt_tokens": 5}}) + valid_response = self._response( + { + "meta_info": { + "prompt_tokens": 5, + "input_top_logprobs": [ + [[-0.1, 1], [-0.2, 2]], + [[-0.3, 3], [-0.4, 4]], + [[-0.5, 5], [-0.6, 6]], + [[-0.7, 7], [-0.8, 8]], + ], + } + } ) - cls.addClassCleanup(cls._stop_teacher) - _wait_for_teacher(cls.teacher_process, cls.teacher_endpoint, cls.teacher_backend) - - @classmethod - def _stop_teacher(cls) -> None: - if cls.teacher_process.poll() is not None: - return - os.killpg(cls.teacher_process.pid, signal.SIGTERM) - try: - cls.teacher_process.wait(timeout=30) - except subprocess.TimeoutExpired: - os.killpg(cls.teacher_process.pid, signal.SIGKILL) - cls.teacher_process.wait() - - @unittest.skipUnless(os.getenv("XTUNER_USE_SGLANG", "0") == "1", "XTUNER_USE_SGLANG=1 is required") - async def test_compute_logprobs_with_sglang_matches_baseline(self) -> None: - self.assertEqual(self.teacher_backend, "sglang") - await self._assert_compute_logprobs_matches_baseline() - - @unittest.skipUnless(os.getenv("XTUNER_USE_LMDEPLOY", "0") == "1", "XTUNER_USE_LMDEPLOY=1 is required") - async def test_compute_logprobs_with_lmdeploy_matches_baseline(self) -> None: - self.assertEqual(self.teacher_backend, "lmdeploy") - await self._assert_compute_logprobs_matches_baseline() - - async def _assert_compute_logprobs_matches_baseline(self) -> None: - client = TeacherLogprobClient(OPDTeacherConfig(name="teacher", endpoint=self.teacher_endpoint)) - self.addAsyncCleanup(client._client.aclose) + client = self._build_topk_client(max_retry_per_sample=1) + client._client.post = AsyncMock(side_effect=[invalid_response, valid_response]) + + with patch("xtuner.v1.rl.distillation.rollout_teacher_client.asyncio.sleep", new=AsyncMock()): + result = await client.compute_logprobs(self._state()) + + self.assertEqual(result.status, Status.COMPLETED) + self.assertEqual(result.teacher_tokens, [[5, 6], [7, 8]]) + self.assertEqual(result.teacher_logprobs, [[-0.5, -0.6], [-0.7, -0.8]]) + self.assertEqual(client._client.post.await_count, 2) + request_payload = client._client.post.await_args.kwargs["json"] + self.assertEqual(request_payload["input_ids"], [10, 11, 12, 13, 14]) + self.assertEqual(request_payload["top_logprobs_num"], 2) + + +class TestTopKTrainingController(unittest.TestCase): + def test_packs_targets_along_sequence_dimension(self) -> None: + controller = TrainingController(workers=[]) + first = { + "seq_ctx": SequenceContext.from_input_ids((torch.tensor([[1, 2]]),), device="cpu"), + "shifted_labels": torch.tensor([[-100, 2]]), + "advantage": [0.0, 0.0], + "rollout_logprobs": torch.zeros(1, 2), + "teacher_logprobs": torch.tensor([[[-0.1, -0.2], [-0.3, -0.4]]]), + "target_token_ids": torch.tensor([[[1, 2], [3, 4]]]), + } + second = { + "seq_ctx": SequenceContext.from_input_ids((torch.tensor([[3]]),), device="cpu"), + "shifted_labels": torch.tensor([[3]]), + "advantage": [0.0], + "rollout_logprobs": torch.zeros(1, 1), + "teacher_logprobs": torch.tensor([[[-0.5, -0.6]]]), + "target_token_ids": torch.tensor([[[5, 6]]]), + } - for sample in self.samples: - with self.subTest(sample_index=sample["sample_index"]): - prompt_ids = torch.as_tensor(sample["prompt_token_ids"], dtype=torch.long).tolist() - response_ids = torch.as_tensor(sample["response_token_ids"], dtype=torch.long).tolist() - expected_logprobs = torch.as_tensor(sample["teacher_log_probs"], dtype=torch.float32) - state = RolloutState( - rollout_id=int(sample["sample_index"]), - group_id=int(sample["group_index"]), - message=[], - prompt_ids=prompt_ids, - tokens=prompt_ids, - response="", - response_ids=response_ids, - status=Status.COMPLETED, - ) + packed = controller._packing([first, second], pack_max_length=4, language_cfg=None) - result = await client.compute_logprobs(state) - - self.assertEqual(result.status, Status.COMPLETED, result.error_msg) - self.assertEqual(result.teacher_tokens, response_ids) - actual_logprobs = torch.tensor(result.teacher_logprobs, dtype=torch.float32) - self.assertEqual(actual_logprobs.shape, expected_logprobs.shape) - absolute_errors = torch.abs(actual_logprobs - expected_logprobs) - mae = absolute_errors.mean().item() - self.assertLessEqual( - mae, - 0.05, - ( - f"Teacher logprob MAE {mae:.8f} exceeds threshold 0.05; " - f"max_abs_error={absolute_errors.max().item():.8f}" - ), - ) + self.assertEqual(len(packed), 1) + self.assertEqual(packed[0]["teacher_logprobs"].shape, (1, 4, 2)) + self.assertEqual(packed[0]["target_token_ids"].shape, (1, 4, 2)) + torch.testing.assert_close(packed[0]["teacher_logprobs"][0, 3], torch.zeros(2)) + torch.testing.assert_close(packed[0]["target_token_ids"][0, 3], torch.zeros(2, dtype=torch.long)) if __name__ == "__main__": diff --git a/tests/rl/test_prepare_train_data.py b/tests/rl/test_prepare_train_data.py index 0a209df5bb..e4fe72ea2f 100644 --- a/tests/rl/test_prepare_train_data.py +++ b/tests/rl/test_prepare_train_data.py @@ -16,6 +16,8 @@ import torch from xtuner.v1.data_proto.rl_data import RolloutState, Status, reset_rollout_response +from xtuner.v1.rl.distillation import DistillationConfig, RolloutTeacherConfig +from xtuner.v1.rl.loss import DistillationLossConfig from xtuner.v1.train.rl_trainer import BaseRLTrainer @@ -33,7 +35,9 @@ class TestPrepareTrainData(unittest.TestCase): def _build_trainer(self, advantages: list[float]): trainer = BaseRLTrainer.__new__(BaseRLTrainer) trainer._advantage_estimator = _FakeAdvantageEstimator(advantages) - trainer._opd_config = None + trainer._distillation_config = None + trainer._distillation_loss_cfg = None + trainer._train_teacher_config = None trainer.tokenizer = MagicMock(return_value={"input_ids": torch.tensor([[999]])}) trainer.logger = MagicMock() return trainer @@ -54,6 +58,10 @@ def _state( position_ids: np.ndarray | None = None, mm_info: dict | None = None, extra_fields: dict | None = None, + input_ids: list[int] | None = None, + labels: list[int] | None = None, + teacher_tokens: list[int] | list[list[int]] | None = None, + teacher_logprobs: list[float] | list[list[float]] | None = None, ) -> RolloutState: return RolloutState( rollout_id=uid, @@ -71,12 +79,25 @@ def _state( position_ids=position_ids, mm_info=mm_info, extra_fields=extra_fields or {}, + input_ids=input_ids, + labels=labels, + teacher_tokens=teacher_tokens, + teacher_logprobs=teacher_logprobs, ) def _prepare(self, trainer, data_groups, pack_max_length=128): with patch("xtuner.v1.train.rl_trainer.XTUNER_DETERMINISTIC", True): return trainer._prepare_train_data(data_groups, pack_max_length=pack_max_length) + @staticmethod + def _enable_rollout_distillation(trainer, loss_config: DistillationLossConfig) -> None: + trainer._distillation_config = DistillationConfig( + loss_config=loss_config, + teachers=[RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"])], + data_source_teacher_map={"agent_math": "teacher"}, + ) + trainer._distillation_loss_cfg = loss_config + def test_text_path_builds_shifted_training_tensors(self): # 文本主路径固定 token 布局：input_ids 去掉 response 最后一个 token，label/logprob 对齐预测位置。 trainer = self._build_trainer([1.5]) @@ -123,7 +144,7 @@ def test_rerolled_state_without_semantic_mask_uses_all_response_tokens(self): self.assertIsNone(state.response_mask) self.assertEqual(data_batches[0]["shifted_labels"].tolist(), [[-100, -100, 30, 31]]) - self.assertEqual(data_batches[0]["advantage"], [1.0, 1.0, 1.0, 1.0, 1.0]) + self.assertEqual(data_batches[0]["advantage"], [0.0, 0.0, 1.0, 1.0]) def test_multi_sample_group_uses_each_sample_reward_and_advantage(self): # 同一个 prompt 下的多个 response 要分别使用自己的 reward 和 advantage。 @@ -169,6 +190,141 @@ def test_vlm_path_uses_train_prompt_ids_and_preserves_multimodal_fields(self): self.assertEqual(seq_ctx.image_grid_thw.dtype, torch.long) self.assertEqual(seq_ctx.image_grid_thw.tolist(), [[1, 2, 3]]) + def test_agentic_topk_targets_include_token_ids_and_logprobs(self): + loss_config = DistillationLossConfig( + policy_loss_cfg={ + "loss_type": "vanilla", + "cliprange_low": 0.2, + "cliprange_high": 0.2, + }, + loss_mode="reverse", + use_policy_gradient=False, + top_k=2, + ) + trainer = self._build_trainer([0.0]) + self._enable_rollout_distillation(trainer, loss_config) + state = self._state( + input_ids=[10, 11, 20, 21, 22], + labels=[-100, -100, 20, -100, 22], + logprobs=[0.0, -0.1, -0.2, -0.3, -0.4], + teacher_tokens=[[100, 101], [102, 103], [104, 105]], + teacher_logprobs=[[-0.5, -0.6], [-0.7, -0.8], [-0.9, -1.0]], + extra_fields={"origin_data_source": "agent_math"}, + ) + + data_batches, _ = self._prepare(trainer, [[state]]) + + self.assertEqual(len(data_batches), 1) + batch = data_batches[0] + self.assertEqual(batch["shifted_labels"].tolist(), [[-100, 20, -100, 22]]) + self.assertEqual( + batch["target_token_ids"].tolist(), + [[[0, 0], [100, 101], [102, 103], [104, 105]]], + ) + torch.testing.assert_close( + batch["teacher_logprobs"], + torch.tensor( + [[[0.0, 0.0], [-0.5, -0.6], [-0.7, -0.8], [-0.9, -1.0]]], + dtype=torch.float32, + ), + ) + with patch("xtuner.v1.rl.loss.distillation_loss.DEVICE", "cpu"): + loss_ctx = loss_config.build( + { + "shifted_labels": batch["shifted_labels"], + "advantages": torch.tensor([batch["advantage"]], dtype=torch.float32), + "old_logprobs": torch.zeros_like(batch["shifted_labels"], dtype=torch.float32), + "teacher_logprobs": batch["teacher_logprobs"], + "target_token_ids": batch["target_token_ids"], + } + ) + assert loss_ctx is not None + type(loss_ctx).build_batches([loss_ctx]) + loss, _ = loss_ctx.loss_fn( + hidden_states=torch.randn(1, 4, 8), + head_weight=torch.randn(128, 8), + head_bias=None, + loss_kwargs=loss_ctx.loss_kwargs, + ) + self.assertTrue(torch.isfinite(loss)) + + def test_plain_topk_targets_include_masked_response_rows(self): + loss_config = DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="forward_kl_topk", + use_policy_gradient=False, + top_k=2, + ) + trainer = self._build_trainer([0.0]) + self._enable_rollout_distillation(trainer, loss_config) + state = self._state( + prompt_ids=[10, 11, 12], + response_ids=[20, 21, 22], + response_mask=[0, 1, 1], + teacher_tokens=[[100, 101], [102, 103], [104, 105]], + teacher_logprobs=[[-0.5, -0.6], [-0.7, -0.8], [-0.9, -1.0]], + extra_fields={"origin_data_source": "agent_math"}, + ) + + data_batches, _ = self._prepare(trainer, [[state]]) + + batch = data_batches[0] + self.assertEqual(batch["shifted_labels"].tolist(), [[-100, -100, -100, 21, 22]]) + self.assertEqual( + batch["target_token_ids"].tolist(), + [[[0, 0], [0, 0], [100, 101], [102, 103], [104, 105]]], + ) + torch.testing.assert_close( + batch["teacher_logprobs"], + torch.tensor( + [[[0.0, 0.0], [0.0, 0.0], [-0.5, -0.6], [-0.7, -0.8], [-0.9, -1.0]]], + dtype=torch.float32, + ), + ) + + def test_sampled_token_targets_align_for_plain_and_agentic_rollouts(self): + loss_config = DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="k1", + use_policy_gradient=True, + ) + trainer = self._build_trainer([0.0, 0.0]) + self._enable_rollout_distillation(trainer, loss_config) + plain_state = self._state( + uid=1, + group_id=1, + prompt_ids=[10, 11, 12], + response_ids=[20, 21, 22], + response_mask=[0, 1, 1], + teacher_tokens=[20, 21, 22], + teacher_logprobs=[-0.5, -0.7, -0.9], + extra_fields={"origin_data_source": "agent_math"}, + ) + agentic_state = self._state( + uid=2, + group_id=2, + input_ids=[30, 31, 40, 41, 42], + labels=[-100, -100, 40, -100, 42], + logprobs=[0.0, -0.1, -0.2, -0.3, -0.4], + teacher_tokens=[40, 41, 42], + teacher_logprobs=[-1.1, -1.2, -1.3], + extra_fields={"origin_data_source": "agent_math"}, + ) + + data_batches, _ = self._prepare(trainer, [[plain_state], [agentic_state]]) + + self.assertEqual(len(data_batches), 2) + torch.testing.assert_close( + data_batches[0]["teacher_logprobs"], + torch.tensor([[0.0, 0.0, -0.5, -0.7, -0.9]], dtype=torch.float32), + ) + torch.testing.assert_close( + data_batches[1]["teacher_logprobs"], + torch.tensor([[0.0, -1.1, -1.2, -1.3]], dtype=torch.float32), + ) + self.assertNotIn("target_token_ids", data_batches[0]) + self.assertNotIn("target_token_ids", data_batches[1]) + def test_invalid_group_is_skipped(self): # FAILED/FILTERED/ABORTED group 不能进入训练 batch，也不能贡献训练样本数。 trainer = self._build_trainer([1.0]) diff --git a/xtuner/v1/rl/distillation/rollout_teacher_client.py b/xtuner/v1/rl/distillation/rollout_teacher_client.py index 3734c490f5..bc3bff4a48 100644 --- a/xtuner/v1/rl/distillation/rollout_teacher_client.py +++ b/xtuner/v1/rl/distillation/rollout_teacher_client.py @@ -125,6 +125,7 @@ async def compute_logprobs(self, state: RolloutState) -> RolloutState: teacher_tokens, teacher_logprobs = self._parse_topk_response( response, response_ids, + logprob_start_len=logprob_start_len, expanded_prompt_len=expanded_prompt_len, ) else: @@ -288,10 +289,11 @@ def _parse_topk_response( response: httpx.Response, response_ids: list[int], *, + logprob_start_len: int, expanded_prompt_len: int | None = None, ) -> tuple[list[list[int]], list[list[float]]]: - meta_info = response.json()["meta_info"] - prompt_token_count = meta_info["prompt_tokens"] + meta_info = self._get_response_meta_info(response) + prompt_token_count = self._get_prompt_token_count(meta_info) if expanded_prompt_len is not None: expected_prompt_token_count = expanded_prompt_len + len(response_ids) if prompt_token_count != expected_prompt_token_count: @@ -302,11 +304,47 @@ def _parse_topk_response( f"got {prompt_token_count}" ) - response_topk = meta_info["input_top_logprobs"][-len(response_ids) :] - teacher_tokens = [[int(token_id) for _, token_id in row] for row in response_topk] - teacher_logprobs = [[float(logprob) for logprob, _ in row] for row in response_topk] - if not all(math.isfinite(logprob) for row in teacher_logprobs for logprob in row): - raise ValueError("Teacher Top-K logprobs contain NaN or Inf") + raw_topk = self._get_meta_list(meta_info, "input_top_logprobs") + expected_rows = prompt_token_count - logprob_start_len - 1 + if len(raw_topk) != expected_rows: + raise ValueError( + "LMDeploy teacher Top-K length mismatch: " + f"expected {expected_rows} rows for logprob_start_len={logprob_start_len}, got {len(raw_topk)}" + ) + if len(raw_topk) < len(response_ids): + raise ValueError( + "LMDeploy teacher Top-K response is shorter than the scored response: " + f"{len(raw_topk)} vs {len(response_ids)}" + ) + + top_k = cast(int, self.loss_config.top_k) + teacher_tokens: list[list[int]] = [] + teacher_logprobs: list[list[float]] = [] + for row_idx, row in enumerate(raw_topk[-len(response_ids) :]): + if not isinstance(row, list) or len(row) != top_k: + raise ValueError(f"Teacher Top-K row {row_idx} must contain exactly {top_k} entries") + + token_row: list[int] = [] + logprob_row: list[float] = [] + for entry_idx, entry in enumerate(row): + if not isinstance(entry, (list, tuple)) or len(entry) != 2: + raise ValueError(f"Teacher Top-K row {row_idx} entry {entry_idx} must be [logprob, token_id]") + raw_logprob, raw_token_id = entry + if isinstance(raw_logprob, bool) or not isinstance(raw_logprob, (int, float)): + raise ValueError(f"Teacher Top-K row {row_idx} entry {entry_idx} has a non-numeric logprob") + if isinstance(raw_token_id, bool) or not isinstance(raw_token_id, int) or raw_token_id < 0: + raise ValueError( + f"Teacher Top-K row {row_idx} entry {entry_idx} has an invalid non-negative token id" + ) + logprob = float(raw_logprob) + if not math.isfinite(logprob): + raise ValueError("Teacher Top-K logprobs contain NaN or Inf") + logprob_row.append(logprob) + token_row.append(raw_token_id) + if len(token_row) != len(set(token_row)): + raise ValueError(f"Teacher Top-K row {row_idx} contains duplicate token ids") + teacher_tokens.append(token_row) + teacher_logprobs.append(logprob_row) return teacher_tokens, teacher_logprobs @staticmethod @@ -317,8 +355,9 @@ def _parse_sglang_response( logprob_start_len: int, expanded_prompt_len: int | None = None, ) -> list[Any]: - raw_logprobs = RolloutTeacherClient._get_input_token_logprobs(response) - prompt_token_count = response.json()["meta_info"]["prompt_tokens"] + meta_info = RolloutTeacherClient._get_response_meta_info(response) + raw_logprobs = RolloutTeacherClient._get_meta_list(meta_info, "input_token_logprobs") + prompt_token_count = RolloutTeacherClient._get_prompt_token_count(meta_info) if expanded_prompt_len is not None: expected_prompt_token_count = expanded_prompt_len + len(response_ids) if prompt_token_count != expected_prompt_token_count: @@ -347,8 +386,9 @@ def _parse_lmdeploy_response( logprob_start_len: int, expanded_prompt_len: int | None = None, ) -> list[Any]: - raw_logprobs = RolloutTeacherClient._get_input_token_logprobs(response) - prompt_token_count = response.json()["meta_info"]["prompt_tokens"] + meta_info = RolloutTeacherClient._get_response_meta_info(response) + raw_logprobs = RolloutTeacherClient._get_meta_list(meta_info, "input_token_logprobs") + prompt_token_count = RolloutTeacherClient._get_prompt_token_count(meta_info) if expanded_prompt_len is not None: expected_prompt_token_count = expanded_prompt_len + len(response_ids) if prompt_token_count != expected_prompt_token_count: @@ -370,25 +410,46 @@ def _parse_lmdeploy_response( return raw_logprobs[-len(response_ids) :] @staticmethod - def _get_input_token_logprobs(response: httpx.Response) -> list[Any]: + def _get_response_meta_info(response: httpx.Response) -> dict[str, Any]: try: - raw_logprobs = response.json()["meta_info"]["input_token_logprobs"] - except (KeyError, TypeError, ValueError) as exc: + payload = response.json() + except (TypeError, ValueError) as exc: raise ValueError("Invalid teacher response") from exc - if not isinstance(raw_logprobs, list): + if not isinstance(payload, dict) or not isinstance(payload.get("meta_info"), dict): raise ValueError("Invalid teacher response") - return raw_logprobs + return cast(dict[str, Any], payload["meta_info"]) + + @staticmethod + def _get_prompt_token_count(meta_info: dict[str, Any]) -> int: + prompt_token_count = meta_info.get("prompt_tokens") + if isinstance(prompt_token_count, bool) or not isinstance(prompt_token_count, int) or prompt_token_count < 1: + raise ValueError("Invalid teacher response prompt_tokens") + return prompt_token_count + + @staticmethod + def _get_meta_list(meta_info: dict[str, Any], field_name: str) -> list[Any]: + value = meta_info.get(field_name) + if not isinstance(value, list): + raise ValueError(f"Invalid teacher response {field_name}") + return value @staticmethod def _validate_response_logprobs( response_logprobs: list[Any], response_ids: list[int], ) -> tuple[list[int], list[float]]: - try: - teacher_tokens = [item[1] for item in response_logprobs] - teacher_logprobs = [float(item[0]) for item in response_logprobs] - except (TypeError, IndexError, ValueError) as exc: - raise ValueError("Invalid teacher response") from exc + teacher_tokens: list[int] = [] + teacher_logprobs: list[float] = [] + for row_idx, item in enumerate(response_logprobs): + if not isinstance(item, (list, tuple)) or len(item) != 2: + raise ValueError(f"Teacher logprob row {row_idx} must be [logprob, token_id]") + raw_logprob, raw_token_id = item + if isinstance(raw_logprob, bool) or not isinstance(raw_logprob, (int, float)): + raise ValueError(f"Teacher logprob row {row_idx} has a non-numeric logprob") + if isinstance(raw_token_id, bool) or not isinstance(raw_token_id, int) or raw_token_id < 0: + raise ValueError(f"Teacher logprob row {row_idx} has an invalid non-negative token id") + teacher_tokens.append(raw_token_id) + teacher_logprobs.append(float(raw_logprob)) if len(teacher_logprobs) != len(response_ids): raise ValueError("Teacher logprob length mismatch") diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 9661120fcd..11724313b7 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -4,6 +4,7 @@ import random import re import time +from collections.abc import Sequence from dataclasses import asdict, dataclass from pathlib import Path from shutil import rmtree @@ -35,6 +36,7 @@ from xtuner.v1.rl.distillation import DistillationConfig from xtuner.v1.rl.evaluator import EvaluatorConfig from xtuner.v1.rl.health_manager import RLHealthManager, _NoOpRLHealthManager +from xtuner.v1.rl.loss import DistillationLossConfig from xtuner.v1.rl.replay_buffer import ( AsyncReplayBufferConfig, SyncReplayBufferConfig, @@ -83,6 +85,66 @@ def _to_cpu_tensor(value: np.ndarray | None, *, dtype: torch.dtype | None = None return torch.as_tensor(value, dtype=dtype, device="cpu") +def _align_rollout_teacher_targets( + state: RolloutState, + loss_config: DistillationLossConfig, + *, + shifted_labels: Sequence[int], + target_start: int, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Left-pad rollout Teacher targets to align with one training sequence.""" + + sequence_length = len(shifted_labels) + context = f"rollout_id={state.rollout_id}, group_id={state.group_id}" + if target_start < 0 or target_start > sequence_length: + raise ValueError( + f"Teacher target_start must be within the shifted sequence: {target_start} vs {sequence_length}; {context}" + ) + if any(label != loss_config.ignore_idx for label in shifted_labels[:target_start]): + raise ValueError(f"Teacher target prefix must contain only ignored labels; {context}") + + expected_target_rows = sequence_length - target_start + raw_teacher_logprobs = state.teacher_logprobs + if raw_teacher_logprobs is None or len(raw_teacher_logprobs) != expected_target_rows: + actual_rows = None if raw_teacher_logprobs is None else len(raw_teacher_logprobs) + raise ValueError( + "Teacher logprobs must align with the target suffix: " + f"expected {expected_target_rows} rows, got {actual_rows}; {context}" + ) + + if loss_config.uses_sampled_token_targets: + sampled_logprobs = cast(list[float], raw_teacher_logprobs) + teacher_logprobs = torch.tensor( + [0.0] * target_start + sampled_logprobs, + dtype=torch.float32, + ).unsqueeze(0) + return teacher_logprobs, None + + top_k = cast(int, loss_config.top_k) + raw_teacher_tokens = state.teacher_tokens + if raw_teacher_tokens is None or len(raw_teacher_tokens) != expected_target_rows: + actual_rows = None if raw_teacher_tokens is None else len(raw_teacher_tokens) + raise ValueError( + "Teacher token ids must align with the target suffix: " + f"expected {expected_target_rows} rows, got {actual_rows}; {context}" + ) + + topk_tokens = cast(list[list[int]], raw_teacher_tokens) + topk_logprobs = cast(list[list[float]], raw_teacher_logprobs) + + prompt_target_tokens = [[0] * top_k for _ in range(target_start)] + prompt_teacher_logprobs = [[0.0] * top_k for _ in range(target_start)] + teacher_logprobs = torch.tensor( + prompt_teacher_logprobs + topk_logprobs, + dtype=torch.float32, + ).unsqueeze(0) + target_token_ids = torch.tensor( + prompt_target_tokens + topk_tokens, + dtype=torch.int64, + ).unsqueeze(0) + return teacher_logprobs, target_token_ids + + def _agent_loop_manager_requires_rollout_proxy( cfg: AgentLoopManagerConfig | DisaggAgentLoopManagerConfig | None, ) -> bool: @@ -1268,30 +1330,22 @@ def _prepare_train_data( input_ids = raw_input_ids[:-1] shifted_labels = labels[1:] teacher_logprobs = None + target_token_ids = None if ( self._distillation_config is not None and self._distillation_loss_cfg is not None and self._distillation_config.rollout_teachers ): - raw_teacher_logprobs = group[i].teacher_logprobs - if raw_teacher_logprobs is None: - raise ValueError( - f"Teacher logprobs cannot be None when distillation is enabled: {group[i]}" - ) teacher_response_start = next( (index for index, label in enumerate(shifted_labels) if label != -100), len(shifted_labels), ) - teacher_response_length = len(shifted_labels) - teacher_response_start - if len(raw_teacher_logprobs) != teacher_response_length: - raise ValueError( - "Teacher logprobs must align with the trainable suffix of shifted agent labels: " - f"{len(raw_teacher_logprobs)} vs {teacher_response_length}, data: {group[i]}" - ) - teacher_logprobs = torch.tensor( - [0.0] * teacher_response_start + raw_teacher_logprobs, - dtype=torch.float32, - ).unsqueeze(0) + teacher_logprobs, target_token_ids = _align_rollout_teacher_targets( + group[i], + self._distillation_loss_cfg, + shifted_labels=shifted_labels, + target_start=teacher_response_start, + ) prompt_len = sum(label == -100 for label in shifted_labels) response_len = len(shifted_labels) - prompt_len prompt_len_list.append(prompt_len) @@ -1324,6 +1378,8 @@ def _prepare_train_data( } if teacher_logprobs is not None: data_dict["teacher_logprobs"] = teacher_logprobs + if target_token_ids is not None: + data_dict["target_token_ids"] = target_token_ids if teacher_index is not None: data_dict["teacher_indices"] = torch.full_like( shifted_labels_t, @@ -1411,28 +1467,15 @@ def _prepare_train_data( and self._distillation_loss_cfg is not None and self._distillation_config.rollout_teachers ): - if self._distillation_loss_cfg.uses_topk_targets: - top_k = cast(int, self._distillation_loss_cfg.top_k) - response_teacher_tokens = cast(list[list[int]], group[i].teacher_tokens) - response_topk_logprobs = cast(list[list[float]], group[i].teacher_logprobs) - assert len(response_teacher_tokens) == len(response_ids) - assert len(response_topk_logprobs) == len(response_ids) - prompt_target_tokens = [[0] * top_k for _ in range(len(prompt_ids) - 1)] - prompt_teacher_logprobs = [[0.0] * top_k for _ in range(len(prompt_ids) - 1)] - data_dict["target_token_ids"] = torch.tensor( - prompt_target_tokens + response_teacher_tokens, - dtype=torch.int64, - ).unsqueeze(0) - data_dict["teacher_logprobs"] = torch.tensor( - prompt_teacher_logprobs + response_topk_logprobs, - dtype=torch.float32, - ).unsqueeze(0) - else: - response_sampled_logprobs = cast(list[float], group[i].teacher_logprobs) - data_dict["teacher_logprobs"] = torch.tensor( - [0.0] * (len(prompt_ids) - 1) + response_sampled_logprobs, - dtype=torch.float32, - ).unsqueeze(0) + teacher_logprobs, target_token_ids = _align_rollout_teacher_targets( + group[i], + self._distillation_loss_cfg, + shifted_labels=shifted_labels, + target_start=len(prompt_ids) - 1, + ) + data_dict["teacher_logprobs"] = teacher_logprobs + if target_token_ids is not None: + data_dict["target_token_ids"] = target_token_ids if teacher_index is not None: data_dict["teacher_indices"] = torch.full_like( shifted_labels_t, From 5f0bb06f00274c188207b67a9bcc8e219f1588ec Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Sun, 30 Aug 2026 12:49:40 +0000 Subject: [PATCH 05/20] Remove K3 OPD and document supported workflows --- docs/en/rl/advanced_tutorial/index.rst | 3 +- .../on_policy_distillation.md | 528 ++++++++++++++++++ docs/zh_cn/rl/advanced_tutorial/index.rst | 1 + .../on_policy_distillation.md | 527 +++++++++++++++++ xtuner/v1/rl/loss/distillation_loss.py | 15 +- xtuner/v1/rl/trainer/worker.py | 4 +- 6 files changed, 1068 insertions(+), 10 deletions(-) create mode 100644 docs/en/rl/advanced_tutorial/on_policy_distillation.md create mode 100644 docs/zh_cn/rl/advanced_tutorial/on_policy_distillation.md diff --git a/docs/en/rl/advanced_tutorial/index.rst b/docs/en/rl/advanced_tutorial/index.rst index dc4a1eeba6..84d796874b 100644 --- a/docs/en/rl/advanced_tutorial/index.rst +++ b/docs/en/rl/advanced_tutorial/index.rst @@ -7,4 +7,5 @@ Reinforcement Learning model.md dataset.md rl_trainer.md - loss.md \ No newline at end of file + loss.md + on_policy_distillation.md diff --git a/docs/en/rl/advanced_tutorial/on_policy_distillation.md b/docs/en/rl/advanced_tutorial/on_policy_distillation.md new file mode 100644 index 0000000000..8297dbdde7 --- /dev/null +++ b/docs/en/rl/advanced_tutorial/on_policy_distillation.md @@ -0,0 +1,528 @@ +# On-Policy Distillation + +This guide describes the On-Policy Distillation (OPD) capabilities currently supported by XTuner, including the available distillation targets, how to combine OPD with policy gradient (PG), how to configure training-objective weights and multi-Teacher routing, and how to choose between a Rollout Teacher and a Training-side Teacher. + +OPD uses sequences generated by the current Student policy as training data. A Teacher then provides token-level targets for those sequences: + +```text +Student rollout + -> Teacher computes sampled-token or Top-K targets + -> Align targets with training labels + -> PG-OPD or direct distillation + -> Update Student +``` + +In this guide, MOPD means Multi-Teacher OPD: data from different sources can be routed to different Teachers. It does not necessarily mean that task reward is mixed into the objective; whether task reward is used is controlled by `task_adv_weight`. + +## Quick Selection + +First choose the distillation target, then choose where the Teacher runs. These two dimensions can be combined independently. + +### Distillation Targets + +| Target | `loss_mode` | `use_policy_gradient` | `top_k` | Optimization | +| --- | --- | ---: | ---: | --- | +| sampled-token | `k1` | `True` | Unset | Add the distillation signal to the advantage, then compute the PG loss | +| Teacher Top-K | `forward` | `False` | Required | Direct backpropagation through Top-K forward KL with tail mass | +| Teacher Top-K | `reverse` | `False` | Required | Direct backpropagation through Top-K reverse KL with tail mass | +| Teacher Top-K | `forward_kl_topk` | `False` | Required | Direct backpropagation through truncated forward KL over the Teacher Top-K | + +Sampled-token currently supports only `k1 + PG`. `k1` cannot use direct backpropagation, and the three Top-K modes cannot use PG. + +### Teacher Runtime + +| Teacher runtime | sampled-token | Top-K | How the Teacher runs | +| --- | ---: | ---: | --- | +| Rollout Teacher + LMDeploy | Supported | Supported | Independent HTTP service | +| Training-side Teacher | Supported | Supported | Frozen model inside each Training Worker | + +With a Rollout Teacher, XTuner can launch the Teacher service together with Student training, or connect to endpoints that are already running. A Training-side Teacher does not require an endpoint; the training engine loads the Teacher checkpoint directly. + +```{note} +All Teachers in one `DistillationConfig` must use the same runtime type. They must all be `RolloutTeacherConfig` instances or all be `TrainTeacherConfig` instances; the two types cannot be mixed. +``` + +## Core Configuration + +OPD has three configuration layers: + +- `DistillationLossConfig` defines the distillation target, PG loss, and objective weights. +- `RolloutTeacherConfig` or `TrainTeacherConfig` defines where the Teacher runs. +- `DistillationConfig` combines the loss, Teacher list, and data-source routing. + +The same `DistillationLossConfig` must be passed to both the Training Worker and `DistillationConfig`: + +```python +from xtuner.v1.rl.distillation import DistillationConfig +from xtuner.v1.rl.loss import DistillationLossConfig +from xtuner.v1.rl.trainer import WorkerConfig + +loss_cfg = DistillationLossConfig(...) + +train_worker_cfg = WorkerConfig( + ..., + loss_cfg=loss_cfg, +) + +distillation_config = DistillationConfig( + loss_config=loss_cfg, + teachers=[...], + data_source_teacher_map={...}, +) +``` + +The Trainer checks that `train_worker_cfg.loss_cfg == distillation_config.loss_config`. Finally, pass `distillation_config` to `RLColocateTrainerConfig` or the corresponding RL Trainer configuration: + +```python +trainer = RLColocateTrainerConfig( + ..., + train_worker_cfg=train_worker_cfg, + distillation_config=distillation_config, +) +``` + +Common RL loss settings such as `policy_loss_cfg`, `use_kl_loss`, rollout importance sampling, `mode`, and `chunk_size` remain available. See [Loss](loss.md) for the shared fields. In particular: + +- `loss_mode` selects the OPD target. +- `use_kl_loss`, `kl_loss_coef`, and `kl_loss_type` control the standard KL penalty against a Reference Model. +- The OPD Teacher and Reference Model have different roles. Enabling OPD does not require enabling Reference KL. + +## Sampled-Token PG-OPD + +Sampled-token OPD does not ask the Teacher to select new tokens. The Student first generates tokens, and the Teacher computes log probabilities only for those same tokens. + +```text +Student rollout + -> Student sampled tokens + -> Teacher computes log probabilities for the same tokens + -> Align with shifted labels + -> Build the k1 distillation signal + -> Combine with task advantage + -> Policy gradient loss +``` + +The training-target contract is: + +```text +teacher_logprobs: [B, T] +target_token_ids: None +``` + +Let the old Student log probability at each valid token be $\ell_{old}$ and the Teacher log probability be $\ell_T$. The `k1` distillation signal is: + +$$ +d_{OPD} = \ell_{old} - \ell_T +$$ + +XTuner detaches this signal and combines it with the task advantage: + +$$ +A_{combined} += w_{task} A_{task} +- w_{distill} d_{OPD} +$$ + +where: + +- $w_{task}$ corresponds to `task_adv_weight`. +- $w_{distill}$ corresponds to `distillation_loss_weight`. +- The final gradient is produced by the policy loss selected by `policy_loss_cfg`. + +A minimal loss configuration is: + +```python +from xtuner.v1.rl.loss import DistillationLossConfig + +loss_cfg = DistillationLossConfig( + policy_loss_cfg={ + "loss_type": "vanilla", + "cliprange_low": 0.2, + "cliprange_high": 0.28, + "clip_ratio_c": 10.0, + "log_prob_diff_min": -20.0, + "log_prob_diff_max": 20.0, + }, + loss_mode="k1", + use_policy_gradient=True, + task_adv_weight=0.0, + distillation_loss_weight=1.0, + use_kl_loss=False, + mode="chunk", + chunk_size=512, +) +``` + +This example is pure OPD: `task_adv_weight=0.0`, so training samples do not need task rewards. Set it to a positive value such as `1.0` to use task rewards as well. + +## Teacher Top-K OPD + +In Top-K OPD, the Teacher selects K tokens at every training position and provides both their token IDs and log probabilities. The Student computes its own log probabilities over the same token IDs, and XTuner then computes a differentiable distillation loss. + +```text +Teacher forward + -> Teacher Top-K token IDs + log probabilities + -> Gather Student log probabilities at the same token IDs + -> Top-K KL + -> Direct backpropagation +``` + +The training-target contract is: + +```text +target_token_ids: [B, T, K] +teacher_logprobs: [B, T, K] +``` + +Let $S_T$ be the Top-K set selected by the Teacher, and let $q_i$ and $p_i$ be the Teacher and Student probabilities for token $i$, respectively. + +### `forward_kl_topk` + +`forward_kl_topk` computes a truncated forward KL term only over the K tokens selected by the Teacher: + +$$ +L_{topk} = \sum_{i \in S_T} q_i(\log q_i - \log p_i) +$$ + +It does not combine probability mass outside the Top-K set into an extra tail bin. `log_prob_min_clamp` can bound Student and Teacher log probabilities from below before computation, while `loss_max_clamp` can bound the per-token loss. + +### `forward` + +In addition to the Teacher Top-K terms, `forward` combines the remaining probability mass into a tail bin: + +$$ +q_{tail}=1-\sum_{i\in S_T}q_i,\qquad +p_{tail}=1-\sum_{i\in S_T}p_i +$$ + +$$ +L_{forward} += \sum_{i\in S_T}q_i(\log q_i-\log p_i) ++q_{tail}(\log q_{tail}-\log p_{tail}) +$$ + +### `reverse` + +`reverse` weights the opposite direction with Student probabilities and also preserves the tail bin: + +$$ +L_{reverse} += \sum_{i\in S_T}p_i(\log p_i-\log q_i) ++p_{tail}(\log p_{tail}-\log q_{tail}) +$$ + +We recommend starting with `forward_kl_topk`, which matches the common Teacher Top-K training semantics: + +```python +loss_cfg = DistillationLossConfig( + policy_loss_cfg={ + "loss_type": "vanilla", + "cliprange_low": 0.2, + "cliprange_high": 0.28, + }, + loss_mode="forward_kl_topk", + use_policy_gradient=False, + top_k=64, + log_prob_min_clamp=-10.0, + loss_max_clamp=10.0, + task_adv_weight=0.0, + distillation_loss_weight=1.0, + use_kl_loss=False, + mode="chunk", + chunk_size=512, +) +``` + +Top-K modes do not place the distillation term in the advantage. Instead, they compute: + +$$ +L_{total} += L_{PG}(w_{task}A_{task}) ++w_{distill}L_{topk} +$$ + +Therefore, even when `use_policy_gradient=False`, `policy_loss_cfg` is still used for the optional task loss. The flag only means that the OPD term itself uses direct backpropagation. + +## Configuring Teachers + +### Let XTuner Launch the Rollout Teacher + +Use `RolloutTeacherConfig.launch_config` to describe the Teacher service that XTuner should launch: + +```python +from xtuner.v1.rl.distillation import ( + DistillationConfig, + RolloutTeacherConfig, + RolloutTeacherLaunchConfig, +) + +distillation_config = DistillationConfig( + loss_config=loss_cfg, + teachers=[ + RolloutTeacherConfig( + name="teacher", + num_replicas=1, + enable_prefix_caching=True, + launch_config=RolloutTeacherLaunchConfig( + model_path=teacher_model_path, + num_workers=1, + server_port=13141, + dtype="bfloat16", + tensor_parallel_size=1, + expert_parallel_size=1, + context_length=max_prompt_length + max_response_length, + max_batch_size=2048, + gpu_memory_utilization=0.8, + ), + ) + ], + data_source_teacher_map={"math": "teacher"}, +) +``` + +The OPD launch utility in this repository: + +1. Reads `launch_config`. +2. Allocates GPUs and ports to each Teacher replica. +3. Reserves the remaining GPUs for Student workers. +4. Starts LMDeploy Teacher services and waits for their health checks. +5. Passes the actual endpoint map to the Trainer through `XTUNER_OPD_TEACHER_ENDPOINTS_JSON`. +6. Starts Student training and stops the Teacher processes launched by the utility when training exits. + +Each Teacher replica must fit entirely on one node. The launcher allocates devices according to `num_workers` and rejects configurations in which Teachers occupy every GPU and leave no GPU for the Student. Keep `num_workers`, TP, EP, and the backend parallel topology consistent. + +For an LMDeploy Rollout Teacher, set: + +```bash +export XTUNER_USE_LMDEPLOY=1 +``` + +LMDeploy Rollout Teachers support both sampled-token and Top-K targets. + +### Connect to Existing Endpoints + +If an external system manages the Teacher services, omit `launch_config` and provide `endpoints` directly: + +```python +distillation_config = DistillationConfig( + loss_config=loss_cfg, + teachers=[ + RolloutTeacherConfig( + name="teacher", + num_replicas=2, + endpoints=[ + "http://teacher-0:8000", + "http://teacher-1:8000", + ], + api_key=None, + request_timeout_s=1200.0, + max_retry_per_sample=2, + max_concurrency=128, + ) + ], + data_source_teacher_map={"math": "teacher"}, +) +``` + +In this mode, XTuner does not start or stop the external services. It only performs health checks, request routing, concurrency limiting, retries, and response parsing. The number of `endpoints` should match `num_replicas`. Each endpoint must be the service root URL without `/generate`. + +Multiple replicas of the same Teacher are routed deterministically using the Teacher name, data source, and rollout group. If a request to one replica fails, the client tries subsequent replicas according to `max_retry_per_sample`. + +### Let the Training Engine Load the Teacher + +Use `TrainTeacherConfig` for a Training-side Teacher: + +```python +from pathlib import Path + +from xtuner.v1.model import get_model_config_from_hf +from xtuner.v1.rl.distillation import DistillationConfig, TrainTeacherConfig + +teacher_model_cfg = get_model_config_from_hf(Path(teacher_model_path)) +teacher_model_cfg.compile_cfg = False + +distillation_config = DistillationConfig( + loss_config=loss_cfg, + teachers=[ + TrainTeacherConfig( + name="teacher", + model_path=teacher_model_path, + model_cfg=teacher_model_cfg, + # When fsdp_cfg is omitted, XTuner creates a default + # configuration with requires_grad=False. + fsdp_cfg=None, + ) + ], + data_source_teacher_map={"math": "teacher"}, +) +``` + +`model_cfg` must match the Teacher checkpoint, and it is recommended to use a configuration object separate from the Student configuration. The Teacher and Student may have different layer counts or hidden sizes, but they need compatible token-ID semantics. Current configuration validation requires them to have the same vocabulary size. + +During Training Worker initialization, XTuner constructs the frozen Teacher and keeps it on CPU. Device transitions for each training phase follow this order: + +```text +Actor and optimizer on GPU + -> Offload actor and optimizer to CPU + -> Onload Teacher to the training device + -> Teacher computes targets + -> Offload Teacher to CPU + -> Restore actor and optimizer + -> Compute old log probabilities and training loss +``` + +A Training-side Teacher does not create an HTTP service and does not need `XTUNER_OPD_TEACHER_ENDPOINTS_JSON`. CPU offload is currently supported; Teacher disk offload is not. + +## MOPD and Data-Source Routing + +`data_source_teacher_map` selects a Teacher using each training sample's `origin_data_source`: + +```python +distillation_config = DistillationConfig( + loss_config=loss_cfg, + teachers=[ + RolloutTeacherConfig(name="math_teacher", ...), + RolloutTeacherConfig(name="code_teacher", ...), + ], + data_source_teacher_map={ + "math": "math_teacher", + "code": "code_teacher", + }, +) +``` + +Teacher names must be unique, and every Teacher referenced by the map must exist. Each data source participating in training should have a mapping; otherwise the Teacher-routing stage cannot generate the distillation targets required by that data. + +### MOPD Behavior with Rollout Teachers + +The rollout path completes routing before issuing an HTTP request: + +```text +origin_data_source + -> teacher_name + -> matching RolloutTeacherClient + -> Request only this Teacher +``` + +A sample does not request every Rollout Teacher. + +### MOPD Behavior with Training-Side Teachers + +The training engine first converts the `origin_data_source` of each sample into `teacher_indices`. The current implementation uses Teacher-major FSDP scheduling: + +```text +for each Teacher: + onload Teacher + for each local pack: + run the Teacher forward pass over the entire pack + select tokens assigned to this Teacher using teacher_indices + write only those positions into the final targets + offload Teacher +``` + +Consequently, current Training-side MOPD runs every Teacher sequentially over every local pack and then selects results by token. It does not run only the routed Teacher for each sample. Both sampled-token and Top-K targets use this schedule. + +This schedule ensures that every FSDP rank enters collectives in the same Teacher and pack order. Even if one rank has no tokens assigned to a particular Teacher, that rank cannot skip the Teacher forward pass independently, because ranks could otherwise enter different collective sequences. + +The direct consequence is that Teacher forward computation grows as the number of Teachers increases. However, Teachers are onloaded and offloaded one at a time, so they do not reside on the training GPU simultaneously. + +## Training-Objective Weights + +OPD uses `task_adv_weight` and `distillation_loss_weight` to control the relative contributions of task reward and the distillation target. + +### `task_adv_weight` + +This value scales the advantage produced by task rewards: + +- `0.0`: pure distillation; training does not depend on task rewards. +- `1.0`: use the full task advantage. +- Any other non-negative value: scale the task objective proportionally. + +### `distillation_loss_weight` + +This value scales the OPD target: + +- sampled-token: scales the `k1` distillation signal added to the advantage. +- Top-K: scales the directly backpropagated Top-K distillation loss. + +Common combinations are: + +| Training objective | `task_adv_weight` | `distillation_loss_weight` | +| --- | ---: | ---: | +| Pure OPD | `0.0` | `1.0` | +| Equal task reward and OPD weights | `1.0` | `1.0` | +| Task reward as the primary objective, OPD as an auxiliary objective | `1.0` | `0.1` | + +These values are not universally optimal. When tuning them, monitor reward, distillation metrics, policy ratio and clipping metrics, and overall training stability together. + +## Launch Examples + +The repository provides sampled-token OPD and MOPD recipes: + +- `recipe/on_policy_distillation/config/rl_dapo_math_opd.py` +- `recipe/on_policy_distillation/config/rl_dapo_math_mopd.py` +- `recipe/on_policy_distillation/scripts/run_sampled_token_opd.sh` +- `recipe/on_policy_distillation/scripts/run_sampled_token_mopd.sh` + +### Single-Teacher Sampled-Token OPD + +```bash +export XTUNER_USE_LMDEPLOY=1 +export WORK_DIR=work_dirs/dapo_math_opd + +bash recipe/on_policy_distillation/scripts/run_sampled_token_opd.sh \ + /path/to/student \ + /path/to/teacher \ + /path/to/train.jsonl +``` + +### Multi-Teacher Sampled-Token MOPD + +```bash +export XTUNER_USE_LMDEPLOY=1 +export GSM8K_TEACHER_MODEL_PATH=/path/to/gsm8k_teacher +export GEO3K_TEACHER_MODEL_PATH=/path/to/geo3k_teacher +export WORK_DIR=work_dirs/dapo_math_mopd + +bash recipe/on_policy_distillation/scripts/run_sampled_token_mopd.sh \ + /path/to/student \ + /path/to/dataset_meta.json +``` + +If the configuration uses external endpoints, the launcher does not create the corresponding Teacher processes. It still waits for endpoint health checks and passes the endpoint map to the Trainer. + +A Training-side Teacher does not need `start_teacher_servers`. Add `TrainTeacherConfig` to a standard RL configuration and use the project's normal RL launch method; the Teacher is initialized together with the Training Worker. + +## Metrics + +Common sampled-token metrics include: + +- `opd_reverse_kl`: mean `old_student_logprob - teacher_logprob` over valid tokens. +- `opd_abs_logprob_loss`: mean absolute log-probability difference. +- `reduced_distillation_kl` and `reduced_distillation_abs_loss`: unified distillation metrics. +- `teacher_score_time_s`: time spent by a Rollout Teacher scoring one state, stored in an extra field of the rollout state. + +Common Top-K metrics include: + +- `reduced_topk_opd_loss`: weighted Top-K distillation loss for the current batch. +- `reduced_topk_opd_kl`: mean Top-K KL over valid tokens. +- `reduced_topk_opd_student_selected_mass`: Student probability mass over the Teacher Top-K support. +- `reduced_topk_opd_teacher_selected_mass`: Teacher Top-K probability mass. +- `reduced_topk_opd_overlap_fraction`: mean overlap between the Student and Teacher Top-K sets. + +Training-side Teachers also record: + +- `time/train_teacher_compute` +- `time/train_teacher_onload` +- `time/train_teacher_offload` + +Final distillation metrics are written under the `distillation/` prefix. Raw step averages from individual workers are also written under `train_metrics/worker_/`. + +## Current Limitations + +- Sampled-token supports only `k1 + Policy Gradient`. +- Rollout Top-K supports only LMDeploy. +- A single `DistillationConfig` cannot mix Rollout Teachers and Training-side Teachers. +- Training-side MOPD runs every Teacher sequentially over every local pack. +- Training-side Teachers support CPU offload but not disk offload. +- Current Top-K distillation uses Teacher-selected Top-K targets, not full-vocabulary distillation. +- A Reference Model and an OPD Teacher have different roles. Their configuration semantics remain different even if they use the same checkpoint. diff --git a/docs/zh_cn/rl/advanced_tutorial/index.rst b/docs/zh_cn/rl/advanced_tutorial/index.rst index d5c53b812c..60b48de72e 100644 --- a/docs/zh_cn/rl/advanced_tutorial/index.rst +++ b/docs/zh_cn/rl/advanced_tutorial/index.rst @@ -10,3 +10,4 @@ rl_trainer.md judger.md loss.md + on_policy_distillation.md diff --git a/docs/zh_cn/rl/advanced_tutorial/on_policy_distillation.md b/docs/zh_cn/rl/advanced_tutorial/on_policy_distillation.md new file mode 100644 index 0000000000..e5315895c7 --- /dev/null +++ b/docs/zh_cn/rl/advanced_tutorial/on_policy_distillation.md @@ -0,0 +1,527 @@ +# On-Policy Distillation + +本文介绍 XTuner 当前的 On-Policy Distillation（OPD）能力，包括支持的蒸馏目标、如何与 policy gradient（PG）组合、如何配置训练目标权重和多 Teacher 路由，以及如何选择 Rollout Teacher 或 Training-side Teacher。 + +OPD 使用当前 Student 策略生成的序列作为训练数据，再让 Teacher 对这些序列提供 token-level 目标： + +```text +Student rollout + -> Teacher 计算 sampled-token 或 Top-K 目标 + -> 目标与训练 labels 对齐 + -> PG-OPD 或直接蒸馏 + -> 更新 Student +``` + +本文中的 MOPD 指 Multi-Teacher OPD：不同数据源可以路由到不同 Teacher。它不表示一定混合 task reward；是否使用 task reward 由 `task_adv_weight` 决定。 + +## 快速选择 + +先选择蒸馏目标，再选择 Teacher 的执行位置。这两个维度可以独立组合。 + +### 蒸馏目标 + +| 目标 | `loss_mode` | `use_policy_gradient` | `top_k` | 优化方式 | +| --- | --- | ---: | ---: | --- | +| sampled-token | `k1` | `True` | 不设置 | 蒸馏信号合入 advantage，再计算 PG loss | +| Teacher Top-K | `forward` | `False` | 必填 | Top-K forward KL（含 tail mass）直接反传 | +| Teacher Top-K | `reverse` | `False` | 必填 | Top-K reverse KL（含 tail mass）直接反传 | +| Teacher Top-K | `forward_kl_topk` | `False` | 必填 | Teacher Top-K 上的截断 forward KL 直接反传 | + +sampled-token 当前只支持 `k1 + PG`。`k1` 不能走直接反传；三个 Top-K mode 不能走 PG。 + +### Teacher 执行位置 + +| Teacher runtime | sampled-token | Top-K | Teacher 如何运行 | +| --- | ---: | ---: | --- | +| Rollout Teacher + LMDeploy | 支持 | 支持 | 独立 HTTP 服务 | +| Training-side Teacher | 支持 | 支持 | Training Worker 内的 frozen model | + +使用 Rollout Teacher 时，可以让 XTuner 的启动脚本同时启动 Teacher 服务和 Student 训练，也可以直接传入已经运行的 endpoint。使用 Training-side Teacher 时不需要 Teacher endpoint，Teacher checkpoint 由训练引擎加载。 + +```{note} +同一个 `DistillationConfig` 中的所有 Teacher 必须使用同一种 runtime：只能全部是 `RolloutTeacherConfig`，或者全部是 `TrainTeacherConfig`，不能混用。 +``` + +## 核心配置 + +OPD 由三个配置层组成： + +- `DistillationLossConfig`：定义蒸馏目标、PG loss 和目标权重。 +- `RolloutTeacherConfig` 或 `TrainTeacherConfig`：定义 Teacher 在哪里执行。 +- `DistillationConfig`：组合 loss、Teacher 列表和数据源路由。 + +同一个 `DistillationLossConfig` 需要同时传给 Training Worker 和 `DistillationConfig`： + +```python +from xtuner.v1.rl.distillation import DistillationConfig +from xtuner.v1.rl.loss import DistillationLossConfig +from xtuner.v1.rl.trainer import WorkerConfig + +loss_cfg = DistillationLossConfig(...) + +train_worker_cfg = WorkerConfig( + ..., + loss_cfg=loss_cfg, +) + +distillation_config = DistillationConfig( + loss_config=loss_cfg, + teachers=[...], + data_source_teacher_map={...}, +) +``` + +Trainer 会检查 `train_worker_cfg.loss_cfg == distillation_config.loss_config`。最后把 `distillation_config` 传给 `RLColocateTrainerConfig` 或相应的 RL Trainer 配置： + +```python +trainer = RLColocateTrainerConfig( + ..., + train_worker_cfg=train_worker_cfg, + distillation_config=distillation_config, +) +``` + +`policy_loss_cfg`、`use_kl_loss`、rollout importance sampling、`mode` 和 `chunk_size` 等公共 RL loss 配置仍然有效。公共字段可参考[损失函数](loss.md)。其中： + +- `loss_mode` 控制 OPD 目标。 +- `use_kl_loss`、`kl_loss_coef` 和 `kl_loss_type` 控制相对于 Reference Model 的普通 KL penalty。 +- OPD Teacher 和 Reference Model 是两个不同角色；开启 OPD 不要求开启 Reference KL。 + +## Sampled-token PG-OPD + +sampled-token OPD 不让 Teacher 重新选 token。Student 先生成 token，Teacher 只计算这些 token 的 logprob。 + +```text +Student rollout + -> Student sampled tokens + -> Teacher 对相同 token 计算 logprob + -> 对齐到 shifted labels + -> 构造 k1 蒸馏信号 + -> 合入 task advantage + -> Policy gradient loss +``` + +它的训练目标契约是： + +```text +teacher_logprobs: [B, T] +target_token_ids: None +``` + +令每个有效 token 上的旧 Student logprob 为 $\ell_{old}$，Teacher logprob 为 $\ell_T$，则 `k1` 蒸馏信号为： + +$$ +d_{OPD} = \ell_{old} - \ell_T +$$ + +XTuner 将它 detach 后与 task advantage 组合： + +$$ +A_{combined} += w_{task} A_{task} +- w_{distill} d_{OPD} +$$ + +其中： + +- $w_{task}$ 对应 `task_adv_weight`。 +- $w_{distill}$ 对应 `distillation_loss_weight`。 +- 最终梯度由 `policy_loss_cfg` 指定的 policy loss 产生。 + +最小 loss 配置如下： + +```python +from xtuner.v1.rl.loss import DistillationLossConfig + +loss_cfg = DistillationLossConfig( + policy_loss_cfg={ + "loss_type": "vanilla", + "cliprange_low": 0.2, + "cliprange_high": 0.28, + "clip_ratio_c": 10.0, + "log_prob_diff_min": -20.0, + "log_prob_diff_max": 20.0, + }, + loss_mode="k1", + use_policy_gradient=True, + task_adv_weight=0.0, + distillation_loss_weight=1.0, + use_kl_loss=False, + mode="chunk", + chunk_size=512, +) +``` + +上例是纯 OPD：`task_adv_weight=0.0`，因此训练样本可以没有 task reward。如果要同时使用 task reward，将它设为正数，例如 `1.0`。 + +## Teacher Top-K OPD + +Top-K OPD 由 Teacher 在每个训练位置选择 K 个 token，并同时提供这些 token 的 id 和 logprob。Student 在同一组 token id 上计算自己的 logprob，再计算可微的蒸馏 loss。 + +```text +Teacher forward + -> Teacher Top-K token ids + logprobs + -> Student 在相同 token ids 上 gather logprobs + -> Top-K KL + -> 直接反向传播 +``` + +它的训练目标契约是： + +```text +target_token_ids: [B, T, K] +teacher_logprobs: [B, T, K] +``` + +设 Teacher 选择的 Top-K 集合为 $S_T$，Teacher 和 Student 在 token $i$ 上的概率分别为 $q_i$ 和 $p_i$。 + +### `forward_kl_topk` + +`forward_kl_topk` 只对 Teacher 选出的 K 个 token 求截断 forward KL 项： + +$$ +L_{topk} = \sum_{i \in S_T} q_i(\log q_i - \log p_i) +$$ + +它不把 Top-K 之外的概率质量合并成额外 tail bin。`log_prob_min_clamp` 可以在计算前限制 Student/Teacher logprob 下界，`loss_max_clamp` 可以限制逐 token loss 的范围。 + +### `forward` + +`forward` 在 Teacher Top-K 项之外，把剩余概率质量作为一个 tail bin： + +$$ +q_{tail}=1-\sum_{i\in S_T}q_i,\qquad +p_{tail}=1-\sum_{i\in S_T}p_i +$$ + +$$ +L_{forward} += \sum_{i\in S_T}q_i(\log q_i-\log p_i) ++q_{tail}(\log q_{tail}-\log p_{tail}) +$$ + +### `reverse` + +`reverse` 使用 Student 概率对相反方向加权，并同样保留 tail bin： + +$$ +L_{reverse} += \sum_{i\in S_T}p_i(\log p_i-\log q_i) ++p_{tail}(\log p_{tail}-\log q_{tail}) +$$ + +推荐从与常见 Teacher Top-K 训练语义一致的 `forward_kl_topk` 开始： + +```python +loss_cfg = DistillationLossConfig( + policy_loss_cfg={ + "loss_type": "vanilla", + "cliprange_low": 0.2, + "cliprange_high": 0.28, + }, + loss_mode="forward_kl_topk", + use_policy_gradient=False, + top_k=64, + log_prob_min_clamp=-10.0, + loss_max_clamp=10.0, + task_adv_weight=0.0, + distillation_loss_weight=1.0, + use_kl_loss=False, + mode="chunk", + chunk_size=512, +) +``` + +Top-K 模式不把蒸馏项放进 advantage，而是计算： + +$$ +L_{total} += L_{PG}(w_{task}A_{task}) ++w_{distill}L_{topk} +$$ + +因此，即使 `use_policy_gradient=False`，`policy_loss_cfg` 仍用于可选的 task loss；它只表示 OPD 项本身使用直接反传。 + +## 配置 Teacher + +### 让 XTuner 同时启动 Rollout Teacher + +使用 `RolloutTeacherConfig.launch_config` 描述需要启动的 Teacher 服务： + +```python +from xtuner.v1.rl.distillation import ( + DistillationConfig, + RolloutTeacherConfig, + RolloutTeacherLaunchConfig, +) + +distillation_config = DistillationConfig( + loss_config=loss_cfg, + teachers=[ + RolloutTeacherConfig( + name="teacher", + num_replicas=1, + enable_prefix_caching=True, + launch_config=RolloutTeacherLaunchConfig( + model_path=teacher_model_path, + num_workers=1, + server_port=13141, + dtype="bfloat16", + tensor_parallel_size=1, + expert_parallel_size=1, + context_length=max_prompt_length + max_response_length, + max_batch_size=2048, + gpu_memory_utilization=0.8, + ), + ) + ], + data_source_teacher_map={"math": "teacher"}, +) +``` + +仓库中的 OPD 启动工具会： + +1. 读取 `launch_config`。 +2. 为各个 Teacher replica 分配 GPU 和端口。 +3. 保留剩余 GPU 给 Student workers。 +4. 启动 LMDeploy Teacher 服务并等待健康检查。 +5. 通过 `XTUNER_OPD_TEACHER_ENDPOINTS_JSON` 把实际 endpoint map 传给 Trainer。 +6. 启动 Student 训练，并在退出时关闭由脚本启动的 Teacher。 + +每个 Teacher replica 必须完整放在一个节点上。启动工具按 `num_workers` 为 replica 分配卡，并拒绝“Teacher 占满所有 GPU、没有 Student GPU”的配置。`num_workers`、TP、EP 和后端的并行拓扑应保持一致。 + +使用 LMDeploy Rollout Teacher 时设置： + +```bash +export XTUNER_USE_LMDEPLOY=1 +``` + +sampled-token 和 Top-K 都支持 LMDeploy Rollout Teacher。 + +### 使用已经启动好的 endpoint + +如果 Teacher 服务由外部系统管理，不设置 `launch_config`，直接填写 `endpoints`： + +```python +distillation_config = DistillationConfig( + loss_config=loss_cfg, + teachers=[ + RolloutTeacherConfig( + name="teacher", + num_replicas=2, + endpoints=[ + "http://teacher-0:8000", + "http://teacher-1:8000", + ], + api_key=None, + request_timeout_s=1200.0, + max_retry_per_sample=2, + max_concurrency=128, + ) + ], + data_source_teacher_map={"math": "teacher"}, +) +``` + +此时 XTuner 不启动或关闭外部服务，只执行健康检查、请求路由、并发限制、失败重试和响应解析。`endpoints` 数量应与 `num_replicas` 一致，每个 endpoint 填写服务根地址，不包含 `/generate`。 + +同一个 Teacher 的多个 replica 会根据 Teacher 名称、数据源和 rollout group 稳定路由；某个 replica 请求失败时，客户端会按照 `max_retry_per_sample` 尝试后续 replica。 + +### 由训练引擎加载 Teacher + +Training-side Teacher 使用 `TrainTeacherConfig`： + +```python +from pathlib import Path + +from xtuner.v1.model import get_model_config_from_hf +from xtuner.v1.rl.distillation import DistillationConfig, TrainTeacherConfig + +teacher_model_cfg = get_model_config_from_hf(Path(teacher_model_path)) +teacher_model_cfg.compile_cfg = False + +distillation_config = DistillationConfig( + loss_config=loss_cfg, + teachers=[ + TrainTeacherConfig( + name="teacher", + model_path=teacher_model_path, + model_cfg=teacher_model_cfg, + # 不传 fsdp_cfg 时，XTuner 会创建 requires_grad=False 的默认配置。 + fsdp_cfg=None, + ) + ], + data_source_teacher_map={"math": "teacher"}, +) +``` + +`model_cfg` 必须与 Teacher checkpoint 匹配，并建议使用独立于 Student 的配置对象。Teacher 与 Student 可以有不同层数或 hidden size，但需要使用兼容的 token id 语义；当前配置校验要求二者 vocabulary size 相同。 + +Training Worker 初始化时会构建 frozen Teacher 并将其留在 CPU。每个训练阶段的设备切换顺序是： + +```text +Actor 和 optimizer 在 GPU + -> Actor 和 optimizer offload 到 CPU + -> Teacher onload 到训练设备 + -> Teacher 计算 targets + -> Teacher offload 到 CPU + -> 恢复 Actor 和 optimizer + -> 计算 old logprobs 和训练 loss +``` + +Training-side Teacher 不创建 HTTP 服务，也不需要 `XTUNER_OPD_TEACHER_ENDPOINTS_JSON`。当前支持 CPU offload，不支持 Teacher disk offload。 + +## MOPD 与数据源路由 + +`data_source_teacher_map` 根据训练样本的 `origin_data_source` 选择 Teacher： + +```python +distillation_config = DistillationConfig( + loss_config=loss_cfg, + teachers=[ + RolloutTeacherConfig(name="math_teacher", ...), + RolloutTeacherConfig(name="code_teacher", ...), + ], + data_source_teacher_map={ + "math": "math_teacher", + "code": "code_teacher", + }, +) +``` + +Teacher 名称必须唯一，映射中的 Teacher 必须存在。每个参与训练的数据源都应有对应映射，否则 Teacher 路由阶段无法生成该数据所需的蒸馏目标。 + +### Rollout Teacher 的 MOPD 行为 + +Rollout 路径在发起 HTTP 请求前完成路由： + +```text +origin_data_source + -> teacher_name + -> 对应 RolloutTeacherClient + -> 只请求这个 Teacher +``` + +一条数据不会请求所有 Rollout Teacher。 + +### Training-side Teacher 的 MOPD 行为 + +训练引擎首先把每条数据的 `origin_data_source` 转换为 `teacher_indices`。当前实现采用 Teacher-major 的 FSDP 调度： + +```text +for each Teacher: + onload Teacher + for each local pack: + Teacher 对整个 pack 执行 forward + 根据 teacher_indices 选择属于当前 Teacher 的 token + 只把这些位置写入最终 targets + offload Teacher +``` + +因此，Training-side MOPD 当前是“所有 Teacher 依次对所有本地 pack 做完整 forward，再按 token 选择结果”，而不是“每条数据只执行对应 Teacher”。sampled-token 和 Top-K 都遵循这一调度。 + +这样做是为了保证所有 FSDP rank 按相同的 Teacher、pack 顺序进入 collective；即使某个 rank 没有属于某个 Teacher 的 token，也不能单独跳过该 Teacher forward，否则各 rank 可能进入不同 collective 序列。 + +它的直接影响是：Teacher 数量增加时，Teacher forward 计算量也会增加；但 Teacher 会逐个 onload/offload，不会同时驻留训练 GPU。 + +## 训练目标权重 + +OPD 使用 `task_adv_weight` 和 `distillation_loss_weight` 控制 task reward 与蒸馏目标的相对贡献。 + +### `task_adv_weight` + +它缩放 task reward 产生的 advantage： + +- `0.0`：纯蒸馏，训练不依赖 task reward。 +- `1.0`：完整使用 task advantage。 +- 其他非负数：按比例调整 task 目标的贡献。 + +### `distillation_loss_weight` + +它缩放 OPD 目标： + +- sampled-token：缩放合入 advantage 的 `k1` 蒸馏信号。 +- Top-K：缩放直接反传的 Top-K distillation loss。 + +常见组合如下： + +| 训练目标 | `task_adv_weight` | `distillation_loss_weight` | +| --- | ---: | ---: | +| 纯 OPD | `0.0` | `1.0` | +| Task reward 与 OPD 等权 | `1.0` | `1.0` | +| Task reward 为主、OPD 为辅助 | `1.0` | `0.1` | + +这些数值不是通用最优值。调整时应同时观察 reward、蒸馏指标、policy ratio/clip 指标和训练稳定性。 + +## 启动示例 + +仓库提供 sampled-token OPD 和 MOPD recipe： + +- `recipe/on_policy_distillation/config/rl_dapo_math_opd.py` +- `recipe/on_policy_distillation/config/rl_dapo_math_mopd.py` +- `recipe/on_policy_distillation/scripts/run_sampled_token_opd.sh` +- `recipe/on_policy_distillation/scripts/run_sampled_token_mopd.sh` + +### 单 Teacher sampled-token OPD + +```bash +export XTUNER_USE_LMDEPLOY=1 +export WORK_DIR=work_dirs/dapo_math_opd + +bash recipe/on_policy_distillation/scripts/run_sampled_token_opd.sh \ + /path/to/student \ + /path/to/teacher \ + /path/to/train.jsonl +``` + +### 多 Teacher sampled-token MOPD + +```bash +export XTUNER_USE_LMDEPLOY=1 +export GSM8K_TEACHER_MODEL_PATH=/path/to/gsm8k_teacher +export GEO3K_TEACHER_MODEL_PATH=/path/to/geo3k_teacher +export WORK_DIR=work_dirs/dapo_math_mopd + +bash recipe/on_policy_distillation/scripts/run_sampled_token_mopd.sh \ + /path/to/student \ + /path/to/dataset_meta.json +``` + +如果配置使用外部 endpoints，上述启动工具不会创建对应 Teacher 进程，但仍会等待 endpoint 健康检查并把 endpoint map 传给 Trainer。 + +Training-side Teacher 不需要 `start_teacher_servers`。在普通 RL 配置中加入 `TrainTeacherConfig` 后，使用项目标准 RL 启动方式即可；Teacher 会随 Training Worker 初始化。 + +## 指标 + +sampled-token 常用指标包括： + +- `opd_reverse_kl`：有效 token 上的平均 `old_student_logprob - teacher_logprob`。 +- `opd_abs_logprob_loss`：两者 logprob 差值的平均绝对值。 +- `reduced_distillation_kl`、`reduced_distillation_abs_loss`：统一蒸馏指标。 +- `teacher_score_time_s`：Rollout Teacher 对单条状态评分的耗时，保存在 rollout 状态的额外字段中。 + +Top-K 常用指标包括： + +- `reduced_topk_opd_loss`：当前 batch 的加权 Top-K 蒸馏 loss。 +- `reduced_topk_opd_kl`：有效 token 上的平均 Top-K KL。 +- `reduced_topk_opd_student_selected_mass`：Student 在 Teacher Top-K support 上的概率质量。 +- `reduced_topk_opd_teacher_selected_mass`：Teacher Top-K 概率质量。 +- `reduced_topk_opd_overlap_fraction`：Student Top-K 与 Teacher Top-K 的平均重叠比例。 + +Training-side Teacher 还会记录： + +- `time/train_teacher_compute` +- `time/train_teacher_onload` +- `time/train_teacher_offload` + +最终蒸馏指标会写入 `distillation/` 前缀下；不同 worker 的原始 step 平均值也会写入 `train_metrics/worker_/`。 + +## 当前限制 + +- sampled-token 只支持 `k1 + Policy Gradient`。 +- Rollout Top-K 只支持 LMDeploy。 +- 一个 `DistillationConfig` 不能混用 Rollout Teacher 和 Training-side Teacher。 +- Training-side MOPD 会让所有 Teacher 依次 forward 所有本地 pack。 +- Training-side Teacher 支持 CPU offload，但不支持 disk offload。 +- 当前 Top-K 是 Teacher-selected Top-K，不是 full-vocabulary distillation。 +- Reference Model 和 OPD Teacher 是不同角色；即使 checkpoint 相同，配置语义也不相同。 diff --git a/xtuner/v1/rl/loss/distillation_loss.py b/xtuner/v1/rl/loss/distillation_loss.py index e655c06ec7..de5d358dd6 100644 --- a/xtuner/v1/rl/loss/distillation_loss.py +++ b/xtuner/v1/rl/loss/distillation_loss.py @@ -26,7 +26,7 @@ logger = get_logger() TopKDistillationMode = Literal["forward", "reverse", "forward_kl_topk"] -DistillationLossMode = Literal["k1", "k3", "forward", "reverse", "forward_kl_topk"] +DistillationLossMode = Literal["k1", "forward", "reverse", "forward_kl_topk"] def compute_topk_distillation_kl( @@ -94,7 +94,7 @@ def validate_loss_options(self) -> Self: @property def uses_sampled_token_targets(self) -> bool: - return self.loss_mode in ("k1", "k3") + return self.loss_mode == "k1" @property def uses_topk_targets(self) -> bool: @@ -125,9 +125,12 @@ def build( is_weights=data.get("rollout_is_weights"), teacher_logprobs=data.get("teacher_logprobs"), target_token_ids=data.get("target_token_ids"), - ).to(DEVICE) + ) + # Rollout teacher targets can be large, so shard them before moving the + # local slice to the accelerator. if sp_mesh is not None and sp_mesh.size() > 1: loss_kwargs = loss_kwargs.sp_split(sp_mesh) + loss_kwargs = loss_kwargs.to(DEVICE) return self.loss_ctx_cls(self, loss_kwargs) @@ -229,11 +232,7 @@ def _compute_sampled_distillation_loss( student_logprobs: torch.Tensor, teacher_logprobs: torch.Tensor, ) -> torch.Tensor: - if self.loss_cfg.loss_mode == "k1": - loss = student_logprobs - teacher_logprobs - else: - log_ratio = (teacher_logprobs - student_logprobs).clamp(min=-20.0, max=20.0) - loss = torch.exp(log_ratio) - log_ratio - 1.0 + loss = student_logprobs - teacher_logprobs if self.loss_cfg.loss_max_clamp is not None: loss = loss.clamp( min=-self.loss_cfg.loss_max_clamp, diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index 079036da84..1f8be9e156 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -693,7 +693,8 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo rollout_logprobs = data.get("rollout_logprobs", None) rollout_logprobs = rollout_logprobs.to(DEVICE) if rollout_logprobs is not None else None if teacher_indices is not None: - teacher_indices = teacher_indices.to(DEVICE) + # Keep the full indices on the host until each SP rank selects + # its local slice. if self.sp_mesh.size() > 1: teacher_indices = sp_split( teacher_indices, @@ -701,6 +702,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo split_dim=1, padding_value=-1, ) + teacher_indices = teacher_indices.to(DEVICE) teacher_indices_list.append(teacher_indices) loss_ctx = loss_cfg.build( data={ From 5e845c4fe648decf0c4ff7cf79b62d4a40dd5584 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Mon, 31 Aug 2026 10:14:00 +0000 Subject: [PATCH 06/20] [Refactor] Centralize AgentLoop post-processing and improve OPD teacher integration --- docs/zh_cn/rl/advanced_tutorial/agent_loop.md | 17 +-- .../verl_agent/common/agent_loop_verl_tool.py | 4 - tests/rl/test_on_policy_distillation.py | 88 ++++++++++++- tests/rl/test_prepare_train_data.py | 27 ++++ tests/rl/test_producer.py | 3 +- tests/rl/test_rl_colocate_trainer.py | 2 +- tests/rl/test_rl_disaggregated_trainer.py | 49 +++++++ tests/rl/test_single_turn_agent_loop.py | 116 ++++++++++++++++- xtuner/v1/rl/agent_loop/agent_loop.py | 122 ++++++------------ xtuner/v1/rl/agent_loop/gsm8k_with_tool.py | 2 - .../agent_in_localhost_loop.py | 22 +++- .../agent_in_sandbox_loop.py | 49 +++---- .../rl/agent_loop/single_turn_agent_loop.py | 8 +- .../rl/agent_loop_manager/disagg_producer.py | 2 +- .../v1/rl/agent_loop_manager/produce_utils.py | 6 +- xtuner/v1/rl/agent_loop_manager/producer.py | 6 +- .../rl/distillation/train_teacher_manager.py | 55 ++++++-- xtuner/v1/rl/trainer/controller.py | 5 - xtuner/v1/rl/trainer/worker.py | 2 + xtuner/v1/train/rl_trainer.py | 107 +++++++++++---- 20 files changed, 495 insertions(+), 197 deletions(-) diff --git a/docs/zh_cn/rl/advanced_tutorial/agent_loop.md b/docs/zh_cn/rl/advanced_tutorial/agent_loop.md index bb937f899a..ac51618aed 100644 --- a/docs/zh_cn/rl/advanced_tutorial/agent_loop.md +++ b/docs/zh_cn/rl/advanced_tutorial/agent_loop.md @@ -58,7 +58,7 @@ async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> l ... ``` -`generate_sample()` 处理单条样本；默认 `generate_group()` 会并发调用多次 `generate_sample()`，并在调用前把 `self.sample_params` 写到组内每条样本上。若配置了 `enable_batch_judge=True`，默认 `generate_group()` 会在组内样本生成完成后调用一次 `self.run_judger(group_samples)`。需要其他组级逻辑时，例如组内共享环境、组内排序过滤，可以覆盖 `generate_group()`。 +`generate_sample()` 只负责生成单条样本；默认 `generate_group()` 会并发调用多次 `generate_sample()`，并负责 Teacher scoring、逐条或批量 Judger 以及组级过滤的执行顺序。调用前，`generate_group()` 会把 `self.sample_params` 写到组内每条样本上。需要其他组级逻辑时，例如组内共享环境、组内排序过滤，可以覆盖 `generate_group()`，但需要同时明确 Teacher、Judger 和过滤的先后关系。 ## 输入输出约定 @@ -139,18 +139,18 @@ agent_loop_config = SingleTurnAgentLoopConfig( ) ``` -开启后，`generate_sample()` 不会逐条调用 Judger；`generate_group()` 会在组内样本全部生成完成后调用一次 `self.run_judger(group_samples)`，内部会转到 `judger.batch_judge(group_samples)`。只有当前 Judger 明确实现 `batch_judge()` 时才应开启。 +开启后，`generate_group()` 会在组内样本全部完成 generation 和必要的 Teacher scoring 后调用一次 `self.run_judger(group_samples)`，内部会转到 `judger.batch_judge(group_samples)`。关闭时，`generate_group()` 会对每条完成的样本调用 `self.run_judger(sample)`。只有当前 Judger 明确实现 `batch_judge()` 时才应开启。 ## 自定义 AgentLoop 自定义 AgentLoop 通常需要做四件事： 1. 继承 `AgentLoop`，实现 `generate_sample()`。 -2. 在 `generate_sample()` 中维护 `tokens`、`sample_params`、`response_ids`、`response`、`logprobs`、`response_mask`、`status`。 -3. 需要打分时，在 response 可用后调用 `self.run_judger(...)`。 +2. 在 `generate_sample()` 中维护 `tokens`、`sample_params`、`response_ids`、`response`、`logprobs`、`response_mask`、`status`，不要在这里调用外部 Judger。 +3. 若覆盖 `generate_group()`，在其中显式编排 Teacher、Judger 和组级过滤；没有 validity check 时应让完成的样本尽早进入 Teacher，有 validity check 时应只把过滤通过的完整组发送给 Teacher。 4. 继承 `AgentLoopConfig`，实现 `build_local()`，这样才能接入 `TaskSpecConfig.agent_loop_config`，并复用 Ray actor 构建逻辑。 -`self.run_judger(...)` 会根据输入形态调用 `judger.judge()` 或 `judger.batch_judge()`，并统一处理 pause/cancel。若自定义 AgentLoop 支持 `enable_batch_judge=True`，`generate_sample()` 中的单条打分需要用 `not self.enable_batch_judge` 保护，避免默认 `generate_group()` 再次批量打分。若自定义 AgentLoop 需要覆盖 `pause()`，应在实现中调用 `await super().pause()`，否则正在执行的 Judger 任务无法复用基类的中断逻辑。 +`self.run_judger(...)` 会根据输入形态调用 `judger.judge()` 或 `judger.batch_judge()`，并统一处理 pause/cancel。默认 `generate_group()` 已经根据 `enable_batch_judge` 选择逐条或批量打分，自定义实现不应再把外部 Judger 塞回 `generate_sample()`。若自定义 AgentLoop 需要覆盖 `pause()`，应在实现中调用 `await super().pause()`，否则正在执行的 Judger 任务无法复用基类的中断逻辑。 ### 最小单轮实现 @@ -169,11 +169,6 @@ class CustomAgentLoop(AgentLoop): rollout_state.sample_params = rollout_state.sample_params or self.sample_params rollout_state = await self.rollout_ctl.generate.remote(rollout_state) - if rollout_state.status != Status.COMPLETED: - return rollout_state - - if self.judger is not None and not self.enable_batch_judge: - rollout_state = await self.run_judger(rollout_state) return rollout_state @@ -244,8 +239,6 @@ class ToolAgentLoop(AgentLoop): assert len(rollout_state.response_ids) == len(rollout_state.logprobs) assert len(rollout_state.response_ids) == len(rollout_state.response_mask) - if rollout_state.status == Status.COMPLETED and self.judger is not None and not self.enable_batch_judge: - rollout_state = await self.run_judger(rollout_state) return rollout_state ``` diff --git a/recipe/verl_agent/common/agent_loop_verl_tool.py b/recipe/verl_agent/common/agent_loop_verl_tool.py index 6956d83f8c..8e9f0e0a65 100644 --- a/recipe/verl_agent/common/agent_loop_verl_tool.py +++ b/recipe/verl_agent/common/agent_loop_verl_tool.py @@ -147,8 +147,4 @@ async def generate_sample(self, rollout_state: RolloutState) -> RolloutState: # raw_prompt is updated in tool_agent_loop: apply_chat_template of tools rollout_state.extra_fields["raw_prompt"] = self.tokenizer.decode(rollout_state.prompt_ids) - # judge rollout_state - if self.judger is not None: - rollout_state = await self.judger.judge(rollout_state) - return rollout_state diff --git a/tests/rl/test_on_policy_distillation.py b/tests/rl/test_on_policy_distillation.py index b9d3a98dfd..fa9d4cb3f9 100644 --- a/tests/rl/test_on_policy_distillation.py +++ b/tests/rl/test_on_policy_distillation.py @@ -1,7 +1,8 @@ import tempfile import unittest from pathlib import Path -from unittest.mock import AsyncMock, patch +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch import httpx import torch @@ -9,11 +10,18 @@ from recipe.on_policy_distillation.build_teacher_server_commands import ( build_teacher_launch_server_commands, ) -from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status from xtuner.v1.data_proto.sequence_context import SequenceContext -from xtuner.v1.rl.distillation import RolloutTeacherClient, RolloutTeacherConfig +from xtuner.v1.rl.distillation import ( + DistillationConfig, + RolloutTeacherClient, + RolloutTeacherConfig, + TrainTeacherManager, + TrainTeacherTimings, +) from xtuner.v1.rl.loss import DistillationLossConfig from xtuner.v1.rl.trainer.controller import TrainingController +from xtuner.v1.train.rl_trainer import BaseRLTrainer, BaseRLTrainerConfig class TestDistillationRecipeConfig(unittest.TestCase): @@ -65,6 +73,80 @@ def test_teacher_launcher_reads_distillation_config(self) -> None: self.assertEqual(records[0][:4], ["teacher[0]", "0", "1", "http://127.0.0.1:13141"]) self.assertIn("lmdeploy", records[0][7:]) + def test_sampled_token_sampling_is_validated_by_trainer_config(self) -> None: + loss_config = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + distillation_config = DistillationConfig( + loss_config=loss_config, + teachers=[RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"])], + data_source_teacher_map={"math": "teacher"}, + ) + config = BaseRLTrainerConfig.model_construct( + train_worker_cfg=SimpleNamespace(loss_cfg=loss_config, model_cfg=MagicMock()), + agent_loop_manager_cfg=SimpleNamespace( + tasks=SimpleNamespace( + task_name="math", + agent_loop_config=SimpleNamespace(sample_params=SampleParams(temperature=0.7)), + ) + ), + distillation_config=distillation_config, + total_train_steps=1, + ) + + with self.assertRaisesRegex(ValueError, "Invalid sample_params for task 'math'.*temperature"): + config._validate_sync_intervals() + + def test_train_teacher_config_is_not_forwarded_to_agent_loop_manager(self) -> None: + trainer = BaseRLTrainer.__new__(BaseRLTrainer) + trainer._rollout_teacher_config = None + trainer._enable_evaluate = False + trainer._resolve_total_train_steps = MagicMock() + trainer.rollout_controller = MagicMock() + trainer.logger = MagicMock() + built_manager = MagicMock() + manager_config = MagicMock() + manager_config.build.return_value = built_manager + config = SimpleNamespace( + tokenizer_path="student", + agent_loop_manager_cfg=manager_config, + sync_weights_interval=1, + ) + + with patch("xtuner.v1.train.rl_trainer.AutoTokenizer.from_pretrained", return_value=MagicMock()): + trainer._build_agent_loop_components(config, replay_buffer=MagicMock()) + + self.assertIs(trainer.agent_loop_manager, built_manager) + self.assertIsNone(manager_config.build.call_args.kwargs["distillation_config"]) + + +class TestTrainTeacherTimings(unittest.TestCase): + def test_lifecycle_timings_are_recorded_for_every_teacher(self) -> None: + manager = TrainTeacherManager.__new__(TrainTeacherManager) + first_teacher = MagicMock() + second_teacher = MagicMock() + timings = TrainTeacherTimings() + + with ( + patch( + "xtuner.v1.rl.distillation.train_teacher_manager.time.perf_counter", + side_effect=[1.0, 2.0, 3.0, 5.0, 6.0, 9.0, 10.0, 14.0, 15.0, 20.0, 21.0, 27.0], + ), + patch.object(manager, "_synchronize_device"), + patch.object(manager, "_offload_to_cpu"), + ): + with manager._teacher_on_device(first_teacher, timings, "teacher_a"): + pass + with manager._teacher_on_device(second_teacher, timings, "teacher_b"): + pass + + self.assertEqual( + timings.to_dict(), + { + "teacher_a": {"compute": 2.0, "onload": 1.0, "offload": 3.0}, + "teacher_b": {"compute": 5.0, "onload": 4.0, "offload": 6.0}, + }, + ) + self.assertEqual((timings.compute, timings.onload, timings.offload), (7.0, 5.0, 9.0)) + class TestRolloutTeacherClient(unittest.IsolatedAsyncioTestCase): @staticmethod diff --git a/tests/rl/test_prepare_train_data.py b/tests/rl/test_prepare_train_data.py index e4fe72ea2f..b6cedb7965 100644 --- a/tests/rl/test_prepare_train_data.py +++ b/tests/rl/test_prepare_train_data.py @@ -190,6 +190,33 @@ def test_vlm_path_uses_train_prompt_ids_and_preserves_multimodal_fields(self): self.assertEqual(seq_ctx.image_grid_thw.dtype, torch.long) self.assertEqual(seq_ctx.image_grid_thw.tolist(), [[1, 2, 3]]) + def test_mixed_agentic_and_vlm_reasoning_use_3d_position_ids(self): + trainer = self._build_trainer([0.5]) + reasoning_state = self._state( + uid=1, + group_id=1, + prompt_ids=[10, 11, 12], + response_ids=[20, 21, 22], + position_ids=np.arange(3, dtype=np.int64).reshape(1, 1, -1).repeat(3, axis=0), + ) + agentic_state = self._state( + uid=2, + group_id=2, + input_ids=[30, 31, 40, 41, 42], + labels=[-100, -100, 40, 41, 42], + logprobs=[0.0, -0.1, -0.2, -0.3, -0.4], + ) + + data_batches, _ = self._prepare(trainer, [[reasoning_state], [agentic_state]]) + + self.assertEqual(tuple(data_batches[0]["seq_ctx"].position_ids.shape), (3, 1, 5)) + agentic_position_ids = data_batches[1]["seq_ctx"].position_ids + self.assertEqual(tuple(agentic_position_ids.shape), (3, 1, 4)) + torch.testing.assert_close( + agentic_position_ids, + torch.arange(4, dtype=torch.long).reshape(1, 1, -1).expand(3, -1, -1), + ) + def test_agentic_topk_targets_include_token_ids_and_logprobs(self): loss_config = DistillationLossConfig( policy_loss_cfg={ diff --git a/tests/rl/test_producer.py b/tests/rl/test_producer.py index 44f8e67c67..8ea9883104 100644 --- a/tests/rl/test_producer.py +++ b/tests/rl/test_producer.py @@ -120,7 +120,6 @@ async def mock_gen(rs, **kwargs): mock_agent_loop.generate_group = mock_gen mock_agent_loop.is_valid_sample_fn = is_valid_sample_fn mock_agent_loop.teacher_clients = {} - mock_agent_loop.collect_rollout_group = AgentLoop.collect_rollout_group.__get__(mock_agent_loop) return mock_agent_loop def _build_context( @@ -560,7 +559,7 @@ async def mock_gen(rs, **kwargs): r.status = Status.COMPLETED return rs - mock_agent_loop.collect_rollout_group = mock_gen + mock_agent_loop.generate_group = mock_gen sampler_cfg = SamplerConfig.model_construct(dataloader_cfg=self.mock_dataloader_cfg) produce_strategy_cfg = AsyncProduceStrategyConfig(over_sample_threshold=1) diff --git a/tests/rl/test_rl_colocate_trainer.py b/tests/rl/test_rl_colocate_trainer.py index 1b6b3481ed..b5ef8e96e8 100644 --- a/tests/rl/test_rl_colocate_trainer.py +++ b/tests/rl/test_rl_colocate_trainer.py @@ -100,7 +100,7 @@ async def generate_group(rollout_states, **kwargs): state.response_model_steps = [model_step] return rollout_states - agent_loop.collect_rollout_group = generate_group + agent_loop.generate_group = generate_group return agent_loop diff --git a/tests/rl/test_rl_disaggregated_trainer.py b/tests/rl/test_rl_disaggregated_trainer.py index 933d9e76fd..add03a5086 100644 --- a/tests/rl/test_rl_disaggregated_trainer.py +++ b/tests/rl/test_rl_disaggregated_trainer.py @@ -472,6 +472,55 @@ def test_log_step_records_disaggregated_effective_e2e_window(self): self.assertAlmostEqual(scalars["throughput/rollout_sgs"], 0.5) self.assertAlmostEqual(scalars["throughput/rollout_tgs"], 15.0) + def test_log_step_records_every_teacher_timing_and_distillation_metric(self): + trainer = self._make_trainer(_FakeManager([])) + trainer._log_step = RLDisaggregatedTrainer._log_step.__get__(trainer, RLDisaggregatedTrainer) + train_info = self._minimal_train_info(training_samples=2, training_tokens=8) + train_info["workers_log_item"] = [ + { + "rollout_is_metrics": {}, + "mismatch_metrics": {}, + "rollout_entropy": 0.0, + "train_entropy": 0.0, + "teacher_timings": { + "teacher_a": {"compute": 1.0, "onload": 2.0, "offload": 3.0}, + "teacher_b": {"compute": 3.0, "onload": 4.0, "offload": 5.0}, + }, + "train_metrics": [ + {"actor_distillation_kl": 0.5, "sampled_opd_loss": 0.25, "policy_loss": 1.0} + ], + "sft_train_metrics": {}, + }, + { + "teacher_timings": { + "teacher_a": {"compute": 4.0, "onload": 1.0, "offload": 2.0}, + "teacher_b": {"compute": 2.0, "onload": 3.0, "offload": 4.0}, + }, + "train_entropy": 0.0, + "train_metrics": [], + "sft_train_metrics": {}, + }, + ] + + trainer._log_step( + train_step=1, + step_timer_dict={}, + produce_result=ProduceBatchResult(rollout_states=[]), + train_info=train_info, + eval_info={}, + ) + + scalars = trainer._exp_tracker.add_scalars.call_args.kwargs["tag_scalar_dict"] + self.assertEqual(scalars["time/train_teacher/teacher_a/compute"], 4.0) + self.assertEqual(scalars["time/train_teacher/teacher_a/onload"], 2.0) + self.assertEqual(scalars["time/train_teacher/teacher_a/offload"], 3.0) + self.assertEqual(scalars["time/train_teacher/teacher_b/compute"], 3.0) + self.assertEqual(scalars["time/train_teacher/total/compute"], 6.0) + self.assertEqual(scalars["time/train_teacher_compute"], 6.0) + self.assertEqual(scalars["distillation/actor_distillation_kl"], 0.5) + self.assertEqual(scalars["distillation/sampled_opd_loss"], 0.25) + self.assertNotIn("distillation/policy_loss", scalars) + def test_update_weights_pauses_generation_without_onloading_rollout(self): manager = _FakeManager([]) trainer = self._make_trainer(manager) diff --git a/tests/rl/test_single_turn_agent_loop.py b/tests/rl/test_single_turn_agent_loop.py index 45d7db9da2..21bf28a261 100644 --- a/tests/rl/test_single_turn_agent_loop.py +++ b/tests/rl/test_single_turn_agent_loop.py @@ -5,6 +5,9 @@ - batch judge 只在整组样本全部 COMPLETED 时触发。 - batch judge 返回结果必须保持输入顺序。 - rollout group validity check 在 batch judge 之后执行。 +- 没有 validity check 时，每条样本先算 Teacher、再进入 Judger。 +- 有 validity check 时，只有通过过滤的 rollout group 才进入 Teacher scoring。 +- 被过滤的 rollout group 不进入 Teacher scoring。 - 组内存在 ABORTED / FAILED 等非 COMPLETED 样本时跳过 judge。 - pause 发生在 slow judger 期间时，整组样本被标记为 ABORTED。 """ @@ -38,16 +41,29 @@ async def remote(self, rollout_state: RolloutState): class _BatchJudger: - def __init__(self): + def __init__(self, events: list[str] | None = None): self.calls: list[list[RolloutState]] = [] + self.events = events async def batch_judge(self, rollout_states): + if self.events is not None: + self.events.append("judge") self.calls.append(rollout_states) for state in rollout_states: state.reward = {"score": float(state.rollout_id)} return rollout_states +class _PerSampleJudger: + def __init__(self, events: list[str]): + self.events = events + + async def judge(self, rollout_state): + self.events.append("judge") + rollout_state.reward = {"score": float(rollout_state.rollout_id)} + return rollout_state + + class _SlowJudger: async def batch_judge(self, rollout_states): await asyncio.sleep(60) @@ -68,7 +84,14 @@ def _state(self, uid: int) -> RolloutState: extra_fields={}, ) - def _build_loop(self, statuses_by_uid: dict[int, Status], judger=None, is_valid_sample_fn=None): + def _build_loop( + self, + statuses_by_uid: dict[int, Status], + judger=None, + is_valid_sample_fn=None, + *, + enable_batch_judge: bool = True, + ): loop = SingleTurnAgentLoop.__new__(SingleTurnAgentLoop) rollout_ctl = MagicMock() rollout_ctl.generate = _RemoteGenerate(statuses_by_uid) @@ -76,12 +99,28 @@ def _build_loop(self, statuses_by_uid: dict[int, Status], judger=None, is_valid_ loop.rollout_ctl = rollout_ctl loop.sample_params = SampleParams(max_tokens=8, temperature=0.7) loop.judger = judger - loop.enable_batch_judge = True + loop.enable_batch_judge = enable_batch_judge loop.is_valid_sample_fn = is_valid_sample_fn + loop.teacher_clients = {} + loop.data_source_teacher_map = {} loop._judger_pause_event = asyncio.Event() loop.logger = MagicMock() return loop + @staticmethod + def _bind_teacher(loop, events: list[str]): + teacher = MagicMock() + + async def compute_logprobs(state): + events.append("teacher") + state.teacher_logprobs = [-0.5] * len(state.response_ids or []) + return state + + teacher.compute_logprobs = AsyncMock(side_effect=compute_logprobs) + loop.teacher_clients = {"teacher": teacher} + loop.data_source_teacher_map = {"math": "teacher"} + return teacher + def test_config_binds_validity_check_to_local_agent_loop_once(self): is_valid_sample_fn = MagicMock() local_loop = MagicMock() @@ -142,6 +181,77 @@ def is_valid_sample_fn(samples): self.assertTrue(all(state.status == Status.FILTERED for state in result)) loop.logger.info.assert_called_once() + async def test_generate_skips_teacher_scoring_for_filtered_group(self): + teacher = MagicMock() + teacher.compute_logprobs = AsyncMock() + loop = self._build_loop( + {1: Status.COMPLETED}, + is_valid_sample_fn=lambda _samples: False, + ) + loop.teacher_clients = {"teacher": teacher} + loop.data_source_teacher_map = {"math": "teacher"} + state = self._state(1) + state.extra_fields["origin_data_source"] = "math" + + result = await loop.generate_group([state]) + + self.assertEqual(result[0].status, Status.FILTERED) + teacher.compute_logprobs.assert_not_awaited() + + async def test_teacher_scoring_precedes_batch_judge_without_filter(self): + events: list[str] = [] + loop = self._build_loop( + {1: Status.COMPLETED, 2: Status.COMPLETED}, + judger=_BatchJudger(events), + ) + teacher = self._bind_teacher(loop, events) + states = [self._state(1), self._state(2)] + for state in states: + state.extra_fields["origin_data_source"] = "math" + + result = await loop.generate_group(states) + + self.assertEqual(events.count("teacher"), 2) + self.assertEqual(events[-1], "judge") + self.assertTrue(all(state.teacher_logprobs == [-0.5] for state in result)) + self.assertEqual(teacher.compute_logprobs.await_count, 2) + + async def test_teacher_scoring_precedes_per_sample_judge_without_filter(self): + events: list[str] = [] + loop = self._build_loop( + {1: Status.COMPLETED}, + judger=_PerSampleJudger(events), + enable_batch_judge=False, + ) + self._bind_teacher(loop, events) + state = self._state(1) + state.extra_fields["origin_data_source"] = "math" + + result = await loop.generate_group([state]) + + self.assertEqual(events, ["teacher", "judge"]) + self.assertEqual(result[0].reward, {"score": 1.0}) + + async def test_valid_group_is_scored_after_batch_judge_and_filter(self): + events: list[str] = [] + + def is_valid_sample_fn(_samples): + events.append("filter") + return True + + loop = self._build_loop( + {1: Status.COMPLETED}, + judger=_BatchJudger(events), + is_valid_sample_fn=is_valid_sample_fn, + ) + self._bind_teacher(loop, events) + state = self._state(1) + state.extra_fields["origin_data_source"] = "math" + + await loop.generate_group([state]) + + self.assertEqual(events, ["judge", "filter", "teacher"]) + async def test_batch_judge_runs_once_when_all_samples_completed_and_preserves_order(self): # 整组样本全部 COMPLETED 时才触发 batch judger；返回顺序必须和输入顺序一致。 judger = _BatchJudger() diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index 1a0c2e34a6..8b06792251 100644 --- a/xtuner/v1/rl/agent_loop/agent_loop.py +++ b/xtuner/v1/rl/agent_loop/agent_loop.py @@ -15,7 +15,6 @@ DistillationConfig, RolloutTeacherClient, route_rollout_teacher_client, - validate_opd_sample_params, ) from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController @@ -44,17 +43,10 @@ def maybe_filter_invalid_sample( is_valid_sample_fn: IsValidSampleFn | None, logger, ) -> list[RolloutState]: - """Finalize rollout-group validity after all generation post-processing. + """Apply task-specific group validation after generation and judging. - Every custom ``AgentLoop.generate_group`` implementation must return through - this helper after generation, judging, flattening, and failure cleanup:: - - # Keep sample validation as the final group-generation step. - return maybe_filter_invalid_sample( - samples, - self.is_valid_sample_fn, - self.logger, - ) + Teacher scoring may run before this helper when no validity check is configured. When a validity check is + configured, call this helper before sending the completed group to the Teacher. """ if get_group_status(group) != Status.COMPLETED: return group @@ -257,7 +249,6 @@ def configure_distillation(self, distillation_config: DistillationConfig | None) if distillation_config is None: return - validate_opd_sample_params(self.sample_params) self.teacher_clients = { teacher.name: RolloutTeacherClient(teacher, distillation_config.loss_config) for teacher in distillation_config.rollout_teachers @@ -265,69 +256,53 @@ def configure_distillation(self, distillation_config: DistillationConfig | None) self.data_source_teacher_map = dict(distillation_config.data_source_teacher_map) @abstractmethod - async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: ... + async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: + """Generate one rollout sample without group-level post-processing.""" + + ... + + async def maybe_compute_teacher_logprob(self, state: RolloutState) -> RolloutState: + if state.status != Status.COMPLETED or not self.teacher_clients: + return state + teacher = route_rollout_teacher_client( + state, + data_source_teacher_map=self.data_source_teacher_map, + teacher_clients=self.teacher_clients, + ) + return await teacher.compute_logprobs(state) + + async def maybe_compute_teacher_logprobs(self, group: list[RolloutState]) -> list[RolloutState]: + return list(await asyncio.gather(*(create_task(self.maybe_compute_teacher_logprob(state)) for state in group))) async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: """Generate one rollout group. Warning: - Subclasses overriding this method MUST call - ``maybe_filter_invalid_sample`` as the final step before returning. - Otherwise ``TaskSpecConfig.is_valid_sample_fn`` will be silently ignored. + Subclasses overriding this method must preserve the Teacher/Judger/filter + ordering described here. Without a validity check, Teacher scoring should + start as soon as each sample finishes generation. With a validity check, + only a completed group that passes filtering should be sent to the Teacher. """ - pending_tasks = [] - for state in rollout_state: - state.sample_params = self.sample_params - task = create_task(self.generate_sample(state, **kwargs)) - pending_tasks.append(task) - generated_samples = asyncio.gather(*pending_tasks) - group_samples = await generated_samples - if self.judger is not None and self.enable_batch_judge: - if all(sample.status == Status.COMPLETED for sample in group_samples): - group_samples = await self.run_judger(group_samples) - # Keep sample validation as the final group-generation step. - return maybe_filter_invalid_sample(group_samples, self.is_valid_sample_fn, self.logger) - - async def collect_rollout_group( - self, - rollout_state: list[RolloutState], - *, - is_valid_sample_func: Callable[[list[RolloutState]], bool] | None = None, - **kwargs, - ) -> list[RolloutState]: - if is_valid_sample_func is None and self.teacher_clients: - teacher = route_rollout_teacher_client( - rollout_state[0], - data_source_teacher_map=self.data_source_teacher_map, - teacher_clients=self.teacher_clients, - ) + filter_before_teacher = self.is_valid_sample_fn is not None - async def generate_and_score(state: RolloutState) -> RolloutState: - state.sample_params = self.sample_params - state = await self.generate_sample(state, **kwargs) - if state.status == Status.COMPLETED: - state = await teacher.compute_logprobs(state) - return state - - group = list(await asyncio.gather(*(create_task(generate_and_score(state)) for state in rollout_state))) - if self.judger is not None and self.enable_batch_judge and get_group_status(group) == Status.COMPLETED: - group = await self.run_judger(group) - return group - - group = await self.generate_group(rollout_state, **kwargs) - if get_group_status(group) != Status.COMPLETED: - return group - if is_valid_sample_func is not None and not is_valid_sample_func(group): - for state in group: - state.status = Status.FILTERED - return group - if self.teacher_clients: - teacher = route_rollout_teacher_client( - group[0], - data_source_teacher_map=self.data_source_teacher_map, - teacher_clients=self.teacher_clients, - ) - group = list(await asyncio.gather(*(create_task(teacher.compute_logprobs(state)) for state in group))) + async def generate_one(state: RolloutState) -> RolloutState: + state.sample_params = self.sample_params + state = await self.generate_sample(state, **kwargs) + # Fast path: overlap Teacher scoring with the remaining generations. + if not filter_before_teacher: + state = await self.maybe_compute_teacher_logprob(state) + if state.status == Status.COMPLETED and self.judger is not None and not self.enable_batch_judge: + state = await self.run_judger(state) + return state + + group = list(await asyncio.gather(*(create_task(generate_one(state)) for state in rollout_state))) + if self.judger is not None and self.enable_batch_judge and get_group_status(group) == Status.COMPLETED: + group = await self.run_judger(group) + + group = maybe_filter_invalid_sample(group, self.is_valid_sample_fn, self.logger) + # Filter path: avoid Teacher work for a group that will be discarded. + if filter_before_teacher and get_group_status(group) == Status.COMPLETED: + group = await self.maybe_compute_teacher_logprobs(group) return group @overload @@ -416,13 +391,6 @@ async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> l finally: await self._release_worker(worker) - async def collect_rollout_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: - worker = await self._pick_worker() - try: - return await worker.collect_rollout_group.remote(rollout_state, **kwargs) - finally: - await self._release_worker(worker) - def get_worker_status(self) -> dict[str, int]: return {str(worker): load for worker, load in self._worker_loads.items()} @@ -470,10 +438,6 @@ async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> Rollou async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: return await self.agent_loop.generate_group(rollout_state, **kwargs) - @ray_method(concurrency_group=AGENT_LOOP_CONCURRENCY_GROUP_GENERATE) - async def collect_rollout_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: - return await self.agent_loop.collect_rollout_group(rollout_state, **kwargs) - @ray_method async def get_rollout_ctl(self): return self.agent_loop.rollout_ctl diff --git a/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py b/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py index 13044d4f35..aeeaf0df3d 100644 --- a/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py +++ b/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py @@ -159,6 +159,4 @@ async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> Rollou assert len(rollout_state.response_ids) == len(rollout_state.response_mask) == len(rollout_state.logprobs), ( f"{len(rollout_state.response_ids)} vs {len(rollout_state.response_mask)} vs {len(rollout_state.logprobs)}" ) - if self.judger is not None and not self.enable_batch_judge: - rollout_state = await self.run_judger(rollout_state) return rollout_state diff --git a/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py b/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py index 9da3701bef..e069be9367 100644 --- a/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py +++ b/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py @@ -9,7 +9,7 @@ from lagent.utils import create_object, ctx_session_id -from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status +from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status from xtuner.v1.rl.agent_loop.sandbox_agent_loop.schemas import ( AgentRolloutItem, RolloutStatus, @@ -134,11 +134,18 @@ def __init__( self.mode = mode async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: + filter_before_teacher = self.is_valid_sample_fn is not None + async def generate_one(state: RolloutState) -> RolloutState: if self._sample_semaphore is None: - return await self.generate_sample(state, **kwargs) - async with self._sample_semaphore: - return await self.generate_sample(state, **kwargs) + state = await self.generate_sample(state, **kwargs) + else: + async with self._sample_semaphore: + state = await self.generate_sample(state, **kwargs) + # Fast path: release the sample slot, then score immediately. + if not filter_before_teacher: + state = await self.maybe_compute_teacher_logprob(state) + return state tasks: list[asyncio.Task[RolloutState]] = [] for state in rollout_state: @@ -148,8 +155,11 @@ async def generate_one(state: RolloutState) -> RolloutState: samples = await asyncio.gather(*tasks) samples = _drop_failed_train_samples(samples, self.mode) - # Keep sample validation as the final group-generation step. - return maybe_filter_invalid_sample(samples, self.is_valid_sample_fn, self.logger) + samples = maybe_filter_invalid_sample(samples, self.is_valid_sample_fn, self.logger) + # Filter path: score only the completed group that survived filtering. + if filter_before_teacher and get_group_status(samples) == Status.COMPLETED: + samples = await self.maybe_compute_teacher_logprobs(samples) + return samples async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: try: diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py index 15d12da47d..913f42d9e1 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py @@ -12,7 +12,6 @@ from lagent.utils import create_object from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status -from xtuner.v1.rl.distillation import route_rollout_teacher_client from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController from xtuner.v1.rl.utils import create_task @@ -230,40 +229,19 @@ def __init__( self._sample_semaphore = asyncio.Semaphore(max_concurrent_samples) if max_concurrent_samples else None self.mode = mode - async def collect_rollout_group( - self, - rollout_state: list[RolloutState], - *, - is_valid_sample_func: Callable[[list[RolloutState]], bool] | None = None, - **kwargs, - ) -> list[RolloutState]: - """Generate agent traces and score every completed segment.""" - group = await self.generate_group(rollout_state, **kwargs) - if get_group_status(group) != Status.COMPLETED: - return group - if is_valid_sample_func is not None and not is_valid_sample_func(group): - for state in group: - state.status = Status.FILTERED - return group - if not self.teacher_clients: - return group - - async def score_trace(state: RolloutState) -> RolloutState: - teacher = route_rollout_teacher_client( - state, - data_source_teacher_map=self.data_source_teacher_map, - teacher_clients=self.teacher_clients, - ) - return await teacher.compute_logprobs(state) - - return list(await asyncio.gather(*(create_task(score_trace(state)) for state in group))) - async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: + filter_before_teacher = self.is_valid_sample_fn is not None + async def generate_one(state: RolloutState) -> list[RolloutState]: if self._sample_semaphore is None: - return await self.generate_sample(state) - async with self._sample_semaphore: - return await self.generate_sample(state) + samples = await self.generate_sample(state) + else: + async with self._sample_semaphore: + samples = await self.generate_sample(state) + # Fast path: score every completed trace segment as soon as its session returns. + if not filter_before_teacher: + samples = await self.maybe_compute_teacher_logprobs(samples) + return samples pending_tasks = [] for state in rollout_state: @@ -274,8 +252,11 @@ async def generate_one(state: RolloutState) -> list[RolloutState]: sample_groups = await generated_samples samples = [sample for sample_group in sample_groups for sample in sample_group] samples = _drop_failed_train_samples(samples, self.mode) - # Keep sample validation as the final group-generation step. - return maybe_filter_invalid_sample(samples, self.is_valid_sample_fn, self.logger) + samples = maybe_filter_invalid_sample(samples, self.is_valid_sample_fn, self.logger) + # Filter path: score only the completed group that survived filtering. + if filter_before_teacher and get_group_status(samples) == Status.COMPLETED: + samples = await self.maybe_compute_teacher_logprobs(samples) + return samples # NOTE: A single sandbox session may yield multiple trainable segments, so this returns a list # rather than the base class's single RolloutState. The base contract is never exercised for diff --git a/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py b/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py index 2368d1c3b5..cc185538ba 100644 --- a/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py +++ b/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py @@ -1,4 +1,4 @@ -from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status +from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController from xtuner.v1.rl.trace.rollout_api import trace_rollout_endpoint, trace_rollout_remote @@ -80,10 +80,4 @@ async def generate_sample( self.rollout_ctl.generate, # type: ignore[attr-defined] rollout_state, ) - # 非 COMPLETED 状态（如被截断、放弃等）直接早退，不触发打分 - if rollout_state.status != Status.COMPLETED: - return rollout_state - if self.judger is not None and not self.enable_batch_judge: - # 如果开启了批量打分，则在 generate_group 里统一打分，不在这里逐条打分 - rollout_state = await self.run_judger(rollout_state) return rollout_state diff --git a/xtuner/v1/rl/agent_loop_manager/disagg_producer.py b/xtuner/v1/rl/agent_loop_manager/disagg_producer.py index 2fdde63ab1..76503ebe17 100644 --- a/xtuner/v1/rl/agent_loop_manager/disagg_producer.py +++ b/xtuner/v1/rl/agent_loop_manager/disagg_producer.py @@ -424,7 +424,7 @@ async def produce_batch(self, ctx: DisaggProduceContext) -> ProduceBatchStatus: async def spawn_one() -> asyncio.Task: rollout_state = await ctx.sample_group(from_expired_pool=sample_expired) return create_task( - ctx.collect_rollout_group( + ctx.generate_group( rollout_state, enable_partial_rollout=self.enable_partial_rollout, ) diff --git a/xtuner/v1/rl/agent_loop_manager/produce_utils.py b/xtuner/v1/rl/agent_loop_manager/produce_utils.py index fef9a0f746..4832e88f2d 100644 --- a/xtuner/v1/rl/agent_loop_manager/produce_utils.py +++ b/xtuner/v1/rl/agent_loop_manager/produce_utils.py @@ -127,7 +127,7 @@ async def sample_group(self, *, from_expired_pool: bool) -> list[RolloutState]: group_status = [Status.EXPIRED, Status.ABORTED] if from_expired_pool else [Status.ABORTED] return await self.sampler.sample(task_name=self.task_name, group_status=group_status) - async def collect_rollout_group( + async def generate_group( self, rollout_state: list[RolloutState], *, @@ -141,13 +141,13 @@ async def collect_rollout_group( start = time.perf_counter() if isinstance(self.agent_loop, ray.actor.ActorHandle): - result = await self.agent_loop.collect_rollout_group.remote( + result = await self.agent_loop.generate_group.remote( rollout_state, is_valid_sample_func=self.is_valid_sample_fn, enable_partial_rollout=enable_partial_rollout, ) else: - result = await self.agent_loop.collect_rollout_group( + result = await self.agent_loop.generate_group( rollout_state, is_valid_sample_func=self.is_valid_sample_fn, enable_partial_rollout=enable_partial_rollout, diff --git a/xtuner/v1/rl/agent_loop_manager/producer.py b/xtuner/v1/rl/agent_loop_manager/producer.py index 75c89bf1d1..2f7f660e70 100644 --- a/xtuner/v1/rl/agent_loop_manager/producer.py +++ b/xtuner/v1/rl/agent_loop_manager/producer.py @@ -283,7 +283,7 @@ async def produce_batch(self, ctx: ProduceContext) -> None: for _ in range(ctx.task_batch_size): rollout_state = await ctx.sampler.sample(task_name=ctx.task_name) - task = create_task(ctx.collect_rollout_group(rollout_state)) + task = create_task(ctx.generate_group(rollout_state)) pending_tasks.add(task) logger.info(f"[SyncProduceStrategy] Started {len(pending_tasks)} initial tasks.") @@ -316,7 +316,7 @@ async def produce_batch(self, ctx: ProduceContext) -> None: completed_sample_count, ctx.task_batch_size ): rollout_state = await ctx.sampler.sample(task_name=ctx.task_name) - task = create_task(ctx.collect_rollout_group(rollout_state)) + task = create_task(ctx.generate_group(rollout_state)) pending_tasks.add(task) progress_displayer.close() @@ -405,7 +405,7 @@ async def produce_batch(self, ctx: ProduceContext) -> None: async def spawn_one() -> asyncio.Task: rollout_state = await ctx.sample_group(from_expired_pool=sample_expired) return create_task( - ctx.collect_rollout_group( + ctx.generate_group( rollout_state, enable_partial_rollout=self.enable_partial_rollout, ) diff --git a/xtuner/v1/rl/distillation/train_teacher_manager.py b/xtuner/v1/rl/distillation/train_teacher_manager.py index b441d5fbb6..f769fa1c19 100644 --- a/xtuner/v1/rl/distillation/train_teacher_manager.py +++ b/xtuner/v1/rl/distillation/train_teacher_manager.py @@ -22,14 +22,36 @@ @dataclass -class TrainTeacherTimings: - """Wall-clock seconds spent in the frozen Teacher lifecycle.""" +class TeacherTiming: + """Wall-clock seconds spent in one frozen Teacher lifecycle.""" compute: float = 0.0 onload: float = 0.0 offload: float = 0.0 +@dataclass +class TrainTeacherTimings(TeacherTiming): + """Wall-clock seconds spent in each frozen Teacher lifecycle.""" + + by_teacher: dict[str, TeacherTiming] = field(default_factory=dict) + + def to_dict(self) -> dict[str, dict[str, float]]: + """Return per-Teacher lifecycle timings. + + Returns: + dict[str, dict[str, float]]: Teacher names mapped to compute, onload, and offload seconds. + """ + return { + teacher_name: { + "compute": timing.compute, + "onload": timing.onload, + "offload": timing.offload, + } + for teacher_name, timing in self.by_teacher.items() + } + + @dataclass class TrainTeacherOutputs: teacher_logprobs: list[torch.Tensor] @@ -76,6 +98,7 @@ def __init__(self, distillation_config: DistillationConfig, *, chunk_size: int | teacher_config.name: teacher_index for teacher_index, teacher_config in enumerate(distillation_config.train_teachers) } + self._teacher_names = [teacher_config.name for teacher_config in distillation_config.train_teachers] def compute_logprobs( self, @@ -133,20 +156,34 @@ def _teacher_on_device( self, teacher: FrozenModel, timings: TrainTeacherTimings, + teacher_name: str | None = None, ) -> Iterator[None]: + teacher_timing = ( + timings.by_teacher.setdefault(teacher_name, TeacherTiming()) if teacher_name is not None else None + ) + onload_begin = time.perf_counter() teacher.to_device(DEVICE) - timings.onload += time.perf_counter() - onload_begin + onload_elapsed = time.perf_counter() - onload_begin + timings.onload += onload_elapsed + if teacher_timing is not None: + teacher_timing.onload += onload_elapsed compute_begin = time.perf_counter() try: yield self._synchronize_device() finally: - timings.compute += time.perf_counter() - compute_begin + compute_elapsed = time.perf_counter() - compute_begin + timings.compute += compute_elapsed + if teacher_timing is not None: + teacher_timing.compute += compute_elapsed offload_begin = time.perf_counter() self._offload_to_cpu(teacher) - timings.offload += time.perf_counter() - offload_begin + offload_elapsed = time.perf_counter() - offload_begin + timings.offload += offload_elapsed + if teacher_timing is not None: + teacher_timing.offload += offload_elapsed @staticmethod def _teacher_seq_ctx(seq_ctx: SequenceContext, *, is_composed: bool) -> SequenceContext: @@ -193,8 +230,8 @@ def _compute_topk_targets( # All ranks iterate Teachers and local packs in the same order so every # FSDP rank enters the same collective sequence, even when a rank has no # tokens routed to a particular Teacher. - for teacher_index, teacher in enumerate(self._teachers): - with self._teacher_on_device(teacher, timings): + for teacher_index, (teacher_name, teacher) in enumerate(zip(self._teacher_names, self._teachers)): + with self._teacher_on_device(teacher, timings, teacher_name): # Every rank forwards every local pack before selecting routed # tokens so all ranks enter the same FSDP collective sequence. for batch_index, (seq_ctx, teacher_indices) in enumerate(zip(seq_ctx_list, teacher_indices_list)): @@ -225,8 +262,8 @@ def _compute_sampled_logprobs( ] # Keep the Teacher-major schedule identical across ranks for FSDP. - for teacher_index, teacher in enumerate(self._teachers): - with self._teacher_on_device(teacher, timings): + for teacher_index, (teacher_name, teacher) in enumerate(zip(self._teacher_names, self._teachers)): + with self._teacher_on_device(teacher, timings, teacher_name): for batch_index, (seq_ctx, shifted_labels, teacher_indices) in enumerate( zip(seq_ctx_list, shifted_labels_list, teacher_indices_list) ): diff --git a/xtuner/v1/rl/trainer/controller.py b/xtuner/v1/rl/trainer/controller.py index 3925631931..b84ff6c4dc 100644 --- a/xtuner/v1/rl/trainer/controller.py +++ b/xtuner/v1/rl/trainer/controller.py @@ -104,11 +104,6 @@ def _packing(self, data_batches, pack_max_length, language_cfg): packed_data_batches = [] is_qwen3_vl = any(data["seq_ctx"].position_ids.ndim == 3 for data in data_batches) - if is_qwen3_vl: - for data in data_batches: - position_ids = data["seq_ctx"].position_ids - if position_ids.ndim == 2: - data["seq_ctx"].position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) has_rollout_routed_experts = False if data_batches[0]["seq_ctx"].rollout_routed_experts is not None: diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index 1f8be9e156..c2d2b79512 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -215,6 +215,7 @@ class WorkerTrainLogItem(TypedDict, total=False): class WorkerLogItem(TypedDict): train_entropy: float + teacher_timings: NotRequired[dict[str, dict[str, float]]] teacher_compute_time: NotRequired[float] teacher_onload_time: NotRequired[float] teacher_offload_time: NotRequired[float] @@ -758,6 +759,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo teacher_indices_list, loss_ctx_list, ) + worker_log_item["teacher_timings"] = teacher_timings.to_dict() worker_log_item["teacher_compute_time"] = teacher_timings.compute worker_log_item["teacher_onload_time"] = teacher_timings.onload worker_log_item["teacher_offload_time"] = teacher_timings.offload diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 11724313b7..989e67d21a 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -33,7 +33,7 @@ ProduceBatchStatus, ) from xtuner.v1.rl.agent_loop_manager.produce_utils import default_should_continue_fn -from xtuner.v1.rl.distillation import DistillationConfig +from xtuner.v1.rl.distillation import DistillationConfig, validate_opd_sample_params from xtuner.v1.rl.evaluator import EvaluatorConfig from xtuner.v1.rl.health_manager import RLHealthManager, _NoOpRLHealthManager from xtuner.v1.rl.loss import DistillationLossConfig @@ -70,13 +70,6 @@ DEVICE = get_device() DEVICE_MODULE = get_torch_device_module() -_DISTILLATION_METRIC_PREFIXES = ( - "reduced_distillation_", - "reduced_topk_opd_", - "opd_", - "topk_opd_", -) - def _to_cpu_tensor(value: np.ndarray | None, *, dtype: torch.dtype | None = None) -> torch.Tensor | None: if value is None: @@ -524,6 +517,19 @@ def _validate_sync_intervals(self): if self.train_worker_cfg.loss_cfg != self.distillation_config.loss_config: raise ValueError("train_worker_cfg.loss_cfg must be distillation_config.loss_config") self.distillation_config.validate_student_model(self.train_worker_cfg.model_cfg) + if self.distillation_config.loss_config.uses_sampled_token_targets: + tasks = self.agent_loop_manager_cfg.tasks + tasks = tasks if isinstance(tasks, list) else [tasks] + for task in tasks: + sample_params = task.agent_loop_config.sample_params + if sample_params is None: + raise ValueError( + f"Task {task.task_name!r} must configure sample_params for sampled-token distillation" + ) + try: + validate_opd_sample_params(sample_params) + except ValueError as exc: + raise ValueError(f"Invalid sample_params for task {task.task_name!r}: {exc}") from exc _validate_sync_intervals( sync_weights_interval=self.sync_weights_interval, checkpoint_interval=self.checkpoint_interval, @@ -877,18 +883,13 @@ def _init_runtime_flags(self, cfg: BaseRLTrainerConfig) -> None: def _build_agent_loop_components(self, cfg: BaseRLTrainerConfig, replay_buffer) -> None: self.tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer_path, trust_remote_code=True) - # Sampled-token objectives require identity Student sampling even when - # Teacher logprobs are produced later by the training workers. - agent_loop_distillation_config = self._rollout_teacher_config - if self._distillation_loss_cfg is not None and self._distillation_loss_cfg.uses_sampled_token_targets: - agent_loop_distillation_config = self._distillation_config agent_loop_manager = cfg.agent_loop_manager_cfg.build( rollout_controller=self.rollout_controller, tokenizer=self.tokenizer, replay_buffer=replay_buffer, logger=self.logger, sync_weights_interval=cfg.sync_weights_interval, - distillation_config=agent_loop_distillation_config, + distillation_config=self._rollout_teacher_config, ) self.agent_loop_manager = cast(AgentLoopManager | DisaggAgentLoopManager, agent_loop_manager) @@ -1217,6 +1218,14 @@ def _prepare_train_data( raw_rewards_sum: float = 0.0, raw_rewards_count: int = 0, ): + has_agentic_data = any(state.input_ids is not None for group in data_groups for state in group) + has_3d_reasoning_data = any( + state.input_ids is None and state.position_ids is not None and state.position_ids.ndim == 3 + for group in data_groups + for state in group + ) + use_3d_agentic_position_ids = has_agentic_data and has_3d_reasoning_data + # Per-session rewards for distribution metrics. Agentic sessions may split into several # trainable segments that share one reward; counting that reward once per session keeps # rewards/* from being weighted by segment count. @@ -1366,6 +1375,13 @@ def _prepare_train_data( ) position_ids = group[i].position_ids + if use_3d_agentic_position_ids: + seq_len = input_ids_t.size(-1) + text_position_ids = np.arange(seq_len, dtype=np.int64).reshape(1, 1, -1) + # Mixed agentic/VLM batches use three-axis MRoPE position IDs. Repeat the + # normal text positions on every axis so agentic samples have the same rank + # as the VLM samples when their sequence contexts are packed together. + position_ids = np.broadcast_to(text_position_ids, (3, 1, seq_len)).copy() multimodal_train_info = group[i].mm_info multi_info_cast = cast(dict | None, multimodal_train_info) seq_ctx = get_train_seq_ctx(input_ids_t, position_ids, multi_info_cast) @@ -1632,7 +1648,8 @@ def _log_step( ) trajectory_str = f"\nTrain step {train_step} data statistics:\n" trajectory_str += "\n".join([f"- {k:<25}: {v:.4f}" for k, v in response_data_info.items()]) - rank0_log_item = train_info["workers_log_item"][0] + worker_log_items = train_info["workers_log_item"] + rank0_log_item = worker_log_items[0] rank0_rollout_is_metrics = rank0_log_item.get("rollout_is_metrics", {}) rank0_mismatch_metrics = rank0_log_item.get("mismatch_metrics", {}) rank0_rollout_entropy = rank0_log_item.get("rollout_entropy", 0.0) @@ -1640,13 +1657,57 @@ def _log_step( all_scalars.update({f"{k}": v for k, v in rank0_mismatch_metrics.items()}) all_scalars.update({"entropy/rollout": rank0_rollout_entropy}) all_scalars.update({"entropy/train": rank0_log_item["train_entropy"]}) - if "teacher_compute_time" in rank0_log_item: - all_scalars["time/train_teacher_compute"] = rank0_log_item["teacher_compute_time"] - if "teacher_onload_time" in rank0_log_item: - all_scalars["time/train_teacher_onload"] = rank0_log_item["teacher_onload_time"] - if "teacher_offload_time" in rank0_log_item: - all_scalars["time/train_teacher_offload"] = rank0_log_item["teacher_offload_time"] - for worker_idx, log_item in enumerate(train_info["workers_log_item"]): + + teacher_timing_phases = ("compute", "onload", "offload") + teacher_timings_by_worker = [log_item.get("teacher_timings", {}) for log_item in worker_log_items] + teacher_names = sorted( + {teacher_name for teacher_timings in teacher_timings_by_worker for teacher_name in teacher_timings} + ) + for teacher_name in teacher_names: + for phase in teacher_timing_phases: + phase_values = [ + float(teacher_timings[teacher_name][phase]) + for teacher_timings in teacher_timings_by_worker + if teacher_name in teacher_timings + ] + all_scalars[f"time/train_teacher/{teacher_name}/{phase}"] = max(phase_values) + + # Training waits for the slowest worker, so report per-worker Teacher totals and then + # take the maximum as the step-level critical-path time. + for phase in teacher_timing_phases: + worker_totals = [ + sum(float(teacher_timing[phase]) for teacher_timing in teacher_timings.values()) + for teacher_timings in teacher_timings_by_worker + ] + if teacher_names: + total = max(worker_totals) + else: + legacy_values: list[float] = [] + if phase == "compute": + legacy_values = [ + float(legacy_time) + for log_item in worker_log_items + if (legacy_time := log_item.get("teacher_compute_time")) is not None + ] + elif phase == "onload": + legacy_values = [ + float(legacy_time) + for log_item in worker_log_items + if (legacy_time := log_item.get("teacher_onload_time")) is not None + ] + else: + legacy_values = [ + float(legacy_time) + for log_item in worker_log_items + if (legacy_time := log_item.get("teacher_offload_time")) is not None + ] + if not legacy_values: + continue + total = max(legacy_values) + all_scalars[f"time/train_teacher/total/{phase}"] = total + all_scalars[f"time/train_teacher_{phase}"] = total + + for worker_idx, log_item in enumerate(worker_log_items): if not self._display_all_workers_log and worker_idx > 0: break mini_batch_metrics: dict[str, List[float]] = {} @@ -1657,7 +1718,7 @@ def _log_step( for key, value in mini_batch_metrics.items(): avg_value = sum(value) / len(value) all_scalars.update({f"train_metrics/worker_{worker_idx}/step_avg_{key}": avg_value}) - if worker_idx == 0 and key.startswith(_DISTILLATION_METRIC_PREFIXES): + if worker_idx == 0 and ("opd" in key or "distillation" in key): all_scalars[f"distillation/{key}"] = avg_value if key in ("opd_reverse_kl", "opd_abs_logprob_loss"): all_scalars[key] = avg_value From 306a66afc9bef070d269527ca62a5e8a7c9e010f Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Fri, 4 Sep 2026 06:02:48 +0000 Subject: [PATCH 07/20] fix merge conflicts --- examples/v1/scripts/run_rl.sh | 6 +++++- examples/v1/scripts/run_rl_submit.sh | 9 ++++++++- tests/rl/test_multi_task_agent_loop_manager.py | 1 + tests/rl/test_rl_colocate_trainer.py | 1 - .../sandbox_agent_loop/agent_in_sandbox_loop.py | 1 - xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py | 2 -- .../rl/agent_loop_manager/disagg_agent_loop_manager.py | 3 +-- xtuner/v1/rl/agent_loop_manager/produce_utils.py | 3 --- xtuner/v1/rl/distillation/train_teacher_manager.py | 3 +-- 9 files changed, 16 insertions(+), 13 deletions(-) diff --git a/examples/v1/scripts/run_rl.sh b/examples/v1/scripts/run_rl.sh index bd313e17bd..01912f931d 100644 --- a/examples/v1/scripts/run_rl.sh +++ b/examples/v1/scripts/run_rl.sh @@ -102,7 +102,11 @@ fi # 根据 NODE_COUNT 分配 num_cpus, 防止内存OOM node_count=${NODE_COUNT:-1} expected_accelerator_count=$((node_count * accelerator_per_node)) -export XTUNER_RL_NUM_WORKERS=${XTUNER_RL_NUM_WORKERS:-$expected_accelerator_count} +# OPD launchers may assign a different number of Student GPUs to each node +# after placing Teacher replicas. Use their already-computed cluster total +# for readiness; the local GPU count is only used to start this Ray node. +expected_accelerator_count=${XTUNER_RL_NUM_WORKERS:-$expected_accelerator_count} +export XTUNER_RL_NUM_WORKERS=${expected_accelerator_count} WORK_DIR=$(realpath "$WORK_DIR") if [ "$RAY_RANK" -eq 0 ]; then diff --git a/examples/v1/scripts/run_rl_submit.sh b/examples/v1/scripts/run_rl_submit.sh index 636bfa116f..405af4ecc9 100644 --- a/examples/v1/scripts/run_rl_submit.sh +++ b/examples/v1/scripts/run_rl_submit.sh @@ -76,6 +76,12 @@ fi # 2. Launch Ray cluster # 根据 NODE_COUNT 分配 num_cpus, 防止内存OOM node_count=${NODE_COUNT:-1} +expected_accelerator_count=$((node_count * ACCELERATOR_PER_NODE)) +# OPD launchers may assign a different number of Student GPUs to each node +# after placing Teacher replicas. Use their already-computed cluster total +# for readiness and the submitted Trainer resource count. +expected_accelerator_count=${XTUNER_RL_NUM_WORKERS:-$expected_accelerator_count} +export XTUNER_RL_NUM_WORKERS=${expected_accelerator_count} if [ "$RAY_RANK" -eq 0 ]; then rm -rf /tmp/ray_log @@ -105,7 +111,6 @@ fi while true; do result=$(ray status | grep ${ACCELERATOR} | cut -d ' ' -f2 | cut -d '/' -f2) - expected_accelerator_count=$((node_count * ${ACCELERATOR_PER_NODE})) if [ "$result" = "$expected_accelerator_count.0" ]; then break else @@ -148,6 +153,7 @@ if [ "$RAY_RANK" -eq 0 ]; then "PYTORCH_CUDA_ALLOC_CONF": "${PYTORCH_CUDA_ALLOC_CONF_VALUE}", "LMDEPLOY_LOG_FILE": "${LMDEPLOY_LOG_FILE}", "XTUNER_RL_MEM_DIR": "${XTUNER_RL_MEM_DIR}", + "XTUNER_RL_NUM_WORKERS": "${XTUNER_RL_NUM_WORKERS}", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1", "SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION": "False" @@ -160,6 +166,7 @@ EOF_JSON --runtime-env-json="$RUNTIME_ENV_JSON" \ -- python xtuner/v1/train/cli/rl.py \ --config $CONFIG_PATH \ + --num-workers "$XTUNER_RL_NUM_WORKERS" \ 2>&1 | tee -a "$LOG_FILE" echo "训练任务提交完成。日志文件: $LOG_FILE" diff --git a/tests/rl/test_multi_task_agent_loop_manager.py b/tests/rl/test_multi_task_agent_loop_manager.py index 4062a8e853..5447f0ad90 100644 --- a/tests/rl/test_multi_task_agent_loop_manager.py +++ b/tests/rl/test_multi_task_agent_loop_manager.py @@ -255,6 +255,7 @@ def test_manager_config_injects_validity_check_when_building_agent_loop(self): judger=None, logger=None, is_valid_sample_fn=is_valid_sample_fn, + distillation_config=None, ) self.assertIs(manager.task_runners[0].agent_loop, built_agent_loop) diff --git a/tests/rl/test_rl_colocate_trainer.py b/tests/rl/test_rl_colocate_trainer.py index b5ef8e96e8..d16446750b 100644 --- a/tests/rl/test_rl_colocate_trainer.py +++ b/tests/rl/test_rl_colocate_trainer.py @@ -214,7 +214,6 @@ def test_fit_accepts_async_strategy_manager_on_colocate_path(self): trainer.fit() trainer.rollout_controller.offload.remote.assert_called_once_with() - trainer.train_controller.onload.assert_called_once_with(target="all") trainer.train_controller.fit.assert_called_once() trainer._release_all_trace_sessions.assert_called_once_with() trainer._release_trace_sessions.assert_not_called() diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py index 913f42d9e1..4415d488c9 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py @@ -6,7 +6,6 @@ import json import traceback import uuid -from collections.abc import Callable from typing import Any, Literal from lagent.utils import create_object diff --git a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py index e988b793f6..733825e031 100644 --- a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py @@ -19,7 +19,6 @@ _MANAGER_STATE_PATH, _STATUS_POLL_INTERVAL_S, _TASK_CHECKPOINT_DIR, - IsValidSampleFn, ProduceBatchResult, _TaskRunner, _TaskSamplerView, @@ -162,7 +161,6 @@ def build( agent_loop=agent_loop, produce_strategy=produce_strategy, sampler=sampler, - is_valid_sample_fn=task_cfg.filter_func, weight=task_cfg.weight, order=order, ) diff --git a/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py index dd245a9695..b6f62b6f80 100644 --- a/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py @@ -10,6 +10,7 @@ from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast from xtuner.v1.data_proto.rl_data import Status from xtuner.v1.rl.agent_loop import AgentLoopConfig, IsValidSampleFn +from xtuner.v1.rl.distillation import DistillationConfig from xtuner.v1.rl.judger import ComposedJudgerConfig, JudgerConfig, build_judger from xtuner.v1.rl.replay_buffer import ReplayBuffer from xtuner.v1.rl.rollout import RolloutController @@ -27,7 +28,6 @@ _MANAGER_STATE_PATH, _STATUS_POLL_INTERVAL_S, _TASK_CHECKPOINT_DIR, - IsValidSampleFn, ProduceBatchResult, ProduceBatchStatus, _TaskRunner, @@ -101,7 +101,6 @@ def build( agent_loop=agent_loop, produce_strategy=produce_strategy, sampler=sampler, - is_valid_sample_fn=task_cfg.filter_func, weight=task_cfg.weight, order=order, ) diff --git a/xtuner/v1/rl/agent_loop_manager/produce_utils.py b/xtuner/v1/rl/agent_loop_manager/produce_utils.py index 4832e88f2d..1a2f8e365f 100644 --- a/xtuner/v1/rl/agent_loop_manager/produce_utils.py +++ b/xtuner/v1/rl/agent_loop_manager/produce_utils.py @@ -143,13 +143,11 @@ async def generate_group( if isinstance(self.agent_loop, ray.actor.ActorHandle): result = await self.agent_loop.generate_group.remote( rollout_state, - is_valid_sample_func=self.is_valid_sample_fn, enable_partial_rollout=enable_partial_rollout, ) else: result = await self.agent_loop.generate_group( rollout_state, - is_valid_sample_func=self.is_valid_sample_fn, enable_partial_rollout=enable_partial_rollout, ) elapsed = time.perf_counter() - start @@ -256,7 +254,6 @@ class _TaskRunner: agent_loop: AgentLoopSpec produce_strategy: Any sampler: Sampler - is_valid_sample_fn: IsValidSampleFn | None = None weight: float = 1.0 order: int = 0 diff --git a/xtuner/v1/rl/distillation/train_teacher_manager.py b/xtuner/v1/rl/distillation/train_teacher_manager.py index f769fa1c19..d341714525 100644 --- a/xtuner/v1/rl/distillation/train_teacher_manager.py +++ b/xtuner/v1/rl/distillation/train_teacher_manager.py @@ -7,7 +7,7 @@ import torch -from xtuner.v1.data_proto.sequence_context import DSATopKCacheState, SequenceContext +from xtuner.v1.data_proto.sequence_context import SequenceContext from xtuner.v1.loss import LogProbConfig, LogProbContext, TopKLogProbConfig from xtuner.v1.model.compose.base import BaseComposeConfig from xtuner.v1.rl.loss import DistillationLossConfig @@ -190,7 +190,6 @@ def _teacher_seq_ctx(seq_ctx: SequenceContext, *, is_composed: bool) -> Sequence overrides = { "rollout_routed_experts": None, "offload_rollout_routed_experts": False, - "dsa_topk_cache": DSATopKCacheState(), } if not is_composed: # A VLM tokenizer supplies 3D M-RoPE position ids for every sample, From eb2b862ba76aa0deef722768510c5706b718748d Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Mon, 7 Sep 2026 06:59:38 +0000 Subject: [PATCH 08/20] fix claude review comments --- xtuner/v1/loss/rl_loss.py | 6 ++++ xtuner/v1/rl/agent_loop/agent_loop.py | 32 +++++++++++++------ .../agent_in_localhost_loop.py | 17 +++------- .../agent_in_sandbox_loop.py | 17 +++------- .../rl/distillation/train_teacher_manager.py | 8 +++-- 5 files changed, 45 insertions(+), 35 deletions(-) diff --git a/xtuner/v1/loss/rl_loss.py b/xtuner/v1/loss/rl_loss.py index fd22b2c629..1356683ede 100644 --- a/xtuner/v1/loss/rl_loss.py +++ b/xtuner/v1/loss/rl_loss.py @@ -116,6 +116,12 @@ def build(self, data: dict, sp_mesh: DeviceMesh | None = None) -> "TopKLogProbCo class TopKLogProbContext(BaseLossContext): """Select model Top-K IDs without storing full log-softmax.""" + # Top-k mode reuses the standard loss-context return protocol: + # - the loss slot carries selected logprobs + # - the logits slot carries token ids + # This keeps the model/worker plumbing unchanged, but callers must + # interpret ModelOutputs.loss/logits according to this contract. + loss_cfg: TopKLogProbConfig loss_kwargs: BaseLossKwargs diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index 8b06792251..9bdea96eca 100644 --- a/xtuner/v1/rl/agent_loop/agent_loop.py +++ b/xtuner/v1/rl/agent_loop/agent_loop.py @@ -274,6 +274,26 @@ async def maybe_compute_teacher_logprob(self, state: RolloutState) -> RolloutSta async def maybe_compute_teacher_logprobs(self, group: list[RolloutState]) -> list[RolloutState]: return list(await asyncio.gather(*(create_task(self.maybe_compute_teacher_logprob(state)) for state in group))) + def _should_filter_before_teacher(self) -> bool: + return self.is_valid_sample_fn is not None + + async def _maybe_score_state(self, state: RolloutState) -> RolloutState: + if not self._should_filter_before_teacher(): + return await self.maybe_compute_teacher_logprob(state) + return state + + async def _maybe_filter_group(self, group: list[RolloutState]) -> list[RolloutState]: + if self.is_valid_sample_fn is None: + return group + return maybe_filter_invalid_sample(group, self.is_valid_sample_fn, self.logger) + + async def _maybe_score_filtered_group(self, group: list[RolloutState]) -> list[RolloutState]: + if not self._should_filter_before_teacher(): + return group + if get_group_status(group) != Status.COMPLETED: + return group + return await self.maybe_compute_teacher_logprobs(group) + async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: """Generate one rollout group. @@ -283,14 +303,11 @@ async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> l start as soon as each sample finishes generation. With a validity check, only a completed group that passes filtering should be sent to the Teacher. """ - filter_before_teacher = self.is_valid_sample_fn is not None async def generate_one(state: RolloutState) -> RolloutState: state.sample_params = self.sample_params state = await self.generate_sample(state, **kwargs) - # Fast path: overlap Teacher scoring with the remaining generations. - if not filter_before_teacher: - state = await self.maybe_compute_teacher_logprob(state) + state = await self._maybe_score_state(state) if state.status == Status.COMPLETED and self.judger is not None and not self.enable_batch_judge: state = await self.run_judger(state) return state @@ -299,11 +316,8 @@ async def generate_one(state: RolloutState) -> RolloutState: if self.judger is not None and self.enable_batch_judge and get_group_status(group) == Status.COMPLETED: group = await self.run_judger(group) - group = maybe_filter_invalid_sample(group, self.is_valid_sample_fn, self.logger) - # Filter path: avoid Teacher work for a group that will be discarded. - if filter_before_teacher and get_group_status(group) == Status.COMPLETED: - group = await self.maybe_compute_teacher_logprobs(group) - return group + group = await self._maybe_filter_group(group) + return await self._maybe_score_filtered_group(group) @overload async def run_judger(self, rollout_state: RolloutState) -> RolloutState: ... diff --git a/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py b/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py index e069be9367..fc88fd3f85 100644 --- a/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py +++ b/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py @@ -9,7 +9,7 @@ from lagent.utils import create_object, ctx_session_id -from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status +from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status from xtuner.v1.rl.agent_loop.sandbox_agent_loop.schemas import ( AgentRolloutItem, RolloutStatus, @@ -20,7 +20,7 @@ from xtuner.v1.rl.rollout.trace_store import get_store from xtuner.v1.rl.utils import create_task -from ..agent_loop import AgentLoop, AgentLoopConfig, maybe_filter_invalid_sample +from ..agent_loop import AgentLoop, AgentLoopConfig def _import_from_path(path: str) -> Any: @@ -134,17 +134,13 @@ def __init__( self.mode = mode async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: - filter_before_teacher = self.is_valid_sample_fn is not None - async def generate_one(state: RolloutState) -> RolloutState: if self._sample_semaphore is None: state = await self.generate_sample(state, **kwargs) else: async with self._sample_semaphore: state = await self.generate_sample(state, **kwargs) - # Fast path: release the sample slot, then score immediately. - if not filter_before_teacher: - state = await self.maybe_compute_teacher_logprob(state) + state = await self._maybe_score_state(state) return state tasks: list[asyncio.Task[RolloutState]] = [] @@ -155,11 +151,8 @@ async def generate_one(state: RolloutState) -> RolloutState: samples = await asyncio.gather(*tasks) samples = _drop_failed_train_samples(samples, self.mode) - samples = maybe_filter_invalid_sample(samples, self.is_valid_sample_fn, self.logger) - # Filter path: score only the completed group that survived filtering. - if filter_before_teacher and get_group_status(samples) == Status.COMPLETED: - samples = await self.maybe_compute_teacher_logprobs(samples) - return samples + samples = await self._maybe_filter_group(samples) + return await self._maybe_score_filtered_group(samples) async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: try: diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py index 4415d488c9..1febb52610 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py @@ -10,14 +10,14 @@ from lagent.utils import create_object -from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status +from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController from xtuner.v1.rl.utils import create_task from ...rollout.chat_template import canonicalize_messages_for_chat_template from ...rollout.trace_store import get_store -from ..agent_loop import AgentLoop, AgentLoopConfig, maybe_filter_invalid_sample +from ..agent_loop import AgentLoop, AgentLoopConfig from .schemas import AgentRolloutItem, RolloutStatus @@ -229,17 +229,13 @@ def __init__( self.mode = mode async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: - filter_before_teacher = self.is_valid_sample_fn is not None - async def generate_one(state: RolloutState) -> list[RolloutState]: if self._sample_semaphore is None: samples = await self.generate_sample(state) else: async with self._sample_semaphore: samples = await self.generate_sample(state) - # Fast path: score every completed trace segment as soon as its session returns. - if not filter_before_teacher: - samples = await self.maybe_compute_teacher_logprobs(samples) + samples = [await self._maybe_score_state(sample) for sample in samples] return samples pending_tasks = [] @@ -251,11 +247,8 @@ async def generate_one(state: RolloutState) -> list[RolloutState]: sample_groups = await generated_samples samples = [sample for sample_group in sample_groups for sample in sample_group] samples = _drop_failed_train_samples(samples, self.mode) - samples = maybe_filter_invalid_sample(samples, self.is_valid_sample_fn, self.logger) - # Filter path: score only the completed group that survived filtering. - if filter_before_teacher and get_group_status(samples) == Status.COMPLETED: - samples = await self.maybe_compute_teacher_logprobs(samples) - return samples + samples = await self._maybe_filter_group(samples) + return await self._maybe_score_filtered_group(samples) # NOTE: A single sandbox session may yield multiple trainable segments, so this returns a list # rather than the base class's single RolloutState. The base contract is never exercised for diff --git a/xtuner/v1/rl/distillation/train_teacher_manager.py b/xtuner/v1/rl/distillation/train_teacher_manager.py index d341714525..b273241982 100644 --- a/xtuner/v1/rl/distillation/train_teacher_manager.py +++ b/xtuner/v1/rl/distillation/train_teacher_manager.py @@ -245,8 +245,12 @@ def _compute_topk_targets( loss_ctx={"lm": loss_ctx}, ) selected = teacher_indices == teacher_index - target_ids[batch_index][selected] = cast(torch.Tensor, output.logits)[selected] - target_logprobs[batch_index][selected] = cast(torch.Tensor, output.loss)[selected] + # Top-k teacher outputs reuse ModelOutputs.logits for token ids + # and ModelOutputs.loss for selected logprobs. + topk_token_ids = cast(torch.Tensor, output.logits) + topk_logprobs = cast(torch.Tensor, output.loss) + target_ids[batch_index][selected] = topk_token_ids[selected] + target_logprobs[batch_index][selected] = topk_logprobs[selected] return target_ids, target_logprobs def _compute_sampled_logprobs( From 2bab84f1e74e75010f78d4464a7fdd824eb3ea99 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Thu, 10 Sep 2026 07:34:00 +0000 Subject: [PATCH 09/20] [Refactor] Consolidate rollout Teacher scoring and metrics --- xtuner/v1/rl/agent_loop/agent_loop.py | 67 +++-------- .../agent_in_localhost_loop.py | 8 +- .../agent_in_sandbox_loop.py | 8 +- xtuner/v1/rl/distillation/__init__.py | 6 +- ...r_client.py => rollout_teacher_manager.py} | 106 ++++++++++++++++-- .../rl/distillation/train_teacher_manager.py | 7 ++ xtuner/v1/rl/loss/__init__.py | 1 + xtuner/v1/rl/loss/metrics.py | 18 +++ xtuner/v1/rl/trainer/worker.py | 8 +- 9 files changed, 147 insertions(+), 82 deletions(-) rename xtuner/v1/rl/distillation/{rollout_teacher_client.py => rollout_teacher_manager.py} (84%) create mode 100644 xtuner/v1/rl/loss/metrics.py diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index 9bdea96eca..f23db62274 100644 --- a/xtuner/v1/rl/agent_loop/agent_loop.py +++ b/xtuner/v1/rl/agent_loop/agent_loop.py @@ -13,8 +13,7 @@ from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status from xtuner.v1.rl.distillation import ( DistillationConfig, - RolloutTeacherClient, - route_rollout_teacher_client, + RolloutTeacherScorer, ) from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController @@ -242,57 +241,18 @@ def __init__( else: self.logger = logger self._judger_pause_event = asyncio.Event() - self.teacher_clients: dict[str, RolloutTeacherClient] = {} - self.data_source_teacher_map: dict[str, str] = {} + self._teacher_scorer = RolloutTeacherScorer.disabled() def configure_distillation(self, distillation_config: DistillationConfig | None) -> None: if distillation_config is None: return - - self.teacher_clients = { - teacher.name: RolloutTeacherClient(teacher, distillation_config.loss_config) - for teacher in distillation_config.rollout_teachers - } - self.data_source_teacher_map = dict(distillation_config.data_source_teacher_map) - - @abstractmethod - async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: - """Generate one rollout sample without group-level post-processing.""" - - ... - - async def maybe_compute_teacher_logprob(self, state: RolloutState) -> RolloutState: - if state.status != Status.COMPLETED or not self.teacher_clients: - return state - teacher = route_rollout_teacher_client( - state, - data_source_teacher_map=self.data_source_teacher_map, - teacher_clients=self.teacher_clients, + self._teacher_scorer = RolloutTeacherScorer.from_distillation_config( + distillation_config, + defers_to_filter=self.is_valid_sample_fn is not None, ) - return await teacher.compute_logprobs(state) - - async def maybe_compute_teacher_logprobs(self, group: list[RolloutState]) -> list[RolloutState]: - return list(await asyncio.gather(*(create_task(self.maybe_compute_teacher_logprob(state)) for state in group))) - - def _should_filter_before_teacher(self) -> bool: - return self.is_valid_sample_fn is not None - async def _maybe_score_state(self, state: RolloutState) -> RolloutState: - if not self._should_filter_before_teacher(): - return await self.maybe_compute_teacher_logprob(state) - return state - - async def _maybe_filter_group(self, group: list[RolloutState]) -> list[RolloutState]: - if self.is_valid_sample_fn is None: - return group - return maybe_filter_invalid_sample(group, self.is_valid_sample_fn, self.logger) - - async def _maybe_score_filtered_group(self, group: list[RolloutState]) -> list[RolloutState]: - if not self._should_filter_before_teacher(): - return group - if get_group_status(group) != Status.COMPLETED: - return group - return await self.maybe_compute_teacher_logprobs(group) + @abstractmethod + async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: ... async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: """Generate one rollout group. @@ -307,17 +267,16 @@ async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> l async def generate_one(state: RolloutState) -> RolloutState: state.sample_params = self.sample_params state = await self.generate_sample(state, **kwargs) - state = await self._maybe_score_state(state) - if state.status == Status.COMPLETED and self.judger is not None and not self.enable_batch_judge: - state = await self.run_judger(state) + state = await self._teacher_scorer.on_sample_ready(state) return state group = list(await asyncio.gather(*(create_task(generate_one(state)) for state in rollout_state))) - if self.judger is not None and self.enable_batch_judge and get_group_status(group) == Status.COMPLETED: - group = await self.run_judger(group) + if self.judger is not None and self.enable_batch_judge: + if all(sample.status == Status.COMPLETED for sample in group): + group = await self.run_judger(group) - group = await self._maybe_filter_group(group) - return await self._maybe_score_filtered_group(group) + group = maybe_filter_invalid_sample(group, self.is_valid_sample_fn, self.logger) + return await self._teacher_scorer.on_group_ready(group) @overload async def run_judger(self, rollout_state: RolloutState) -> RolloutState: ... diff --git a/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py b/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py index fc88fd3f85..bdba686585 100644 --- a/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py +++ b/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py @@ -20,7 +20,7 @@ from xtuner.v1.rl.rollout.trace_store import get_store from xtuner.v1.rl.utils import create_task -from ..agent_loop import AgentLoop, AgentLoopConfig +from ..agent_loop import AgentLoop, AgentLoopConfig, maybe_filter_invalid_sample def _import_from_path(path: str) -> Any: @@ -140,7 +140,7 @@ async def generate_one(state: RolloutState) -> RolloutState: else: async with self._sample_semaphore: state = await self.generate_sample(state, **kwargs) - state = await self._maybe_score_state(state) + state = await self._teacher_scorer.on_sample_ready(state) return state tasks: list[asyncio.Task[RolloutState]] = [] @@ -151,8 +151,8 @@ async def generate_one(state: RolloutState) -> RolloutState: samples = await asyncio.gather(*tasks) samples = _drop_failed_train_samples(samples, self.mode) - samples = await self._maybe_filter_group(samples) - return await self._maybe_score_filtered_group(samples) + samples = maybe_filter_invalid_sample(samples, self.is_valid_sample_fn, self.logger) + return await self._teacher_scorer.on_group_ready(samples) async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: try: diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py index 1febb52610..14dec5bf74 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py @@ -17,7 +17,7 @@ from ...rollout.chat_template import canonicalize_messages_for_chat_template from ...rollout.trace_store import get_store -from ..agent_loop import AgentLoop, AgentLoopConfig +from ..agent_loop import AgentLoop, AgentLoopConfig, maybe_filter_invalid_sample from .schemas import AgentRolloutItem, RolloutStatus @@ -235,7 +235,7 @@ async def generate_one(state: RolloutState) -> list[RolloutState]: else: async with self._sample_semaphore: samples = await self.generate_sample(state) - samples = [await self._maybe_score_state(sample) for sample in samples] + samples = [await self._teacher_scorer.on_sample_ready(sample) for sample in samples] return samples pending_tasks = [] @@ -247,8 +247,8 @@ async def generate_one(state: RolloutState) -> list[RolloutState]: sample_groups = await generated_samples samples = [sample for sample_group in sample_groups for sample in sample_group] samples = _drop_failed_train_samples(samples, self.mode) - samples = await self._maybe_filter_group(samples) - return await self._maybe_score_filtered_group(samples) + samples = maybe_filter_invalid_sample(samples, self.is_valid_sample_fn, self.logger) + return await self._teacher_scorer.on_group_ready(samples) # NOTE: A single sandbox session may yield multiple trainable segments, so this returns a list # rather than the base class's single RolloutState. The base contract is never exercised for diff --git a/xtuner/v1/rl/distillation/__init__.py b/xtuner/v1/rl/distillation/__init__.py index 85c604ac17..4e92b7bf80 100644 --- a/xtuner/v1/rl/distillation/__init__.py +++ b/xtuner/v1/rl/distillation/__init__.py @@ -5,10 +5,10 @@ TeacherConfig, TrainTeacherConfig, ) -from .rollout_teacher_client import ( +from .rollout_teacher_manager import ( RolloutTeacherClient, RolloutTeacherReplicaRouter, - route_rollout_teacher_client, + RolloutTeacherScorer, validate_opd_sample_params, ) from .train_teacher_manager import TrainTeacherManager, TrainTeacherOutputs, TrainTeacherTimings @@ -22,7 +22,7 @@ "TrainTeacherConfig", "RolloutTeacherClient", "RolloutTeacherReplicaRouter", - "route_rollout_teacher_client", + "RolloutTeacherScorer", "TrainTeacherManager", "TrainTeacherOutputs", "TrainTeacherTimings", diff --git a/xtuner/v1/rl/distillation/rollout_teacher_client.py b/xtuner/v1/rl/distillation/rollout_teacher_manager.py similarity index 84% rename from xtuner/v1/rl/distillation/rollout_teacher_client.py rename to xtuner/v1/rl/distillation/rollout_teacher_manager.py index bc3bff4a48..0bcade5a66 100644 --- a/xtuner/v1/rl/distillation/rollout_teacher_client.py +++ b/xtuner/v1/rl/distillation/rollout_teacher_manager.py @@ -9,11 +9,12 @@ import httpx -from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status +from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status from xtuner.v1.rl.loss import DistillationLossConfig +from xtuner.v1.rl.utils import create_task from xtuner.v1.utils import get_logger -from .config import RolloutTeacherConfig +from .config import DistillationConfig, RolloutTeacherConfig logger = get_logger() @@ -81,6 +82,10 @@ def __init__(self, config: RolloutTeacherConfig, loss_config: DistillationLossCo headers["Authorization"] = f"Bearer {config.api_key}" self._client = httpx.AsyncClient(headers=headers, timeout=config.request_timeout_s) + async def aclose(self) -> None: + """Close the underlying HTTP client.""" + await self._client.aclose() + async def compute_logprobs(self, state: RolloutState) -> RolloutState: start = time.perf_counter() try: @@ -460,12 +465,91 @@ def _validate_response_logprobs( return teacher_tokens, teacher_logprobs -def route_rollout_teacher_client( - state: RolloutState, - *, - data_source_teacher_map: dict[str, str], - teacher_clients: dict[str, RolloutTeacherClient], -) -> RolloutTeacherClient: - data_source = state.extra_fields["origin_data_source"] - teacher_name = data_source_teacher_map[data_source] - return teacher_clients[teacher_name] +class RolloutTeacherScorer: + """Coordinate rollout-side Teacher routing and asynchronous scoring. + + Args: + teacher_clients (dict[str, RolloutTeacherClient]): Teacher clients keyed + by configured Teacher name. + data_source_teacher_map (dict[str, str]): Mapping from rollout data + source to Teacher name. + defers_to_filter (bool): Whether scoring waits for group filtering. + """ + + def __init__( + self, + teacher_clients: dict[str, RolloutTeacherClient], + data_source_teacher_map: dict[str, str], + *, + defers_to_filter: bool, + ) -> None: + self._teacher_clients = teacher_clients + self._data_source_teacher_map = data_source_teacher_map + self._defers_to_filter = defers_to_filter + self._closed = False + + @classmethod + def from_distillation_config( + cls, + distillation_config: DistillationConfig | None, + *, + defers_to_filter: bool, + ) -> RolloutTeacherScorer: + """Build a scorer from the rollout-side portion of distillation + config.""" + if distillation_config is None or not distillation_config.rollout_teachers: + return cls.disabled() + + teacher_clients = { + teacher.name: RolloutTeacherClient(teacher, distillation_config.loss_config) + for teacher in distillation_config.rollout_teachers + } + return cls( + teacher_clients=teacher_clients, + data_source_teacher_map=dict(distillation_config.data_source_teacher_map), + defers_to_filter=defers_to_filter, + ) + + @staticmethod + def disabled() -> RolloutTeacherScorer: + """Return a scorer that leaves all rollout states unchanged.""" + return RolloutTeacherScorer({}, {}, defers_to_filter=False) + + async def on_sample_ready(self, state: RolloutState) -> RolloutState: + """Score a completed sample when Teacher scoring does not defer to + filtering.""" + if self._closed or self._defers_to_filter or state.status != Status.COMPLETED or not self._teacher_clients: + return state + return await self._compute_logprobs(state) + + async def on_group_ready(self, group: list[RolloutState]) -> list[RolloutState]: + """Score a completed group after validity filtering when configured to + defer.""" + if ( + self._closed + or not self._defers_to_filter + or not self._teacher_clients + or get_group_status(group) != Status.COMPLETED + ): + return group + return list( + await asyncio.gather( + *(create_task(self._compute_logprobs(state)) for state in group), + ) + ) + + async def aclose(self) -> None: + """Close all HTTP clients owned by this scorer.""" + if self._closed: + return + self._closed = True + await asyncio.gather(*(client.aclose() for client in self._teacher_clients.values())) + + async def _compute_logprobs(self, state: RolloutState) -> RolloutState: + teacher = self._route_teacher_client(state) + return await teacher.compute_logprobs(state) + + def _route_teacher_client(self, state: RolloutState) -> RolloutTeacherClient: + data_source = state.extra_fields["origin_data_source"] + teacher_name = self._data_source_teacher_map[data_source] + return self._teacher_clients[teacher_name] diff --git a/xtuner/v1/rl/distillation/train_teacher_manager.py b/xtuner/v1/rl/distillation/train_teacher_manager.py index b273241982..a8f45d2e47 100644 --- a/xtuner/v1/rl/distillation/train_teacher_manager.py +++ b/xtuner/v1/rl/distillation/train_teacher_manager.py @@ -233,6 +233,9 @@ def _compute_topk_targets( with self._teacher_on_device(teacher, timings, teacher_name): # Every rank forwards every local pack before selecting routed # tokens so all ranks enter the same FSDP collective sequence. + # TODO(perf): Use a collective-safe compacted-token path to + # avoid computing teacher logprobs for tokens routed to other + # Teachers. for batch_index, (seq_ctx, teacher_indices) in enumerate(zip(seq_ctx_list, teacher_indices_list)): loss_ctx = topk_logprob_config.build(data={}) assert loss_ctx is not None @@ -267,6 +270,10 @@ def _compute_sampled_logprobs( # Keep the Teacher-major schedule identical across ranks for FSDP. for teacher_index, (teacher_name, teacher) in enumerate(zip(self._teacher_names, self._teachers)): with self._teacher_on_device(teacher, timings, teacher_name): + # TODO(perf): Forward only tokens routed to this Teacher once + # a collective-safe compacted-token path is available; the + # current full-pack forward computes logprobs for other Teachers + # as well. for batch_index, (seq_ctx, shifted_labels, teacher_indices) in enumerate( zip(seq_ctx_list, shifted_labels_list, teacher_indices_list) ): diff --git a/xtuner/v1/rl/loss/__init__.py b/xtuner/v1/rl/loss/__init__.py index 45d3e1d069..d796db0859 100644 --- a/xtuner/v1/rl/loss/__init__.py +++ b/xtuner/v1/rl/loss/__init__.py @@ -14,4 +14,5 @@ ) from .grpo_loss import GRPOLossConfig, GRPOLossContext, GRPOLossKwargs from .loss_fn import check_config, get_policy_loss_fn, kl_penalty, pg_loss_fn, register_policy_loss, sft_loss_fn +from .metrics import finalize_train_metrics from .oreal_loss import OrealLossConfig, OrealLossContext, OrealLossKwargs diff --git a/xtuner/v1/rl/loss/metrics.py b/xtuner/v1/rl/loss/metrics.py new file mode 100644 index 0000000000..0b8b93b594 --- /dev/null +++ b/xtuner/v1/rl/loss/metrics.py @@ -0,0 +1,18 @@ +from typing import Any + +import torch + +from .base_loss import BaseRLLossConfig, finalize_train_policy_metrics +from .distillation_loss import DistillationLossConfig, finalize_distillation_metrics + + +def finalize_train_metrics( + extra_info_dict: dict[str, Any], + device: str | torch.device, + loss_cfg: BaseRLLossConfig, +) -> dict[str, Any]: + """Finalize policy metrics and optional distillation metrics.""" + extra_info_dict = finalize_train_policy_metrics(extra_info_dict, device) + if isinstance(loss_cfg, DistillationLossConfig): + extra_info_dict = finalize_distillation_metrics(extra_info_dict, device) + return extra_info_dict diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index c2d2b79512..7b31f26bae 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -47,9 +47,7 @@ from xtuner.v1.rl.loss import ( BaseRLLossConfig, BaseRLLossContext, - DistillationLossConfig, - finalize_distillation_metrics, - finalize_train_policy_metrics, + finalize_train_metrics, kl_penalty, ) from xtuner.v1.rl.model_utils import build_frozen_model @@ -961,9 +959,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo for k, v in extra_info_dict.items() if isinstance(v, (torch.Tensor, int, float)) } - extra_info_dict = finalize_train_policy_metrics(extra_info_dict, DEVICE) - if isinstance(loss_cfg, DistillationLossConfig): - extra_info_dict = finalize_distillation_metrics(extra_info_dict, DEVICE) + extra_info_dict = finalize_train_metrics(extra_info_dict, DEVICE, loss_cfg) train_step_info.pop("total_loss") # type: ignore[misc] max_memory = DEVICE_MODULE.max_memory_allocated() / (1024**3) # type: ignore[attr-defined] reserved_memory = DEVICE_MODULE.max_memory_reserved() / (1024**3) # type: ignore[attr-defined] From 400d6e83876faf26b524671755e45fcd067d50aa Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Thu, 10 Sep 2026 09:00:24 +0000 Subject: [PATCH 10/20] Refactor teacher scorer configuration boundaries --- xtuner/v1/rl/agent_loop/agent_loop.py | 47 ++++--- .../agent_loop_manager/agent_loop_manager.py | 6 +- .../disagg_agent_loop_manager.py | 6 +- xtuner/v1/rl/distillation/__init__.py | 6 + xtuner/v1/rl/distillation/config.py | 119 +++++++++++++++++- .../distillation/rollout_teacher_manager.py | 32 ++--- .../rl/distillation/train_teacher_manager.py | 26 ++-- xtuner/v1/rl/trainer/worker.py | 14 ++- xtuner/v1/train/rl_trainer.py | 26 ++-- 9 files changed, 201 insertions(+), 81 deletions(-) diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index f23db62274..1a82ecf127 100644 --- a/xtuner/v1/rl/agent_loop/agent_loop.py +++ b/xtuner/v1/rl/agent_loop/agent_loop.py @@ -12,8 +12,8 @@ from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status from xtuner.v1.rl.distillation import ( - DistillationConfig, RolloutTeacherScorer, + RolloutTeacherScorerConfig, ) from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController @@ -75,7 +75,7 @@ def build( logger=None, *, is_valid_sample_fn: IsValidSampleFn | None = None, - distillation_config: DistillationConfig | None = None, + rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None, ) -> AgentLoopSpec: if self.cpu_resources is None: agent_loop = self.build_local( @@ -84,7 +84,13 @@ def build( logger=logger, ) agent_loop.is_valid_sample_fn = is_valid_sample_fn - agent_loop.configure_distillation(distillation_config) + agent_loop._set_teacher_scorer( + rollout_teacher_scorer_config.build( + defers_to_filter=is_valid_sample_fn is not None, + ) + if rollout_teacher_scorer_config is not None + else RolloutTeacherScorer.disabled(), + ) return agent_loop concurrency = AGENT_LOOP_RAY_GENERATE_MAX_CONCURRENCY @@ -102,7 +108,7 @@ def build( judger=judger, logger=logger, is_valid_sample_fn=is_valid_sample_fn, - distillation_config=distillation_config, + rollout_teacher_scorer_config=rollout_teacher_scorer_config, ) return self._build_ray_actor( rollout_controller=rollout_controller, @@ -111,7 +117,7 @@ def build( judger=judger, logger=logger, is_valid_sample_fn=is_valid_sample_fn, - distillation_config=distillation_config, + rollout_teacher_scorer_config=rollout_teacher_scorer_config, ) @abstractmethod @@ -131,7 +137,7 @@ def _build_ray_actor( judger: Judger | None = None, logger=None, is_valid_sample_fn: IsValidSampleFn | None = None, - distillation_config: DistillationConfig | None = None, + rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None, ) -> RayAgentLoopProxy: ray_agent_loop = ray.remote( concurrency_groups={ @@ -151,7 +157,7 @@ def _build_ray_actor( actor_memory=cpu_resources.cpu_memory_per_worker, capture_child_tasks=True, is_valid_sample_fn=is_valid_sample_fn, - distillation_config=distillation_config, + rollout_teacher_scorer_config=rollout_teacher_scorer_config, ), ) @@ -165,7 +171,7 @@ def _build_ray_actors( logger=None, start_bundle_idx: int = 0, is_valid_sample_fn: IsValidSampleFn | None = None, - distillation_config: DistillationConfig | None = None, + rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None, ) -> list[RayAgentLoopProxy]: ray_agent_loop = ray.remote( concurrency_groups={ @@ -186,7 +192,7 @@ def _build_ray_actors( actor_memory_per_worker=cpu_resources.cpu_memory_per_worker, capture_child_tasks=True, is_valid_sample_fn=is_valid_sample_fn, - distillation_config=distillation_config, + rollout_teacher_scorer_config=rollout_teacher_scorer_config, ), ) @@ -200,7 +206,7 @@ def _build_router( logger=None, start_bundle_idx: int = 0, is_valid_sample_fn: IsValidSampleFn | None = None, - distillation_config: DistillationConfig | None = None, + rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None, ) -> RouterAgentLoop: return RouterAgentLoop( workers=self._build_ray_actors( @@ -212,7 +218,7 @@ def _build_router( logger=logger, start_bundle_idx=start_bundle_idx, is_valid_sample_fn=is_valid_sample_fn, - distillation_config=distillation_config, + rollout_teacher_scorer_config=rollout_teacher_scorer_config, ), rollout_ctl=rollout_controller, ) @@ -243,13 +249,8 @@ def __init__( self._judger_pause_event = asyncio.Event() self._teacher_scorer = RolloutTeacherScorer.disabled() - def configure_distillation(self, distillation_config: DistillationConfig | None) -> None: - if distillation_config is None: - return - self._teacher_scorer = RolloutTeacherScorer.from_distillation_config( - distillation_config, - defers_to_filter=self.is_valid_sample_fn is not None, - ) + def _set_teacher_scorer(self, teacher_scorer: RolloutTeacherScorer) -> None: + self._teacher_scorer = teacher_scorer @abstractmethod async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: ... @@ -393,7 +394,7 @@ def __init__( logger=None, is_valid_sample_fn: IsValidSampleFn | None = None, *, - distillation_config: DistillationConfig | None = None, + rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None, ): self.agent_loop = agent_loop_config.build_local( rollout_controller=rollout_controller, @@ -401,7 +402,13 @@ def __init__( logger=logger, ) self.agent_loop.is_valid_sample_fn = is_valid_sample_fn - self.agent_loop.configure_distillation(distillation_config) + self.agent_loop._set_teacher_scorer( + rollout_teacher_scorer_config.build( + defers_to_filter=is_valid_sample_fn is not None, + ) + if rollout_teacher_scorer_config is not None + else RolloutTeacherScorer.disabled(), + ) @ray_method(concurrency_group=AGENT_LOOP_CONCURRENCY_GROUP_GENERATE) async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: diff --git a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py index 733825e031..831c5d0639 100644 --- a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py @@ -9,7 +9,7 @@ from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast from xtuner.v1.data_proto.rl_data import Status from xtuner.v1.rl.agent_loop import AgentLoopConfig, IsValidSampleFn -from xtuner.v1.rl.distillation import DistillationConfig +from xtuner.v1.rl.distillation import RolloutTeacherScorerConfig from xtuner.v1.rl.judger import ComposedJudgerConfig, JudgerConfig, build_judger from xtuner.v1.rl.replay_buffer import ReplayBuffer from xtuner.v1.rl.rollout import RolloutController @@ -130,7 +130,7 @@ def build( replay_buffer: ReplayBuffer, logger=None, sync_weights_interval: int = 1, - distillation_config: DistillationConfig | None = None, + rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None, ) -> "AgentLoopManager": tasks = self.tasks if isinstance(self.tasks, list) else [self.tasks] if not tasks: @@ -148,7 +148,7 @@ def build( judger=build_judger(task_cfg.judger_config) if task_cfg.judger_config is not None else None, logger=logger, is_valid_sample_fn=task_cfg.is_valid_sample_fn, - distillation_config=distillation_config, + rollout_teacher_scorer_config=rollout_teacher_scorer_config, ) produce_strategy = task_cfg.produce_strategy_config.build( sync_weights_interval=sync_weights_interval, diff --git a/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py index b6f62b6f80..8d896a2950 100644 --- a/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py @@ -10,7 +10,7 @@ from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast from xtuner.v1.data_proto.rl_data import Status from xtuner.v1.rl.agent_loop import AgentLoopConfig, IsValidSampleFn -from xtuner.v1.rl.distillation import DistillationConfig +from xtuner.v1.rl.distillation import RolloutTeacherScorerConfig from xtuner.v1.rl.judger import ComposedJudgerConfig, JudgerConfig, build_judger from xtuner.v1.rl.replay_buffer import ReplayBuffer from xtuner.v1.rl.rollout import RolloutController @@ -70,7 +70,7 @@ def build( replay_buffer: ReplayBuffer, logger=None, sync_weights_interval: int = 1, - distillation_config: DistillationConfig | None = None, + rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None, ) -> "DisaggAgentLoopManager": tasks = self.tasks if isinstance(self.tasks, list) else [self.tasks] if not tasks: @@ -88,7 +88,7 @@ def build( judger=build_judger(task_cfg.judger_config) if task_cfg.judger_config is not None else None, logger=logger, is_valid_sample_fn=task_cfg.is_valid_sample_fn, - distillation_config=distillation_config, + rollout_teacher_scorer_config=rollout_teacher_scorer_config, ) produce_strategy = task_cfg.produce_strategy_config.build( sync_weights_interval=sync_weights_interval, diff --git a/xtuner/v1/rl/distillation/__init__.py b/xtuner/v1/rl/distillation/__init__.py index 4e92b7bf80..0f8bcea864 100644 --- a/xtuner/v1/rl/distillation/__init__.py +++ b/xtuner/v1/rl/distillation/__init__.py @@ -2,8 +2,11 @@ DistillationConfig, RolloutTeacherConfig, RolloutTeacherLaunchConfig, + RolloutTeacherScorerConfig, TeacherConfig, + TeacherTargetConfig, TrainTeacherConfig, + TrainTeacherManagerConfig, ) from .rollout_teacher_manager import ( RolloutTeacherClient, @@ -18,7 +21,10 @@ "DistillationConfig", "RolloutTeacherConfig", "RolloutTeacherLaunchConfig", + "RolloutTeacherScorerConfig", "TeacherConfig", + "TeacherTargetConfig", + "TrainTeacherManagerConfig", "TrainTeacherConfig", "RolloutTeacherClient", "RolloutTeacherReplicaRouter", diff --git a/xtuner/v1/rl/distillation/config.py b/xtuner/v1/rl/distillation/config.py index f9557637db..c4b77faaa8 100644 --- a/xtuner/v1/rl/distillation/config.py +++ b/xtuner/v1/rl/distillation/config.py @@ -1,7 +1,7 @@ from __future__ import annotations from pathlib import Path -from typing import Literal, cast +from typing import TYPE_CHECKING, Literal, cast from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator @@ -11,6 +11,11 @@ from xtuner.v1.rl.loss.distillation_loss import DistillationLossConfig +if TYPE_CHECKING: + from .rollout_teacher_manager import RolloutTeacherScorer + from .train_teacher_manager import TrainTeacherManager + + class RolloutTeacherLaunchConfig(BaseModel): model_config = ConfigDict(extra="forbid") @@ -58,6 +63,91 @@ class TrainTeacherConfig(BaseModel): TeacherConfig = RolloutTeacherConfig | TrainTeacherConfig +class TeacherTargetConfig(BaseModel): + """Configuration for the output protocol requested from a Teacher.""" + + model_config = ConfigDict(extra="forbid") + + mode: Literal["sampled_token", "topk"] + top_k: int | None = Field(default=None, gt=0) + + @model_validator(mode="after") + def validate_target_config(self) -> TeacherTargetConfig: + if self.mode == "topk" and self.top_k is None: + raise ValueError("top_k is required for top-k Teacher targets") + if self.mode == "sampled_token" and self.top_k is not None: + raise ValueError("top_k is only valid for top-k Teacher targets") + return self + + @property + def uses_sampled_token_targets(self) -> bool: + return self.mode == "sampled_token" + + @property + def uses_topk_targets(self) -> bool: + return self.mode == "topk" + + +class RolloutTeacherScorerConfig(BaseModel): + """Serializable configuration used to build a rollout Teacher scorer.""" + + model_config = ConfigDict(extra="forbid") + + teachers: list[RolloutTeacherConfig] = Field(min_length=1) + data_source_teacher_map: dict[str, str] = Field(min_length=1) + target_config: TeacherTargetConfig + + @model_validator(mode="after") + def validate_config(self) -> RolloutTeacherScorerConfig: + teacher_names_list = [teacher.name for teacher in self.teachers] + if len(teacher_names_list) != len(set(teacher_names_list)): + raise ValueError("Rollout Teacher names must be unique") + teacher_names = set(teacher_names_list) + unknown_teachers = set(self.data_source_teacher_map.values()) - teacher_names + if unknown_teachers: + raise ValueError(f"data_source_teacher_map references unknown teachers: {sorted(unknown_teachers)}") + return self + + def build(self, *, defers_to_filter: bool) -> RolloutTeacherScorer: + from .rollout_teacher_manager import RolloutTeacherScorer + + return RolloutTeacherScorer.from_config(self, defers_to_filter=defers_to_filter) + + +class TrainTeacherManagerConfig(BaseModel): + """Serializable configuration used to build a training Teacher manager.""" + + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) + + teachers: list[TrainTeacherConfig] = Field(min_length=1) + data_source_teacher_map: dict[str, str] = Field(min_length=1) + target_config: TeacherTargetConfig + + @model_validator(mode="after") + def validate_config(self) -> TrainTeacherManagerConfig: + teacher_names_list = [teacher.name for teacher in self.teachers] + if len(teacher_names_list) != len(set(teacher_names_list)): + raise ValueError("Train Teacher names must be unique") + teacher_names = set(teacher_names_list) + unknown_teachers = set(self.data_source_teacher_map.values()) - teacher_names + if unknown_teachers: + raise ValueError(f"data_source_teacher_map references unknown teachers: {sorted(unknown_teachers)}") + return self + + @property + def teacher_index_by_data_source(self) -> dict[str, int]: + teacher_index_by_name = {teacher.name: index for index, teacher in enumerate(self.teachers)} + return { + data_source: teacher_index_by_name[teacher_name] + for data_source, teacher_name in self.data_source_teacher_map.items() + } + + def build(self, *, chunk_size: int | None) -> TrainTeacherManager: + from .train_teacher_manager import TrainTeacherManager + + return TrainTeacherManager(self, chunk_size=chunk_size) + + class DistillationConfig(BaseModel): """Distillation objective, Teacher runtimes, and data-source routing.""" @@ -102,6 +192,33 @@ def train_teachers(self) -> list[TrainTeacherConfig]: def teacher_index_by_data_source(self) -> dict[str, int]: return self._teacher_index_by_data_source + @property + def teacher_target_config(self) -> TeacherTargetConfig: + return TeacherTargetConfig( + mode="topk" if self.loss_config.uses_topk_targets else "sampled_token", + top_k=self.loss_config.top_k if self.loss_config.uses_topk_targets else None, + ) + + @property + def rollout_teacher_scorer_config(self) -> RolloutTeacherScorerConfig | None: + if not self.rollout_teachers: + return None + return RolloutTeacherScorerConfig( + teachers=list(self.rollout_teachers), + data_source_teacher_map=dict(self.data_source_teacher_map), + target_config=self.teacher_target_config, + ) + + @property + def train_teacher_manager_config(self) -> TrainTeacherManagerConfig | None: + if not self.train_teachers: + return None + return TrainTeacherManagerConfig( + teachers=list(self.train_teachers), + data_source_teacher_map=dict(self.data_source_teacher_map), + target_config=self.teacher_target_config, + ) + def validate_student_model(self, student_model_cfg: TransformerConfig | BaseComposeConfig) -> None: """Validate contracts that require both Student and TrainTeacher configs.""" diff --git a/xtuner/v1/rl/distillation/rollout_teacher_manager.py b/xtuner/v1/rl/distillation/rollout_teacher_manager.py index 0bcade5a66..3d3222899e 100644 --- a/xtuner/v1/rl/distillation/rollout_teacher_manager.py +++ b/xtuner/v1/rl/distillation/rollout_teacher_manager.py @@ -10,11 +10,10 @@ import httpx from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status -from xtuner.v1.rl.loss import DistillationLossConfig from xtuner.v1.rl.utils import create_task from xtuner.v1.utils import get_logger -from .config import DistillationConfig, RolloutTeacherConfig +from .config import RolloutTeacherConfig, RolloutTeacherScorerConfig, TeacherTargetConfig logger = get_logger() @@ -66,12 +65,12 @@ def resolve_replica_idx( class RolloutTeacherClient: """Asynchronous teacher client scoped to one AgentLoop.""" - def __init__(self, config: RolloutTeacherConfig, loss_config: DistillationLossConfig) -> None: + def __init__(self, config: RolloutTeacherConfig, target_config: TeacherTargetConfig) -> None: self.config = config - self.loss_config = loss_config + self.target_config = target_config self.name = config.name self.backend = self._resolve_backend_from_env() - if self.loss_config.uses_topk_targets and self.backend != "lmdeploy": + if self.target_config.uses_topk_targets and self.backend != "lmdeploy": raise RuntimeError("Rollout Teacher Top-K targets currently require LMDeploy") self.urls = [f"{endpoint.rstrip('/')}/generate" for endpoint in config.endpoints] self._semaphores = [asyncio.Semaphore(config.max_concurrency) for _ in self.urls] @@ -126,7 +125,7 @@ async def compute_logprobs(self, state: RolloutState) -> RolloutState: response.raise_for_status() teacher_tokens: list[int] | list[list[int]] teacher_logprobs: list[float] | list[list[float]] - if self.loss_config.uses_topk_targets: + if self.target_config.uses_topk_targets: teacher_tokens, teacher_logprobs = self._parse_topk_response( response, response_ids, @@ -226,8 +225,8 @@ def _construct_payload( raise RuntimeError(f"Unsupported teacher backend: {self.backend}") if image_data: payload["image_data"] = image_data - if self.loss_config.uses_topk_targets: - payload["top_logprobs_num"] = cast(int, self.loss_config.top_k) + if self.target_config.uses_topk_targets: + payload["top_logprobs_num"] = cast(int, self.target_config.top_k) return payload @staticmethod @@ -322,7 +321,7 @@ def _parse_topk_response( f"{len(raw_topk)} vs {len(response_ids)}" ) - top_k = cast(int, self.loss_config.top_k) + top_k = cast(int, self.target_config.top_k) teacher_tokens: list[list[int]] = [] teacher_logprobs: list[list[float]] = [] for row_idx, row in enumerate(raw_topk[-len(response_ids) :]): @@ -489,24 +488,19 @@ def __init__( self._closed = False @classmethod - def from_distillation_config( + def from_config( cls, - distillation_config: DistillationConfig | None, + config: RolloutTeacherScorerConfig, *, defers_to_filter: bool, ) -> RolloutTeacherScorer: - """Build a scorer from the rollout-side portion of distillation - config.""" - if distillation_config is None or not distillation_config.rollout_teachers: - return cls.disabled() - + """Build a scorer from its serializable configuration.""" teacher_clients = { - teacher.name: RolloutTeacherClient(teacher, distillation_config.loss_config) - for teacher in distillation_config.rollout_teachers + teacher.name: RolloutTeacherClient(teacher, config.target_config) for teacher in config.teachers } return cls( teacher_clients=teacher_clients, - data_source_teacher_map=dict(distillation_config.data_source_teacher_map), + data_source_teacher_map=dict(config.data_source_teacher_map), defers_to_filter=defers_to_filter, ) diff --git a/xtuner/v1/rl/distillation/train_teacher_manager.py b/xtuner/v1/rl/distillation/train_teacher_manager.py index a8f45d2e47..b180e7c567 100644 --- a/xtuner/v1/rl/distillation/train_teacher_manager.py +++ b/xtuner/v1/rl/distillation/train_teacher_manager.py @@ -10,11 +10,10 @@ from xtuner.v1.data_proto.sequence_context import SequenceContext from xtuner.v1.loss import LogProbConfig, LogProbContext, TopKLogProbConfig from xtuner.v1.model.compose.base import BaseComposeConfig -from xtuner.v1.rl.loss import DistillationLossConfig from xtuner.v1.rl.model_utils import FrozenModel, build_frozen_model from xtuner.v1.utils import get_device, get_torch_device_module -from .config import DistillationConfig +from .config import TrainTeacherManagerConfig DEVICE = get_device() @@ -67,14 +66,14 @@ class TrainTeacherManager: swapping the Actor and optimizer around the Teacher phase. """ - def __init__(self, distillation_config: DistillationConfig, *, chunk_size: int | None) -> None: - self.loss_config = cast(DistillationLossConfig, distillation_config.loss_config) + def __init__(self, config: TrainTeacherManagerConfig, *, chunk_size: int | None) -> None: + self.target_config = config.target_config mode = "chunk" if chunk_size is not None else "eager" self.logprob_config = LogProbConfig(chunk_size=chunk_size, mode=mode) self.topk_logprob_config: TopKLogProbConfig | None = None - if self.loss_config.uses_topk_targets: + if self.target_config.uses_topk_targets: self.topk_logprob_config = TopKLogProbConfig( - top_k=cast(int, self.loss_config.top_k), + top_k=cast(int, self.target_config.top_k), chunk_size=chunk_size, mode=mode, ) @@ -83,7 +82,7 @@ def __init__(self, distillation_config: DistillationConfig, *, chunk_size: int | # errors fail during worker initialization. Each Teacher is offloaded # immediately, preventing multiple full models from co-residing on GPU. self._teachers: list[FrozenModel] = [] - for teacher_config in distillation_config.train_teachers: + for teacher_config in config.teachers: teacher = build_frozen_model( teacher_config.model_cfg, teacher_config.model_path, @@ -91,14 +90,11 @@ def __init__(self, distillation_config: DistillationConfig, *, chunk_size: int | ) self._teachers.append(teacher) - self._teacher_is_composed = [ - isinstance(teacher.model_cfg, BaseComposeConfig) for teacher in distillation_config.train_teachers - ] + self._teacher_is_composed = [isinstance(teacher.model_cfg, BaseComposeConfig) for teacher in config.teachers] self._teacher_index_by_name = { - teacher_config.name: teacher_index - for teacher_index, teacher_config in enumerate(distillation_config.train_teachers) + teacher_config.name: teacher_index for teacher_index, teacher_config in enumerate(config.teachers) } - self._teacher_names = [teacher_config.name for teacher_config in distillation_config.train_teachers] + self._teacher_names = [teacher_config.name for teacher_config in config.teachers] def compute_logprobs( self, @@ -108,7 +104,7 @@ def compute_logprobs( teacher_indices_list: list[torch.Tensor], ) -> TrainTeacherOutputs: timings = TrainTeacherTimings() - if self.loss_config.uses_sampled_token_targets: + if self.target_config.uses_sampled_token_targets: return TrainTeacherOutputs( teacher_logprobs=self._compute_sampled_logprobs( seq_ctx_list, @@ -215,7 +211,7 @@ def _compute_topk_targets( teacher_indices_list: list[torch.Tensor], timings: TrainTeacherTimings, ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: - top_k = cast(int, self.loss_config.top_k) + top_k = cast(int, self.target_config.top_k) topk_logprob_config = cast(TopKLogProbConfig, self.topk_logprob_config) target_ids = [ torch.zeros((*teacher_indices.shape, top_k), dtype=torch.long, device=DEVICE) diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index 7b31f26bae..0a11eb2384 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -43,7 +43,11 @@ from xtuner.v1.model.compose.base import BaseComposeConfig from xtuner.v1.model.utils.misc import ModelForwardExtraLogInfo from xtuner.v1.profiler import profiling_memory, profiling_time -from xtuner.v1.rl.distillation import DistillationConfig, TrainTeacherManager, TrainTeacherTimings +from xtuner.v1.rl.distillation import ( + TrainTeacherManager, + TrainTeacherManagerConfig, + TrainTeacherTimings, +) from xtuner.v1.rl.loss import ( BaseRLLossConfig, BaseRLLossContext, @@ -163,7 +167,7 @@ class WorkerConfig(BaseModel): free_rollout_routed_experts_in_worker: bool = True # 默认不需要用户配置 offload_rollout_routed_experts: bool = False offload_old_logprobs: bool = False - distillation_config: DistillationConfig | None = None + train_teacher_manager_config: TrainTeacherManagerConfig | None = None # sft config sft_dataloader_cfg: DataloaderConfig | None = None @@ -274,10 +278,8 @@ def __init__( ) self._train_teacher_manager: TrainTeacherManager | None = None - distillation_config = worker_cfg.distillation_config - if distillation_config is not None and distillation_config.train_teachers: - self._train_teacher_manager = TrainTeacherManager( - distillation_config, + if worker_cfg.train_teacher_manager_config is not None: + self._train_teacher_manager = worker_cfg.train_teacher_manager_config.build( chunk_size=worker_cfg.loss_cfg.chunk_size, ) diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 989e67d21a..1f00b7efcb 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -738,18 +738,14 @@ def _init_common(self, cfg: BaseRLTrainerConfig, *, meta_path: str, logger_tag: cfg.distillation_config = cfg.distillation_config.resolve_teacher_endpoints(endpoint_map) self._distillation_config = cfg.distillation_config - self._distillation_loss_cfg = ( - self._distillation_config.loss_config if self._distillation_config is not None else None + self._rollout_teacher_scorer_config = ( + self._distillation_config.rollout_teacher_scorer_config if self._distillation_config is not None else None ) - self._train_teacher_config = ( - self._distillation_config - if self._distillation_config is not None and self._distillation_config.train_teachers - else None + self._train_teacher_manager_config = ( + self._distillation_config.train_teacher_manager_config if self._distillation_config is not None else None ) - self._rollout_teacher_config = ( - self._distillation_config - if self._distillation_config is not None and self._distillation_config.rollout_teachers - else None + self._distillation_loss_cfg = ( + self._distillation_config.loss_config if self._distillation_config is not None else None ) check_fa3() @@ -855,7 +851,7 @@ def _init_train_worker_config(self, cfg: BaseRLTrainerConfig, log_dir: Path) -> cfg.train_worker_cfg.free_rollout_routed_experts_in_worker = False cfg.train_worker_cfg.load_from = cfg.load_from cfg.train_worker_cfg.log_dir = log_dir - cfg.train_worker_cfg.distillation_config = cfg.distillation_config + cfg.train_worker_cfg.train_teacher_manager_config = self._train_teacher_manager_config self._train_worker_cfg = cfg.train_worker_cfg def _init_rollout_config(self, cfg: BaseRLTrainerConfig, log_dir: Path) -> None: @@ -889,7 +885,7 @@ def _build_agent_loop_components(self, cfg: BaseRLTrainerConfig, replay_buffer) replay_buffer=replay_buffer, logger=self.logger, sync_weights_interval=cfg.sync_weights_interval, - distillation_config=self._rollout_teacher_config, + rollout_teacher_scorer_config=self._rollout_teacher_scorer_config, ) self.agent_loop_manager = cast(AgentLoopManager | DisaggAgentLoopManager, agent_loop_manager) @@ -904,7 +900,7 @@ def _build_agent_loop_components(self, cfg: BaseRLTrainerConfig, replay_buffer) replay_buffer=replay_buffer, logger=self.logger, sync_weights_interval=cfg.sync_weights_interval, - distillation_config=None, + rollout_teacher_scorer_config=None, ), ) @@ -1240,7 +1236,9 @@ def _prepare_train_data( data_batches = [] teacher_index_by_data_source = ( - self._train_teacher_config.teacher_index_by_data_source if self._train_teacher_config is not None else None + self._train_teacher_manager_config.teacher_index_by_data_source + if self._train_teacher_manager_config is not None + else None ) for j, group in enumerate(data_groups): From f28bdbf1c7254db67b3694e593d9734a7805fca0 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Fri, 11 Sep 2026 03:28:43 +0000 Subject: [PATCH 11/20] Simplify loss metric finalization --- .../rl/distillation/train_teacher_manager.py | 50 ++++++++++- xtuner/v1/rl/loss/__init__.py | 1 - xtuner/v1/rl/loss/base_loss.py | 5 +- xtuner/v1/rl/loss/distillation_loss.py | 16 ++++ xtuner/v1/rl/loss/metrics.py | 18 ---- xtuner/v1/rl/trainer/controller.py | 2 +- xtuner/v1/rl/trainer/worker.py | 86 +++++-------------- xtuner/v1/train/rl_trainer.py | 23 +---- 8 files changed, 90 insertions(+), 111 deletions(-) delete mode 100644 xtuner/v1/rl/loss/metrics.py diff --git a/xtuner/v1/rl/distillation/train_teacher_manager.py b/xtuner/v1/rl/distillation/train_teacher_manager.py index b180e7c567..9b5b5eb8bb 100644 --- a/xtuner/v1/rl/distillation/train_teacher_manager.py +++ b/xtuner/v1/rl/distillation/train_teacher_manager.py @@ -3,7 +3,7 @@ import time from contextlib import contextmanager from dataclasses import dataclass, field -from typing import Iterator, cast +from typing import TYPE_CHECKING, Any, Iterator, cast import torch @@ -16,6 +16,10 @@ from .config import TrainTeacherManagerConfig +if TYPE_CHECKING: + from xtuner.v1.rl.loss.base_loss import BaseRLLossContext + + DEVICE = get_device() DEVICE_MODULE = get_torch_device_module() @@ -62,8 +66,8 @@ class TrainTeacherManager: """Execute training-side Teachers in one TrainingWorker process. The manager owns Teacher model construction, deterministic Teacher-major scheduling, CPU/device residency, and - sampled-token or top-k output calculation. The caller remains responsible for preparing distributed inputs and - swapping the Actor and optimizer around the Teacher phase. + sampled-token or top-k output calculation. The caller remains responsible for swapping the Actor and optimizer + around the Teacher phase. """ def __init__(self, config: TrainTeacherManagerConfig, *, chunk_size: int | None) -> None: @@ -126,6 +130,46 @@ def compute_logprobs( timings=timings, ) + def compute_teacher_outputs( + self, + seq_ctx_list: list[SequenceContext], + loss_ctx_list: list[BaseRLLossContext], + ) -> TrainTeacherTimings: + """Compute Teacher outputs for the given loss contexts.""" + teacher_indices_list: list[torch.Tensor] = [] + shifted_labels_list: list[torch.Tensor] = [] + for loss_ctx in loss_ctx_list: + loss_kwargs = cast(Any, loss_ctx.loss_kwargs) + teacher_indices = loss_kwargs.teacher_indices + if teacher_indices is None: + raise ValueError("teacher_indices are required when the training Teacher manager is enabled") + teacher_indices_list.append(teacher_indices) + shifted_labels_list.append(loss_kwargs.shifted_labels) + + try: + outputs = self.compute_logprobs( + seq_ctx_list=seq_ctx_list, + shifted_labels_list=shifted_labels_list, + teacher_indices_list=teacher_indices_list, + ) + target_token_ids_list: list[torch.Tensor | None] = ( + list(outputs.target_token_ids) + if outputs.target_token_ids is not None + else [None] * len(outputs.teacher_logprobs) + ) + for loss_ctx, teacher_logprobs, target_token_ids in zip( + loss_ctx_list, + outputs.teacher_logprobs, + target_token_ids_list, + ): + loss_kwargs = cast(Any, loss_ctx.loss_kwargs) + loss_kwargs.teacher_logprobs = teacher_logprobs + if target_token_ids is not None: + loss_kwargs.target_token_ids = target_token_ids + return outputs.timings + finally: + self.offload_all_to_cpu() + def offload_all_to_cpu(self) -> None: for teacher in self._teachers: self._offload_to_cpu(teacher) diff --git a/xtuner/v1/rl/loss/__init__.py b/xtuner/v1/rl/loss/__init__.py index d796db0859..45d3e1d069 100644 --- a/xtuner/v1/rl/loss/__init__.py +++ b/xtuner/v1/rl/loss/__init__.py @@ -14,5 +14,4 @@ ) from .grpo_loss import GRPOLossConfig, GRPOLossContext, GRPOLossKwargs from .loss_fn import check_config, get_policy_loss_fn, kl_penalty, pg_loss_fn, register_policy_loss, sft_loss_fn -from .metrics import finalize_train_metrics from .oreal_loss import OrealLossConfig, OrealLossContext, OrealLossKwargs diff --git a/xtuner/v1/rl/loss/base_loss.py b/xtuner/v1/rl/loss/base_loss.py index e2eccc9676..7805b967f1 100644 --- a/xtuner/v1/rl/loss/base_loss.py +++ b/xtuner/v1/rl/loss/base_loss.py @@ -94,6 +94,10 @@ def loss_ctx_cls(self) -> type["BaseRLLossContext"]: def _loss_kwargs_cls(self) -> type["BaseRLLossKwargs"]: raise NotImplementedError + def finalize_metrics(self, extra_info_dict: dict[str, Any], device: str | torch.device) -> dict[str, Any]: + """Finalize metrics emitted by the policy loss.""" + return finalize_train_policy_metrics(extra_info_dict, device) + def build( self, data: dict, @@ -304,5 +308,4 @@ def reduce_values(keys, op): extra_info_dict["reduced_train_policy_ratio_min"] = ratio_min # legacy metric，keep here extra_info_dict["max_ratio"] = ratio_max - return extra_info_dict diff --git a/xtuner/v1/rl/loss/distillation_loss.py b/xtuner/v1/rl/loss/distillation_loss.py index de5d358dd6..4f8292f687 100644 --- a/xtuner/v1/rl/loss/distillation_loss.py +++ b/xtuner/v1/rl/loss/distillation_loss.py @@ -108,6 +108,11 @@ def loss_ctx_cls(self) -> type["DistillationLossContext"]: def _loss_kwargs_cls(self) -> type["DistillationLossKwargs"]: return DistillationLossKwargs + def finalize_metrics(self, extra_info_dict: dict[str, Any], device: str | torch.device) -> dict[str, Any]: + """Finalize policy and distillation metrics emitted by this loss.""" + extra_info_dict = super().finalize_metrics(extra_info_dict, device) + return finalize_distillation_metrics(extra_info_dict, device) + def build( self, data: dict, @@ -125,6 +130,7 @@ def build( is_weights=data.get("rollout_is_weights"), teacher_logprobs=data.get("teacher_logprobs"), target_token_ids=data.get("target_token_ids"), + teacher_indices=data.get("teacher_indices"), ) # Rollout teacher targets can be large, so shard them before moving the # local slice to the accelerator. @@ -137,6 +143,7 @@ def build( class DistillationLossKwargs(BaseRLLossKwargs): teacher_logprobs: torch.Tensor | None = None target_token_ids: torch.Tensor | None = None + teacher_indices: torch.Tensor | None = None distillation_loss_weight: torch.Tensor | None = None def sp_split(self, sp_mesh: DeviceMesh) -> Self: @@ -155,6 +162,13 @@ def sp_split(self, sp_mesh: DeviceMesh) -> Self: split_dim=1, padding_value=0, ) + if self.teacher_indices is not None: + self.teacher_indices = sp_split( + self.teacher_indices, + sp_mesh=sp_mesh, + split_dim=1, + padding_value=-1, + ) return self def to(self, device: torch.device | str) -> Self: @@ -163,6 +177,8 @@ def to(self, device: torch.device | str) -> Self: self.teacher_logprobs = self.teacher_logprobs.to(device) if self.target_token_ids is not None: self.target_token_ids = self.target_token_ids.to(device) + if self.teacher_indices is not None: + self.teacher_indices = self.teacher_indices.to(device) if self.distillation_loss_weight is not None: self.distillation_loss_weight = self.distillation_loss_weight.to(device) return self diff --git a/xtuner/v1/rl/loss/metrics.py b/xtuner/v1/rl/loss/metrics.py deleted file mode 100644 index 0b8b93b594..0000000000 --- a/xtuner/v1/rl/loss/metrics.py +++ /dev/null @@ -1,18 +0,0 @@ -from typing import Any - -import torch - -from .base_loss import BaseRLLossConfig, finalize_train_policy_metrics -from .distillation_loss import DistillationLossConfig, finalize_distillation_metrics - - -def finalize_train_metrics( - extra_info_dict: dict[str, Any], - device: str | torch.device, - loss_cfg: BaseRLLossConfig, -) -> dict[str, Any]: - """Finalize policy metrics and optional distillation metrics.""" - extra_info_dict = finalize_train_policy_metrics(extra_info_dict, device) - if isinstance(loss_cfg, DistillationLossConfig): - extra_info_dict = finalize_distillation_metrics(extra_info_dict, device) - return extra_info_dict diff --git a/xtuner/v1/rl/trainer/controller.py b/xtuner/v1/rl/trainer/controller.py index b84ff6c4dc..70ab23a48f 100644 --- a/xtuner/v1/rl/trainer/controller.py +++ b/xtuner/v1/rl/trainer/controller.py @@ -25,7 +25,7 @@ class ColateItem(TypedDict): rollout_logprobs: torch.Tensor | None teacher_logprobs: NotRequired[torch.Tensor | None] target_token_ids: NotRequired[torch.Tensor | None] - teacher_indices: NotRequired[torch.Tensor] + teacher_indices: NotRequired[torch.Tensor | None] def _summarize_process_group_results(results: list[dict[str, Any]]) -> str: diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index 0a11eb2384..099eff2336 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -38,7 +38,6 @@ from xtuner.v1.loss import BaseLossContext, CELossConfig, LogProbConfig from xtuner.v1.loss.ce_loss import CELossContext, LMHeadLossContext from xtuner.v1.loss.mtp_loss import MTPLossConfig, MTPLossContext -from xtuner.v1.loss.utils import sp_split from xtuner.v1.model.base import ModelItem, TransformerConfig from xtuner.v1.model.compose.base import BaseComposeConfig from xtuner.v1.model.utils.misc import ModelForwardExtraLogInfo @@ -51,7 +50,6 @@ from xtuner.v1.rl.loss import ( BaseRLLossConfig, BaseRLLossContext, - finalize_train_metrics, kl_penalty, ) from xtuner.v1.rl.model_utils import build_frozen_model @@ -203,7 +201,7 @@ class WorkerInputItem(TypedDict): rollout_logprobs: torch.Tensor | None teacher_logprobs: NotRequired[torch.Tensor | None] target_token_ids: NotRequired[torch.Tensor | None] - teacher_indices: NotRequired[torch.Tensor] + teacher_indices: NotRequired[torch.Tensor | None] class WorkerTrainLogItem(TypedDict, total=False): @@ -218,9 +216,6 @@ class WorkerTrainLogItem(TypedDict, total=False): class WorkerLogItem(TypedDict): train_entropy: float teacher_timings: NotRequired[dict[str, dict[str, float]]] - teacher_compute_time: NotRequired[float] - teacher_onload_time: NotRequired[float] - teacher_offload_time: NotRequired[float] rollout_entropy: NotRequired[float] mismatch_metrics: NotRequired[dict[str, float]] rollout_is_metrics: NotRequired[dict[str, float]] @@ -498,43 +493,27 @@ def compute_ref_logprobs( self._ref_model.to_device("cpu") return ref_logprobs_list - def _compute_train_teacher_outputs( + def _maybe_compute_train_teacher_outputs( self, seq_ctx_list: list[SequenceContext], - teacher_indices_list: list[torch.Tensor], loss_ctx_list: list[BaseRLLossContext], - ) -> TrainTeacherTimings: + ) -> TrainTeacherTimings | None: + if self._train_teacher_manager is None: + return None + self._engine.put_model_to_device("cpu") self._engine.put_optimizer_to_device("cpu") if hasattr(DEVICE_MODULE, "empty_cache"): DEVICE_MODULE.empty_cache() try: - assert self._train_teacher_manager is not None - outputs = self._train_teacher_manager.compute_logprobs( - seq_ctx_list=seq_ctx_list, - shifted_labels_list=[loss_ctx.loss_kwargs.shifted_labels for loss_ctx in loss_ctx_list], - teacher_indices_list=teacher_indices_list, + teacher_timings = self._train_teacher_manager.compute_teacher_outputs( + seq_ctx_list, + loss_ctx_list, ) - if outputs.target_token_ids is None: - for loss_ctx, teacher_logprobs in zip(loss_ctx_list, outputs.teacher_logprobs): - loss_ctx.loss_kwargs.teacher_logprobs = teacher_logprobs - else: - for loss_ctx, token_ids, teacher_logprobs in zip( - loss_ctx_list, - outputs.target_token_ids, - outputs.teacher_logprobs, - ): - loss_ctx.loss_kwargs.target_token_ids = token_ids - loss_ctx.loss_kwargs.teacher_logprobs = teacher_logprobs - return outputs.timings + return teacher_timings finally: - if self._train_teacher_manager is not None: - self._train_teacher_manager.offload_all_to_cpu() - self._onload_actor_and_optimizer() - - def _onload_actor_and_optimizer(self) -> None: - self._engine.put_model_to_device(DEVICE) - self._engine.put_optimizer_to_device(DEVICE) + self._engine.put_model_to_device(DEVICE) + self._engine.put_optimizer_to_device(DEVICE) def _add_rollout_routed_experts( self, seq_ctx: SequenceContext, rollout_routed_experts: torch.Tensor | list[torch.Tensor | ray.ObjectRef] @@ -645,8 +624,6 @@ def _maybe_profiling(self, global_train_step: int, phase: str): def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLogItem: # NOTE: sglang会清除logger handle, 重新创建 self.logger = get_logger(log_dir=self.log_dir, tag="TrainingWorker") - if self._train_teacher_manager is None: - self._onload_actor_and_optimizer() loss_cfg: BaseRLLossConfig = self.config.loss_cfg num_batches = len(data_batches) iters_per_step = math.ceil(num_batches / self._optimizer_steps) @@ -659,7 +636,6 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo # Init loss_ctx: shifted_labels, advantages, rollout_logprobs seq_ctx_list: list[SequenceContext] = [] loss_ctx_list: list[BaseRLLossContext] = [] - teacher_indices_list: list[torch.Tensor] = [] mtp_loss_ctx_list: list[list[MTPLossContext]] = [] prepare_inputs_begin = time.perf_counter() for data in data_batches: @@ -693,18 +669,6 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo advantages = data["advantages"].to(DEVICE) rollout_logprobs = data.get("rollout_logprobs", None) rollout_logprobs = rollout_logprobs.to(DEVICE) if rollout_logprobs is not None else None - if teacher_indices is not None: - # Keep the full indices on the host until each SP rank selects - # its local slice. - if self.sp_mesh.size() > 1: - teacher_indices = sp_split( - teacher_indices, - sp_mesh=self.sp_mesh, - split_dim=1, - padding_value=-1, - ) - teacher_indices = teacher_indices.to(DEVICE) - teacher_indices_list.append(teacher_indices) loss_ctx = loss_cfg.build( data={ "shifted_labels": shifted_labels, @@ -712,6 +676,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo "rollout_logprobs": rollout_logprobs, "teacher_logprobs": data.get("teacher_logprobs", None), "target_token_ids": data.get("target_token_ids", None), + "teacher_indices": teacher_indices, }, sp_mesh=self.sp_mesh, ) @@ -743,32 +708,19 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo f"{time.perf_counter() - prepare_inputs_begin:.4f}s" ) - del data_batches - # When sp_mesh.size() > 1, get the sp_split shifted_labels and rollout_logprobs shifted_labels_list = [loss_ctx.loss_kwargs.shifted_labels for loss_ctx in loss_ctx_list] rollout_logprobs_list = [loss_ctx.loss_kwargs.rollout_logprobs for loss_ctx in loss_ctx_list] - worker_log_item: WorkerLogItem = {"train_entropy": 0.0, "train_metrics": [], "sft_train_metrics": {}} - if self._train_teacher_manager is not None: - # Training-side Teachers share the training workers' devices. Keep - # Actor and optimizer on CPU while each frozen Teacher produces its - # targets, then restore Actor state for old-logprob and train forward. - teacher_timings = self._compute_train_teacher_outputs( - seq_ctx_list, - teacher_indices_list, - loss_ctx_list, - ) - worker_log_item["teacher_timings"] = teacher_timings.to_dict() - worker_log_item["teacher_compute_time"] = teacher_timings.compute - worker_log_item["teacher_onload_time"] = teacher_timings.onload - worker_log_item["teacher_offload_time"] = teacher_timings.offload + del data_batches # compute old logprobs old_logprobs_list = self.compute_actor_logprobs(seq_ctx_list, shifted_labels_list) for old_logprobs, loss_ctx in zip(old_logprobs_list, loss_ctx_list): loss_ctx.loss_kwargs.old_logprobs = old_logprobs + worker_log_item: WorkerLogItem = {"train_entropy": 0.0, "train_metrics": [], "sft_train_metrics": {}} + logger_msg = f"Rollout {rollout_idx}: " # compute entropy @@ -857,6 +809,10 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo # run, which bounds their GPU residency by the optimizer group size. self._maybe_offload_logprob(loss_ctx_list) del old_logprobs_list + + teacher_timings = self._maybe_compute_train_teacher_outputs(seq_ctx_list, loss_ctx_list) + if teacher_timings is not None: + worker_log_item.update({"teacher_timings": teacher_timings.to_dict()}) # compute batched loss context batched_loss_ctx_list: list[BaseRLLossContext] = [] @@ -961,7 +917,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo for k, v in extra_info_dict.items() if isinstance(v, (torch.Tensor, int, float)) } - extra_info_dict = finalize_train_metrics(extra_info_dict, DEVICE, loss_cfg) + extra_info_dict = loss_cfg.finalize_metrics(extra_info_dict, DEVICE) train_step_info.pop("total_loss") # type: ignore[misc] max_memory = DEVICE_MODULE.max_memory_allocated() / (1024**3) # type: ignore[attr-defined] reserved_memory = DEVICE_MODULE.max_memory_reserved() / (1024**3) # type: ignore[attr-defined] diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 1f00b7efcb..af510d5a60 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -1680,28 +1680,7 @@ def _log_step( if teacher_names: total = max(worker_totals) else: - legacy_values: list[float] = [] - if phase == "compute": - legacy_values = [ - float(legacy_time) - for log_item in worker_log_items - if (legacy_time := log_item.get("teacher_compute_time")) is not None - ] - elif phase == "onload": - legacy_values = [ - float(legacy_time) - for log_item in worker_log_items - if (legacy_time := log_item.get("teacher_onload_time")) is not None - ] - else: - legacy_values = [ - float(legacy_time) - for log_item in worker_log_items - if (legacy_time := log_item.get("teacher_offload_time")) is not None - ] - if not legacy_values: - continue - total = max(legacy_values) + continue all_scalars[f"time/train_teacher/total/{phase}"] = total all_scalars[f"time/train_teacher_{phase}"] = total From 61fdee66677c831a5b60649d32d9c2d19e6b1b7d Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Fri, 11 Sep 2026 06:05:23 +0000 Subject: [PATCH 12/20] Centralize distillation trainer validation --- xtuner/v1/rl/distillation/__init__.py | 2 +- xtuner/v1/rl/distillation/config.py | 53 ++++++++++++++++++- .../distillation/rollout_teacher_manager.py | 23 +------- xtuner/v1/train/rl_trainer.py | 19 +------ 4 files changed, 56 insertions(+), 41 deletions(-) diff --git a/xtuner/v1/rl/distillation/__init__.py b/xtuner/v1/rl/distillation/__init__.py index 0f8bcea864..fe3cdf4a48 100644 --- a/xtuner/v1/rl/distillation/__init__.py +++ b/xtuner/v1/rl/distillation/__init__.py @@ -7,12 +7,12 @@ TeacherTargetConfig, TrainTeacherConfig, TrainTeacherManagerConfig, + validate_opd_sample_params, ) from .rollout_teacher_manager import ( RolloutTeacherClient, RolloutTeacherReplicaRouter, RolloutTeacherScorer, - validate_opd_sample_params, ) from .train_teacher_manager import TrainTeacherManager, TrainTeacherOutputs, TrainTeacherTimings diff --git a/xtuner/v1/rl/distillation/config.py b/xtuner/v1/rl/distillation/config.py index c4b77faaa8..b2f74990f3 100644 --- a/xtuner/v1/rl/distillation/config.py +++ b/xtuner/v1/rl/distillation/config.py @@ -1,21 +1,47 @@ from __future__ import annotations from pathlib import Path -from typing import TYPE_CHECKING, Literal, cast +from typing import TYPE_CHECKING, Any, Literal, cast from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator from xtuner.v1.config.fsdp import FSDPConfig +from xtuner.v1.data_proto.rl_data import SampleParams from xtuner.v1.model.base import TransformerConfig from xtuner.v1.model.compose.base import BaseComposeConfig from xtuner.v1.rl.loss.distillation_loss import DistillationLossConfig if TYPE_CHECKING: + from xtuner.v1.train.rl_trainer import BaseRLTrainerConfig + from .rollout_teacher_manager import RolloutTeacherScorer from .train_teacher_manager import TrainTeacherManager +def validate_opd_sample_params(sample_params: SampleParams) -> None: + """Validate Student sampling required by sampled-token PG-OPD.""" + + identity_sampling_params: dict[str, Any] = { + "temperature": 1.0, + "top_p": 1.0, + "top_k": 0, + "repetition_penalty": 1.0, + "presence_penalty": 0.0, + "frequency_penalty": 0.0, + "min_tokens": 0, + } + non_identity_params = { + name: getattr(sample_params, name) + for name, expected in identity_sampling_params.items() + if getattr(sample_params, name) != expected + } + if non_identity_params: + raise ValueError(f"PG-OPD requires identity student sampling, got {non_identity_params}") + if not sample_params.return_logprob or not sample_params.return_token_ids: + raise ValueError("PG-OPD requires return_logprob=True and return_token_ids=True") + + class RolloutTeacherLaunchConfig(BaseModel): model_config = ConfigDict(extra="forbid") @@ -241,6 +267,31 @@ def validate_student_model(self, student_model_cfg: TransformerConfig | BaseComp f"student vocab_size={student_vocab_size}" ) + def validate_trainer(self, trainer_cfg: BaseRLTrainerConfig) -> None: + """Validate contracts that span the distillation and trainer + configs.""" + + if trainer_cfg.train_worker_cfg.loss_cfg != self.loss_config: + raise ValueError("train_worker_cfg.loss_cfg must be distillation_config.loss_config") + + self.validate_student_model(trainer_cfg.train_worker_cfg.model_cfg) + + if not self.loss_config.uses_sampled_token_targets: + return + + tasks = trainer_cfg.agent_loop_manager_cfg.tasks + task_configs = tasks if isinstance(tasks, list) else [tasks] + for task in task_configs: + sample_params = task.agent_loop_config.sample_params + if sample_params is None: + raise ValueError( + f"Task {task.task_name!r} must configure sample_params for sampled-token distillation" + ) + try: + validate_opd_sample_params(sample_params) + except ValueError as exc: + raise ValueError(f"Invalid sample_params for task {task.task_name!r}: {exc}") from exc + def resolve_teacher_endpoints( self, endpoint_map: dict[str, list[str]], diff --git a/xtuner/v1/rl/distillation/rollout_teacher_manager.py b/xtuner/v1/rl/distillation/rollout_teacher_manager.py index 3d3222899e..4cd5c12a9a 100644 --- a/xtuner/v1/rl/distillation/rollout_teacher_manager.py +++ b/xtuner/v1/rl/distillation/rollout_teacher_manager.py @@ -9,7 +9,7 @@ import httpx -from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status +from xtuner.v1.data_proto.rl_data import RolloutState, Status, get_group_status from xtuner.v1.rl.utils import create_task from xtuner.v1.utils import get_logger @@ -19,27 +19,6 @@ logger = get_logger() -def validate_opd_sample_params(sample_params: SampleParams) -> None: - identity_sampling_params: dict[str, Any] = { - "temperature": 1.0, - "top_p": 1.0, - "top_k": 0, - "repetition_penalty": 1.0, - "presence_penalty": 0.0, - "frequency_penalty": 0.0, - "min_tokens": 0, - } - non_identity_params = { - name: getattr(sample_params, name) - for name, expected in identity_sampling_params.items() - if getattr(sample_params, name) != expected - } - if non_identity_params: - raise ValueError(f"PG-OPD requires identity student sampling, got {non_identity_params}") - if not sample_params.return_logprob or not sample_params.return_token_ids: - raise ValueError("PG-OPD requires return_logprob=True and return_token_ids=True") - - class RolloutTeacherReplicaRouter: """Resolve a physical teacher replica for one scoring request.""" diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index af510d5a60..681abd449e 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -33,7 +33,7 @@ ProduceBatchStatus, ) from xtuner.v1.rl.agent_loop_manager.produce_utils import default_should_continue_fn -from xtuner.v1.rl.distillation import DistillationConfig, validate_opd_sample_params +from xtuner.v1.rl.distillation import DistillationConfig from xtuner.v1.rl.evaluator import EvaluatorConfig from xtuner.v1.rl.health_manager import RLHealthManager, _NoOpRLHealthManager from xtuner.v1.rl.loss import DistillationLossConfig @@ -514,22 +514,7 @@ def _validate_sync_intervals(self): if self.total_epochs is not None and self.total_epochs <= 0: raise ValueError(f"total_epochs must be positive, got {self.total_epochs}.") if self.distillation_config is not None: - if self.train_worker_cfg.loss_cfg != self.distillation_config.loss_config: - raise ValueError("train_worker_cfg.loss_cfg must be distillation_config.loss_config") - self.distillation_config.validate_student_model(self.train_worker_cfg.model_cfg) - if self.distillation_config.loss_config.uses_sampled_token_targets: - tasks = self.agent_loop_manager_cfg.tasks - tasks = tasks if isinstance(tasks, list) else [tasks] - for task in tasks: - sample_params = task.agent_loop_config.sample_params - if sample_params is None: - raise ValueError( - f"Task {task.task_name!r} must configure sample_params for sampled-token distillation" - ) - try: - validate_opd_sample_params(sample_params) - except ValueError as exc: - raise ValueError(f"Invalid sample_params for task {task.task_name!r}: {exc}") from exc + self.distillation_config.validate_trainer(self) _validate_sync_intervals( sync_weights_interval=self.sync_weights_interval, checkpoint_interval=self.checkpoint_interval, From 9f2bc38c8535f9c545cf6841385f4746bd7674de Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Fri, 11 Sep 2026 11:18:09 +0000 Subject: [PATCH 13/20] Extract distillation trainer adapter --- xtuner/v1/rl/distillation/__init__.py | 2 + xtuner/v1/rl/distillation/trainer_adapter.py | 189 +++++++++++++++++++ xtuner/v1/train/rl_trainer.py | 174 +++-------------- 3 files changed, 217 insertions(+), 148 deletions(-) create mode 100644 xtuner/v1/rl/distillation/trainer_adapter.py diff --git a/xtuner/v1/rl/distillation/__init__.py b/xtuner/v1/rl/distillation/__init__.py index fe3cdf4a48..ee4505ea65 100644 --- a/xtuner/v1/rl/distillation/__init__.py +++ b/xtuner/v1/rl/distillation/__init__.py @@ -15,6 +15,7 @@ RolloutTeacherScorer, ) from .train_teacher_manager import TrainTeacherManager, TrainTeacherOutputs, TrainTeacherTimings +from .trainer_adapter import DistillationTrainerAdapter __all__ = [ @@ -32,5 +33,6 @@ "TrainTeacherManager", "TrainTeacherOutputs", "TrainTeacherTimings", + "DistillationTrainerAdapter", "validate_opd_sample_params", ] diff --git a/xtuner/v1/rl/distillation/trainer_adapter.py b/xtuner/v1/rl/distillation/trainer_adapter.py new file mode 100644 index 0000000000..9ca8c38bb2 --- /dev/null +++ b/xtuner/v1/rl/distillation/trainer_adapter.py @@ -0,0 +1,189 @@ +"""Trainer-side adapter for on-policy distillation.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, cast + +import torch + +from xtuner.v1.data_proto.rl_data import RolloutState + +from .config import DistillationConfig, RolloutTeacherScorerConfig, TrainTeacherManagerConfig + + +if TYPE_CHECKING: + from xtuner.v1.rl.trainer.worker import WorkerLogItem + + +class DistillationTrainerAdapter: + """Adapt distillation semantics to the Trainer integration points. + + The adapter is intentionally stateless. It does not own Teacher or model + lifecycle and does not keep cross-batch mutable state. A ``None`` config + is represented by the same object and produces the non-distillation + defaults. + """ + + def __init__(self, config: DistillationConfig | None): + self._config = config + + @property + def rollout_teacher_scorer_config(self) -> RolloutTeacherScorerConfig | None: + if self._config is None: + return None + return self._config.rollout_teacher_scorer_config + + @property + def train_teacher_manager_config(self) -> TrainTeacherManagerConfig | None: + if self._config is None: + return None + return self._config.train_teacher_manager_config + + @property + def task_adv_weight(self) -> float: + if self._config is None: + return 1.0 + return self._config.loss_config.task_adv_weight + + @property + def teacher_index_by_data_source(self) -> Mapping[str, int] | None: + if self._config is None or not self._config.train_teachers: + return None + return self._config.teacher_index_by_data_source + + def rollout_teacher_targets( + self, + state: RolloutState, + *, + shifted_labels: Sequence[int], + target_start: int, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Return aligned rollout Teacher targets, or no fields when + disabled.""" + + if self._config is None or not self._config.rollout_teachers: + return None, None + return self._align_rollout_teacher_targets( + state, + shifted_labels=shifted_labels, + target_start=target_start, + ) + + def reward_scalars( + self, + observations: Sequence[tuple[RolloutState, float]], + ) -> dict[str, float]: + """Aggregate representative session rewards by their configured + Teacher.""" + + if self._config is None: + return {} + + teacher_rewards: dict[str, list[float]] = {} + for representative, reward in observations: + data_source = representative.extra_fields.get("origin_data_source") + if not isinstance(data_source, str): + continue + teacher_name = self._config.data_source_teacher_map.get(data_source) + if teacher_name is not None: + teacher_rewards.setdefault(teacher_name, []).append(float(reward)) + + return { + f"rewards/{teacher_name}/mean": float(torch.tensor(rewards, dtype=torch.float32).mean().item()) + for teacher_name, rewards in teacher_rewards.items() + if rewards + } + + def timing_scalars(self, worker_log_items: Sequence[WorkerLogItem]) -> dict[str, float]: + """Aggregate train Teacher timings across workers on the critical + path.""" + + phases = ("compute", "onload", "offload") + teacher_timings_by_worker = [log_item.get("teacher_timings", {}) for log_item in worker_log_items] + teacher_names = sorted( + {teacher_name for teacher_timings in teacher_timings_by_worker for teacher_name in teacher_timings} + ) + scalars: dict[str, float] = {} + for teacher_name in teacher_names: + for phase in phases: + phase_values = [ + float(teacher_timings[teacher_name][phase]) + for teacher_timings in teacher_timings_by_worker + if teacher_name in teacher_timings + ] + scalars[f"time/train_teacher/{teacher_name}/{phase}"] = max(phase_values) + + # Training waits for the slowest worker, so report per-worker Teacher + # totals and then take the maximum as the step-level critical-path time. + for phase in phases: + worker_totals = [ + sum(float(teacher_timing[phase]) for teacher_timing in teacher_timings.values()) + for teacher_timings in teacher_timings_by_worker + ] + if not teacher_names: + continue + total = max(worker_totals) + scalars[f"time/train_teacher/total/{phase}"] = total + scalars[f"time/train_teacher_{phase}"] = total + return scalars + + def _align_rollout_teacher_targets( + self, + state: RolloutState, + *, + shifted_labels: Sequence[int], + target_start: int, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + assert self._config is not None + loss_config = self._config.loss_config + sequence_length = len(shifted_labels) + context = f"rollout_id={state.rollout_id}, group_id={state.group_id}" + if target_start < 0 or target_start > sequence_length: + raise ValueError( + "Teacher target_start must be within the shifted sequence: " + f"{target_start} vs {sequence_length}; {context}" + ) + if any(label != loss_config.ignore_idx for label in shifted_labels[:target_start]): + raise ValueError(f"Teacher target prefix must contain only ignored labels; {context}") + + expected_target_rows = sequence_length - target_start + raw_teacher_logprobs = state.teacher_logprobs + if raw_teacher_logprobs is None or len(raw_teacher_logprobs) != expected_target_rows: + actual_rows = None if raw_teacher_logprobs is None else len(raw_teacher_logprobs) + raise ValueError( + "Teacher logprobs must align with the target suffix: " + f"expected {expected_target_rows} rows, got {actual_rows}; {context}" + ) + + if loss_config.uses_sampled_token_targets: + sampled_logprobs = cast(list[float], raw_teacher_logprobs) + teacher_logprobs = torch.tensor( + [0.0] * target_start + sampled_logprobs, + dtype=torch.float32, + ).unsqueeze(0) + return teacher_logprobs, None + + top_k = cast(int, loss_config.top_k) + raw_teacher_tokens = state.teacher_tokens + if raw_teacher_tokens is None or len(raw_teacher_tokens) != expected_target_rows: + actual_rows = None if raw_teacher_tokens is None else len(raw_teacher_tokens) + raise ValueError( + "Teacher token ids must align with the target suffix: " + f"expected {expected_target_rows} rows, got {actual_rows}; {context}" + ) + + topk_tokens = cast(list[list[int]], raw_teacher_tokens) + topk_logprobs = cast(list[list[float]], raw_teacher_logprobs) + + prompt_target_tokens = [[0] * top_k for _ in range(target_start)] + prompt_teacher_logprobs = [[0.0] * top_k for _ in range(target_start)] + teacher_logprobs = torch.tensor( + prompt_teacher_logprobs + topk_logprobs, + dtype=torch.float32, + ).unsqueeze(0) + target_token_ids = torch.tensor( + prompt_target_tokens + topk_tokens, + dtype=torch.int64, + ).unsqueeze(0) + return teacher_logprobs, target_token_ids diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 681abd449e..57ee1a7033 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -4,7 +4,6 @@ import random import re import time -from collections.abc import Sequence from dataclasses import asdict, dataclass from pathlib import Path from shutil import rmtree @@ -33,10 +32,9 @@ ProduceBatchStatus, ) from xtuner.v1.rl.agent_loop_manager.produce_utils import default_should_continue_fn -from xtuner.v1.rl.distillation import DistillationConfig +from xtuner.v1.rl.distillation import DistillationConfig, DistillationTrainerAdapter from xtuner.v1.rl.evaluator import EvaluatorConfig from xtuner.v1.rl.health_manager import RLHealthManager, _NoOpRLHealthManager -from xtuner.v1.rl.loss import DistillationLossConfig from xtuner.v1.rl.replay_buffer import ( AsyncReplayBufferConfig, SyncReplayBufferConfig, @@ -78,66 +76,6 @@ def _to_cpu_tensor(value: np.ndarray | None, *, dtype: torch.dtype | None = None return torch.as_tensor(value, dtype=dtype, device="cpu") -def _align_rollout_teacher_targets( - state: RolloutState, - loss_config: DistillationLossConfig, - *, - shifted_labels: Sequence[int], - target_start: int, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """Left-pad rollout Teacher targets to align with one training sequence.""" - - sequence_length = len(shifted_labels) - context = f"rollout_id={state.rollout_id}, group_id={state.group_id}" - if target_start < 0 or target_start > sequence_length: - raise ValueError( - f"Teacher target_start must be within the shifted sequence: {target_start} vs {sequence_length}; {context}" - ) - if any(label != loss_config.ignore_idx for label in shifted_labels[:target_start]): - raise ValueError(f"Teacher target prefix must contain only ignored labels; {context}") - - expected_target_rows = sequence_length - target_start - raw_teacher_logprobs = state.teacher_logprobs - if raw_teacher_logprobs is None or len(raw_teacher_logprobs) != expected_target_rows: - actual_rows = None if raw_teacher_logprobs is None else len(raw_teacher_logprobs) - raise ValueError( - "Teacher logprobs must align with the target suffix: " - f"expected {expected_target_rows} rows, got {actual_rows}; {context}" - ) - - if loss_config.uses_sampled_token_targets: - sampled_logprobs = cast(list[float], raw_teacher_logprobs) - teacher_logprobs = torch.tensor( - [0.0] * target_start + sampled_logprobs, - dtype=torch.float32, - ).unsqueeze(0) - return teacher_logprobs, None - - top_k = cast(int, loss_config.top_k) - raw_teacher_tokens = state.teacher_tokens - if raw_teacher_tokens is None or len(raw_teacher_tokens) != expected_target_rows: - actual_rows = None if raw_teacher_tokens is None else len(raw_teacher_tokens) - raise ValueError( - "Teacher token ids must align with the target suffix: " - f"expected {expected_target_rows} rows, got {actual_rows}; {context}" - ) - - topk_tokens = cast(list[list[int]], raw_teacher_tokens) - topk_logprobs = cast(list[list[float]], raw_teacher_logprobs) - - prompt_target_tokens = [[0] * top_k for _ in range(target_start)] - prompt_teacher_logprobs = [[0.0] * top_k for _ in range(target_start)] - teacher_logprobs = torch.tensor( - prompt_teacher_logprobs + topk_logprobs, - dtype=torch.float32, - ).unsqueeze(0) - target_token_ids = torch.tensor( - prompt_target_tokens + topk_tokens, - dtype=torch.int64, - ).unsqueeze(0) - return teacher_logprobs, target_token_ids - - def _agent_loop_manager_requires_rollout_proxy( cfg: AgentLoopManagerConfig | DisaggAgentLoopManagerConfig | None, ) -> bool: @@ -716,22 +654,14 @@ class BaseRLTrainer: agent_loop_manager: AgentLoopManager | DisaggAgentLoopManager eval_agent_loop_manager: AgentLoopManager _debug_train_files: dict[int, Path] + _distillation: DistillationTrainerAdapter def _init_common(self, cfg: BaseRLTrainerConfig, *, meta_path: str, logger_tag: str) -> None: if cfg.distillation_config is not None and cfg.distillation_config.rollout_teachers: endpoint_map = json.loads(os.environ.get("XTUNER_OPD_TEACHER_ENDPOINTS_JSON", "{}")) cfg.distillation_config = cfg.distillation_config.resolve_teacher_endpoints(endpoint_map) - self._distillation_config = cfg.distillation_config - self._rollout_teacher_scorer_config = ( - self._distillation_config.rollout_teacher_scorer_config if self._distillation_config is not None else None - ) - self._train_teacher_manager_config = ( - self._distillation_config.train_teacher_manager_config if self._distillation_config is not None else None - ) - self._distillation_loss_cfg = ( - self._distillation_config.loss_config if self._distillation_config is not None else None - ) + self._distillation = DistillationTrainerAdapter(cfg.distillation_config) check_fa3() self._init_work_dir_and_meta(cfg, meta_path) @@ -836,7 +766,7 @@ def _init_train_worker_config(self, cfg: BaseRLTrainerConfig, log_dir: Path) -> cfg.train_worker_cfg.free_rollout_routed_experts_in_worker = False cfg.train_worker_cfg.load_from = cfg.load_from cfg.train_worker_cfg.log_dir = log_dir - cfg.train_worker_cfg.train_teacher_manager_config = self._train_teacher_manager_config + cfg.train_worker_cfg.train_teacher_manager_config = self._distillation.train_teacher_manager_config self._train_worker_cfg = cfg.train_worker_cfg def _init_rollout_config(self, cfg: BaseRLTrainerConfig, log_dir: Path) -> None: @@ -870,7 +800,7 @@ def _build_agent_loop_components(self, cfg: BaseRLTrainerConfig, replay_buffer) replay_buffer=replay_buffer, logger=self.logger, sync_weights_interval=cfg.sync_weights_interval, - rollout_teacher_scorer_config=self._rollout_teacher_scorer_config, + rollout_teacher_scorer_config=self._distillation.rollout_teacher_scorer_config, ) self.agent_loop_manager = cast(AgentLoopManager | DisaggAgentLoopManager, agent_loop_manager) @@ -1211,7 +1141,7 @@ def _prepare_train_data( # trainable segments that share one reward; counting that reward once per session keeps # rewards/* from being weighted by segment count. cluster_rewards_list: list[float] = [] - teacher_rewards: dict[str, list[float]] = {} + teacher_reward_observations: list[tuple[RolloutState, float]] = [] advantages_list = [] prompt_len_list = [] response_len_list = [] @@ -1220,11 +1150,7 @@ def _prepare_train_data( training_samples = 0 data_batches = [] - teacher_index_by_data_source = ( - self._train_teacher_manager_config.teacher_index_by_data_source - if self._train_teacher_manager_config is not None - else None - ) + teacher_index_by_data_source = self._distillation.teacher_index_by_data_source for j, group in enumerate(data_groups): if not is_valid_for_training(group, self.logger): @@ -1232,9 +1158,7 @@ def _prepare_train_data( continue training_samples += len(group) - task_adv_weight = ( - self._distillation_loss_cfg.task_adv_weight if self._distillation_loss_cfg is not None else 1.0 - ) + task_adv_weight = self._distillation.task_adv_weight prompt_ids = None if any(data.input_ids is None for data in group): is_vlm_model = "train_prompt_ids" in group[0].extra_fields @@ -1274,14 +1198,7 @@ def _prepare_train_data( sample_cluster_indices.append(cluster_index) cluster_rewards_list.extend(cluster_rewards) - if self._distillation_config is not None: - for reward, representative in zip(cluster_rewards, cluster_representatives): - data_source = representative.extra_fields.get("origin_data_source") - if not isinstance(data_source, str): - continue - teacher_name = self._distillation_config.data_source_teacher_map.get(data_source) - if teacher_name is not None: - teacher_rewards.setdefault(teacher_name, []).append(reward) + teacher_reward_observations.extend(zip(cluster_representatives, cluster_rewards)) if task_adv_weight == 0: sample_advantages = [0.0] * len(group) @@ -1323,21 +1240,15 @@ def _prepare_train_data( shifted_labels = labels[1:] teacher_logprobs = None target_token_ids = None - if ( - self._distillation_config is not None - and self._distillation_loss_cfg is not None - and self._distillation_config.rollout_teachers - ): - teacher_response_start = next( - (index for index, label in enumerate(shifted_labels) if label != -100), - len(shifted_labels), - ) - teacher_logprobs, target_token_ids = _align_rollout_teacher_targets( - group[i], - self._distillation_loss_cfg, - shifted_labels=shifted_labels, - target_start=teacher_response_start, - ) + teacher_response_start = next( + (index for index, label in enumerate(shifted_labels) if label != -100), + len(shifted_labels), + ) + teacher_logprobs, target_token_ids = self._distillation.rollout_teacher_targets( + group[i], + shifted_labels=shifted_labels, + target_start=teacher_response_start, + ) prompt_len = sum(label == -100 for label in shifted_labels) response_len = len(shifted_labels) - prompt_len prompt_len_list.append(prompt_len) @@ -1461,17 +1372,12 @@ def _prepare_train_data( "advantage": actual_advantages, "rollout_logprobs": rollout_logprobs, } - if ( - self._distillation_config is not None - and self._distillation_loss_cfg is not None - and self._distillation_config.rollout_teachers - ): - teacher_logprobs, target_token_ids = _align_rollout_teacher_targets( - group[i], - self._distillation_loss_cfg, - shifted_labels=shifted_labels, - target_start=len(prompt_ids) - 1, - ) + teacher_logprobs, target_token_ids = self._distillation.rollout_teacher_targets( + group[i], + shifted_labels=shifted_labels, + target_start=len(prompt_ids) - 1, + ) + if teacher_logprobs is not None: data_dict["teacher_logprobs"] = teacher_logprobs if target_token_ids is not None: data_dict["target_token_ids"] = target_token_ids @@ -1514,9 +1420,7 @@ def _prepare_train_data( "prompt_len/min": prompt_len_t.min().item(), "prompt_len/max": prompt_len_t.max().item(), } - for teacher_name, rewards in teacher_rewards.items(): - if rewards: - info_dict[f"rewards/{teacher_name}/mean"] = torch.tensor(rewards, dtype=torch.float32).mean().item() + info_dict.update(self._distillation.reward_scalars(teacher_reward_observations)) if tool_turns_list: tool_turns_t = torch.tensor(tool_turns_list, dtype=torch.float32) info_dict["tool_turns/mean"] = tool_turns_t.mean().item() @@ -1641,33 +1545,7 @@ def _log_step( all_scalars.update({"entropy/rollout": rank0_rollout_entropy}) all_scalars.update({"entropy/train": rank0_log_item["train_entropy"]}) - teacher_timing_phases = ("compute", "onload", "offload") - teacher_timings_by_worker = [log_item.get("teacher_timings", {}) for log_item in worker_log_items] - teacher_names = sorted( - {teacher_name for teacher_timings in teacher_timings_by_worker for teacher_name in teacher_timings} - ) - for teacher_name in teacher_names: - for phase in teacher_timing_phases: - phase_values = [ - float(teacher_timings[teacher_name][phase]) - for teacher_timings in teacher_timings_by_worker - if teacher_name in teacher_timings - ] - all_scalars[f"time/train_teacher/{teacher_name}/{phase}"] = max(phase_values) - - # Training waits for the slowest worker, so report per-worker Teacher totals and then - # take the maximum as the step-level critical-path time. - for phase in teacher_timing_phases: - worker_totals = [ - sum(float(teacher_timing[phase]) for teacher_timing in teacher_timings.values()) - for teacher_timings in teacher_timings_by_worker - ] - if teacher_names: - total = max(worker_totals) - else: - continue - all_scalars[f"time/train_teacher/total/{phase}"] = total - all_scalars[f"time/train_teacher_{phase}"] = total + all_scalars.update(self._distillation.timing_scalars(worker_log_items)) for worker_idx, log_item in enumerate(worker_log_items): if not self._display_all_workers_log and worker_idx > 0: From 8daf57feb523b70f02e5da266e8198504ce554d3 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Sun, 13 Sep 2026 14:21:00 +0000 Subject: [PATCH 14/20] Unify rollout teacher targets --- tests/rl/test_on_policy_distillation.py | 265 ++++++++++++++++-- tests/rl/test_prepare_train_data.py | 30 +- tests/rl/test_single_turn_agent_loop.py | 75 ++--- xtuner/v1/data_proto/rl_data.py | 40 ++- .../distillation/rollout_teacher_manager.py | 9 +- xtuner/v1/rl/distillation/trainer_adapter.py | 113 ++++---- 6 files changed, 390 insertions(+), 142 deletions(-) diff --git a/tests/rl/test_on_policy_distillation.py b/tests/rl/test_on_policy_distillation.py index fa9d4cb3f9..2af33ccd3e 100644 --- a/tests/rl/test_on_policy_distillation.py +++ b/tests/rl/test_on_policy_distillation.py @@ -14,12 +14,15 @@ from xtuner.v1.data_proto.sequence_context import SequenceContext from xtuner.v1.rl.distillation import ( DistillationConfig, + DistillationTrainerAdapter, RolloutTeacherClient, RolloutTeacherConfig, + TeacherTargetConfig, TrainTeacherManager, + TrainTeacherOutputs, TrainTeacherTimings, ) -from xtuner.v1.rl.loss import DistillationLossConfig +from xtuner.v1.rl.loss import DistillationLossConfig, DistillationLossKwargs from xtuner.v1.rl.trainer.controller import TrainingController from xtuner.v1.train.rl_trainer import BaseRLTrainer, BaseRLTrainerConfig @@ -95,9 +98,125 @@ def test_sampled_token_sampling_is_validated_by_trainer_config(self) -> None: with self.assertRaisesRegex(ValueError, "Invalid sample_params for task 'math'.*temperature"): config._validate_sync_intervals() + def test_distillation_trainer_validation_handles_multiple_tasks(self) -> None: + loss_config = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + distillation_config = DistillationConfig( + loss_config=loss_config, + teachers=[RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"])], + data_source_teacher_map={"math": "teacher", "code": "teacher"}, + ) + trainer_cfg = SimpleNamespace( + train_worker_cfg=SimpleNamespace(loss_cfg=loss_config, model_cfg=MagicMock()), + agent_loop_manager_cfg=SimpleNamespace( + tasks=[ + SimpleNamespace( + task_name="math", + agent_loop_config=SimpleNamespace(sample_params=SampleParams()), + ), + SimpleNamespace( + task_name="code", + agent_loop_config=SimpleNamespace(sample_params=SampleParams()), + ), + ] + ), + ) + + distillation_config.validate_trainer(trainer_cfg) + + def test_distillation_trainer_validation_rejects_missing_sample_params(self) -> None: + loss_config = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + distillation_config = DistillationConfig( + loss_config=loss_config, + teachers=[RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"])], + data_source_teacher_map={"math": "teacher"}, + ) + trainer_cfg = SimpleNamespace( + train_worker_cfg=SimpleNamespace(loss_cfg=loss_config, model_cfg=MagicMock()), + agent_loop_manager_cfg=SimpleNamespace( + tasks=SimpleNamespace( + task_name="math", + agent_loop_config=SimpleNamespace(sample_params=None), + ) + ), + ) + + with self.assertRaisesRegex(ValueError, "Task 'math' must configure sample_params"): + distillation_config.validate_trainer(trainer_cfg) + + def test_distillation_trainer_validation_skips_sampling_constraints_for_topk(self) -> None: + loss_config = DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="forward_kl_topk", + use_policy_gradient=False, + top_k=2, + ) + distillation_config = DistillationConfig( + loss_config=loss_config, + teachers=[RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"])], + data_source_teacher_map={"math": "teacher"}, + ) + trainer_cfg = SimpleNamespace( + train_worker_cfg=SimpleNamespace(loss_cfg=loss_config, model_cfg=MagicMock()), + agent_loop_manager_cfg=SimpleNamespace( + tasks=SimpleNamespace( + task_name="math", + agent_loop_config=SimpleNamespace(sample_params=SampleParams(temperature=0.7)), + ) + ), + ) + + distillation_config.validate_trainer(trainer_cfg) + + def test_distillation_trainer_validation_rejects_mismatched_loss_config(self) -> None: + distillation_loss_config = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + worker_loss_config = DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + task_adv_weight=1.0, + ) + distillation_config = DistillationConfig( + loss_config=distillation_loss_config, + teachers=[RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"])], + data_source_teacher_map={"math": "teacher"}, + ) + trainer_cfg = SimpleNamespace( + train_worker_cfg=SimpleNamespace(loss_cfg=worker_loss_config, model_cfg=MagicMock()), + agent_loop_manager_cfg=SimpleNamespace(tasks=SimpleNamespace()), + ) + + with self.assertRaisesRegex(ValueError, "train_worker_cfg.loss_cfg must be distillation_config.loss_config"): + distillation_config.validate_trainer(trainer_cfg) + + def test_rollout_teacher_scorer_config_is_projected_from_distillation_config(self) -> None: + loss_config = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + distillation_config = DistillationConfig( + loss_config=loss_config, + teachers=[RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"])], + data_source_teacher_map={"math": "teacher"}, + ) + + scorer_config = distillation_config.rollout_teacher_scorer_config + + self.assertIsNotNone(scorer_config) + assert scorer_config is not None + self.assertEqual(scorer_config.target_config, distillation_config.teacher_target_config) + self.assertEqual([teacher.name for teacher in scorer_config.teachers], ["teacher"]) + self.assertEqual(scorer_config.data_source_teacher_map, {"math": "teacher"}) + self.assertTrue(scorer_config.target_config.uses_sampled_token_targets) + self.assertIsNone(scorer_config.target_config.top_k) + self.assertIsNone(distillation_config.train_teacher_manager_config) + + def test_distillation_trainer_adapter_uses_null_defaults_without_config(self) -> None: + adapter = DistillationTrainerAdapter(None) + + self.assertIsNone(adapter.rollout_teacher_scorer_config) + self.assertIsNone(adapter.train_teacher_manager_config) + self.assertEqual(adapter.task_adv_weight, 1.0) + self.assertIsNone(adapter.teacher_index_by_data_source) + self.assertEqual(adapter.timing_scalars([]), {}) + def test_train_teacher_config_is_not_forwarded_to_agent_loop_manager(self) -> None: trainer = BaseRLTrainer.__new__(BaseRLTrainer) - trainer._rollout_teacher_config = None + trainer._distillation = DistillationTrainerAdapter(None) trainer._enable_evaluate = False trainer._resolve_total_train_steps = MagicMock() trainer.rollout_controller = MagicMock() @@ -115,7 +234,7 @@ def test_train_teacher_config_is_not_forwarded_to_agent_loop_manager(self) -> No trainer._build_agent_loop_components(config, replay_buffer=MagicMock()) self.assertIs(trainer.agent_loop_manager, built_manager) - self.assertIsNone(manager_config.build.call_args.kwargs["distillation_config"]) + self.assertIsNone(manager_config.build.call_args.kwargs["rollout_teacher_scorer_config"]) class TestTrainTeacherTimings(unittest.TestCase): @@ -147,6 +266,111 @@ def test_lifecycle_timings_are_recorded_for_every_teacher(self) -> None: ) self.assertEqual((timings.compute, timings.onload, timings.offload), (7.0, 5.0, 9.0)) + def test_compute_teacher_outputs_writes_teacher_outputs_to_loss_contexts(self) -> None: + manager = TrainTeacherManager.__new__(TrainTeacherManager) + teacher_indices = torch.tensor([[0, 1]], dtype=torch.long) + shifted_labels = torch.tensor([[-100, 10]], dtype=torch.long) + teacher_logprobs = [torch.tensor([[-0.1, -0.2]])] + target_token_ids = [torch.tensor([[[1, 2], [3, 4]]], dtype=torch.long)] + timings = TrainTeacherTimings() + seq_ctx = MagicMock() + loss_ctx = SimpleNamespace( + loss_kwargs=SimpleNamespace( + teacher_indices=teacher_indices, + shifted_labels=shifted_labels, + ) + ) + + with ( + patch.object( + manager, + "compute_logprobs", + return_value=TrainTeacherOutputs( + teacher_logprobs=teacher_logprobs, + target_token_ids=target_token_ids, + timings=timings, + ), + ) as compute_logprobs, + patch.object(manager, "offload_all_to_cpu") as offload_all_to_cpu, + ): + result = manager.compute_teacher_outputs([seq_ctx], [loss_ctx]) + + self.assertIs(result, timings) + compute_logprobs.assert_called_once_with( + seq_ctx_list=[seq_ctx], + shifted_labels_list=[shifted_labels], + teacher_indices_list=[teacher_indices], + ) + torch.testing.assert_close(loss_ctx.loss_kwargs.teacher_logprobs, teacher_logprobs[0]) + torch.testing.assert_close(loss_ctx.loss_kwargs.target_token_ids, target_token_ids[0]) + offload_all_to_cpu.assert_called_once_with() + + def test_distillation_loss_kwargs_carry_teacher_indices(self) -> None: + loss_cfg = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + teacher_indices = torch.tensor([[0, 1]], dtype=torch.long) + loss_ctx = loss_cfg.build( + data={ + "shifted_labels": torch.tensor([[-100, 10]], dtype=torch.long), + "advantages": torch.ones(1, 2), + "teacher_indices": teacher_indices, + } + ) + + assert loss_ctx is not None + torch.testing.assert_close(loss_ctx.loss_kwargs.teacher_indices, teacher_indices) + + def test_distillation_loss_kwargs_split_teacher_indices_with_ignore_padding(self) -> None: + loss_kwargs = DistillationLossKwargs( + shifted_labels=torch.ones(1, 2, dtype=torch.long), + advantages=torch.ones(1, 2), + teacher_logprobs=torch.ones(1, 2), + target_token_ids=torch.ones(1, 2, dtype=torch.long), + teacher_indices=torch.tensor([[0, 1]], dtype=torch.long), + ) + sp_mesh = SimpleNamespace(size=lambda: 2) + + with ( + patch("xtuner.v1.loss.ce_loss.sp_split", side_effect=lambda tensor, **_: tensor), + patch("xtuner.v1.rl.loss.base_loss.sp_split", side_effect=lambda tensor, **_: tensor), + patch( + "xtuner.v1.rl.loss.distillation_loss.sp_split", + side_effect=lambda tensor, **_: tensor, + ) as split, + ): + loss_kwargs.sp_split(sp_mesh) + + padding_values = [call.kwargs["padding_value"] for call in split.call_args_list] + self.assertEqual(padding_values, [0.0, 0, -1]) + + def test_finalize_train_metrics_handles_policy_and_distillation_keys(self) -> None: + loss_cfg = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + metrics = loss_cfg.finalize_metrics( + { + "reduced_train_policy_ratio_abs_dev_sum": 2.0, + "reduced_train_policy_kl1_sum": 4.0, + "reduced_train_policy_kl3_sum": 6.0, + "reduced_train_policy_valid_count": 2.0, + "reduced_train_policy_ratio_max": 1.5, + "reduced_train_policy_ratio_min": 0.5, + "reduced_distillation_kl_sum": 8.0, + "reduced_distillation_abs_loss_sum": 10.0, + "reduced_distillation_valid_count": 2.0, + "reduced_opd_reverse_kl_sum": 8.0, + "reduced_opd_abs_logprob_loss_sum": 10.0, + }, + "cpu", + ) + + self.assertEqual(metrics["reduced_train_policy_ratio_abs_dev_mean"], 1.0) + self.assertEqual(metrics["reduced_train_policy_kl1"], 2.0) + self.assertEqual(metrics["reduced_train_policy_kl3"], 3.0) + self.assertEqual(metrics["reduced_train_policy_ratio_max"], 1.5) + self.assertEqual(metrics["reduced_train_policy_ratio_min"], 0.5) + self.assertEqual(metrics["reduced_distillation_kl"], 4.0) + self.assertEqual(metrics["reduced_distillation_abs_loss"], 5.0) + self.assertEqual(metrics["opd_reverse_kl"], 4.0) + self.assertEqual(metrics["opd_abs_logprob_loss"], 5.0) + class TestRolloutTeacherClient(unittest.IsolatedAsyncioTestCase): @staticmethod @@ -158,10 +382,8 @@ def _response(payload: dict) -> httpx.Response: ) def _build_topk_client(self, *, max_retry_per_sample: int = 0) -> RolloutTeacherClient: - loss_config = DistillationLossConfig( - policy_loss_cfg={"loss_type": "vanilla"}, - loss_mode="forward_kl_topk", - use_policy_gradient=False, + target_config = TeacherTargetConfig( + mode="topk", top_k=2, ) with patch.dict( @@ -174,7 +396,7 @@ def _build_topk_client(self, *, max_retry_per_sample: int = 0) -> RolloutTeacher endpoints=["http://teacher"], max_retry_per_sample=max_retry_per_sample, ), - loss_config, + target_config, ) self.addAsyncCleanup(client._client.aclose) return client @@ -206,14 +428,14 @@ async def test_compute_sampled_token_logprobs_uses_current_interface(self) -> No } }, ) - loss_config = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + target_config = TeacherTargetConfig(mode="sampled_token") with patch.dict( "os.environ", {"XTUNER_USE_LMDEPLOY": "1", "XTUNER_USE_SGLANG": "0", "XTUNER_USE_VLLM": "0"}, ): client = RolloutTeacherClient( RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"]), - loss_config, + target_config, ) self.addAsyncCleanup(client._client.aclose) client._client.post = AsyncMock(return_value=response) @@ -229,8 +451,9 @@ async def test_compute_sampled_token_logprobs_uses_current_interface(self) -> No result = await client.compute_logprobs(state) self.assertEqual(result.status, Status.COMPLETED) - self.assertEqual(result.teacher_tokens, [13, 14]) - self.assertEqual(result.teacher_logprobs, [-0.3, -0.4]) + self.assertEqual(result.teacher_targets.kind, "sampled") + self.assertEqual(result.teacher_targets.tokens, [13, 14]) + self.assertEqual(result.teacher_targets.logprobs, [-0.3, -0.4]) self.assertIn("teacher_score_time_s", result.extra_fields) async def test_malformed_sampled_response_becomes_failed_state(self) -> None: @@ -246,7 +469,7 @@ async def test_malformed_sampled_response_becomes_failed_state(self) -> None: } } ) - loss_config = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + target_config = TeacherTargetConfig(mode="sampled_token") with patch.dict( "os.environ", {"XTUNER_USE_LMDEPLOY": "1", "XTUNER_USE_SGLANG": "0", "XTUNER_USE_VLLM": "0"}, @@ -257,7 +480,7 @@ async def test_malformed_sampled_response_becomes_failed_state(self) -> None: endpoints=["http://teacher"], max_retry_per_sample=0, ), - loss_config, + target_config, ) self.addAsyncCleanup(client._client.aclose) client._client.post = AsyncMock(return_value=response) @@ -281,7 +504,7 @@ async def test_sampled_response_rejects_non_numeric_logprob(self) -> None: } } ) - loss_config = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + target_config = TeacherTargetConfig(mode="sampled_token") with patch.dict( "os.environ", {"XTUNER_USE_LMDEPLOY": "1", "XTUNER_USE_SGLANG": "0", "XTUNER_USE_VLLM": "0"}, @@ -292,7 +515,7 @@ async def test_sampled_response_rejects_non_numeric_logprob(self) -> None: endpoints=["http://teacher"], max_retry_per_sample=0, ), - loss_config, + target_config, ) self.addAsyncCleanup(client._client.aclose) client._client.post = AsyncMock(return_value=response) @@ -331,8 +554,7 @@ async def test_malformed_topk_responses_become_failed_states(self) -> None: result = await client.compute_logprobs(self._state()) self.assertEqual(result.status, Status.FAILED) - self.assertIsNone(result.teacher_tokens) - self.assertIsNone(result.teacher_logprobs) + self.assertIsNone(result.teacher_targets) self.assertIn("last_error=", result.error_msg or "") client._client.post.assert_awaited_once() @@ -372,12 +594,13 @@ async def test_invalid_topk_response_is_retried_before_success(self) -> None: client = self._build_topk_client(max_retry_per_sample=1) client._client.post = AsyncMock(side_effect=[invalid_response, valid_response]) - with patch("xtuner.v1.rl.distillation.rollout_teacher_client.asyncio.sleep", new=AsyncMock()): + with patch("xtuner.v1.rl.distillation.rollout_teacher_manager.asyncio.sleep", new=AsyncMock()): result = await client.compute_logprobs(self._state()) self.assertEqual(result.status, Status.COMPLETED) - self.assertEqual(result.teacher_tokens, [[5, 6], [7, 8]]) - self.assertEqual(result.teacher_logprobs, [[-0.5, -0.6], [-0.7, -0.8]]) + self.assertEqual(result.teacher_targets.kind, "topk") + self.assertEqual(result.teacher_targets.tokens, [[5, 6], [7, 8]]) + self.assertEqual(result.teacher_targets.logprobs, [[-0.5, -0.6], [-0.7, -0.8]]) self.assertEqual(client._client.post.await_count, 2) request_payload = client._client.post.await_args.kwargs["json"] self.assertEqual(request_payload["input_ids"], [10, 11, 12, 13, 14]) diff --git a/tests/rl/test_prepare_train_data.py b/tests/rl/test_prepare_train_data.py index b6cedb7965..c382f3f5d7 100644 --- a/tests/rl/test_prepare_train_data.py +++ b/tests/rl/test_prepare_train_data.py @@ -15,8 +15,8 @@ import numpy as np import torch -from xtuner.v1.data_proto.rl_data import RolloutState, Status, reset_rollout_response -from xtuner.v1.rl.distillation import DistillationConfig, RolloutTeacherConfig +from xtuner.v1.data_proto.rl_data import RolloutState, Status, TeacherTargets, reset_rollout_response +from xtuner.v1.rl.distillation import DistillationConfig, DistillationTrainerAdapter, RolloutTeacherConfig from xtuner.v1.rl.loss import DistillationLossConfig from xtuner.v1.train.rl_trainer import BaseRLTrainer @@ -35,9 +35,7 @@ class TestPrepareTrainData(unittest.TestCase): def _build_trainer(self, advantages: list[float]): trainer = BaseRLTrainer.__new__(BaseRLTrainer) trainer._advantage_estimator = _FakeAdvantageEstimator(advantages) - trainer._distillation_config = None - trainer._distillation_loss_cfg = None - trainer._train_teacher_config = None + trainer._distillation = DistillationTrainerAdapter(None) trainer.tokenizer = MagicMock(return_value={"input_ids": torch.tensor([[999]])}) trainer.logger = MagicMock() return trainer @@ -81,8 +79,15 @@ def _state( extra_fields=extra_fields or {}, input_ids=input_ids, labels=labels, - teacher_tokens=teacher_tokens, - teacher_logprobs=teacher_logprobs, + teacher_targets=( + TeacherTargets( + kind="topk" if teacher_tokens and isinstance(teacher_tokens[0], list) else "sampled", + tokens=teacher_tokens, + logprobs=teacher_logprobs, + ) + if teacher_tokens is not None and teacher_logprobs is not None + else None + ), ) def _prepare(self, trainer, data_groups, pack_max_length=128): @@ -91,12 +96,13 @@ def _prepare(self, trainer, data_groups, pack_max_length=128): @staticmethod def _enable_rollout_distillation(trainer, loss_config: DistillationLossConfig) -> None: - trainer._distillation_config = DistillationConfig( - loss_config=loss_config, - teachers=[RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"])], - data_source_teacher_map={"agent_math": "teacher"}, + trainer._distillation = DistillationTrainerAdapter( + DistillationConfig( + loss_config=loss_config, + teachers=[RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"])], + data_source_teacher_map={"agent_math": "teacher"}, + ) ) - trainer._distillation_loss_cfg = loss_config def test_text_path_builds_shifted_training_tensors(self): # 文本主路径固定 token 布局：input_ids 去掉 response 最后一个 token，label/logprob 对齐预测位置。 diff --git a/tests/rl/test_single_turn_agent_loop.py b/tests/rl/test_single_turn_agent_loop.py index 21bf28a261..d37a950e26 100644 --- a/tests/rl/test_single_turn_agent_loop.py +++ b/tests/rl/test_single_turn_agent_loop.py @@ -5,7 +5,7 @@ - batch judge 只在整组样本全部 COMPLETED 时触发。 - batch judge 返回结果必须保持输入顺序。 - rollout group validity check 在 batch judge 之后执行。 -- 没有 validity check 时，每条样本先算 Teacher、再进入 Judger。 +- 没有 validity check 时，每条样本先算 Teacher。 - 有 validity check 时，只有通过过滤的 rollout group 才进入 Teacher scoring。 - 被过滤的 rollout group 不进入 Teacher scoring。 - 组内存在 ABORTED / FAILED 等非 COMPLETED 样本时跳过 judge。 @@ -16,7 +16,7 @@ import unittest from unittest.mock import AsyncMock, MagicMock, patch -from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status +from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, TeacherTargets, get_group_status from xtuner.v1.rl.agent_loop import AgentLoopActor from xtuner.v1.rl.agent_loop.single_turn_agent_loop import SingleTurnAgentLoop, SingleTurnAgentLoopConfig @@ -54,22 +54,36 @@ async def batch_judge(self, rollout_states): return rollout_states -class _PerSampleJudger: - def __init__(self, events: list[str]): - self.events = events - - async def judge(self, rollout_state): - self.events.append("judge") - rollout_state.reward = {"score": float(rollout_state.rollout_id)} - return rollout_state - - class _SlowJudger: async def batch_judge(self, rollout_states): await asyncio.sleep(60) return rollout_states +class _PassthroughTeacherScorer: + async def on_sample_ready(self, state): + return state + + async def on_group_ready(self, group): + return group + + +class _TeacherScorer: + def __init__(self, teacher, *, defers_to_filter: bool): + self.teacher = teacher + self.defers_to_filter = defers_to_filter + + async def on_sample_ready(self, state): + if self.defers_to_filter: + return state + return await self.teacher.compute_logprobs(state) + + async def on_group_ready(self, group): + if not self.defers_to_filter or get_group_status(group) != Status.COMPLETED: + return group + return list(await asyncio.gather(*(self.teacher.compute_logprobs(state) for state in group))) + + class TestSingleTurnAgentLoop(unittest.IsolatedAsyncioTestCase): def _state(self, uid: int) -> RolloutState: return RolloutState( @@ -101,8 +115,7 @@ def _build_loop( loop.judger = judger loop.enable_batch_judge = enable_batch_judge loop.is_valid_sample_fn = is_valid_sample_fn - loop.teacher_clients = {} - loop.data_source_teacher_map = {} + loop._teacher_scorer = _PassthroughTeacherScorer() loop._judger_pause_event = asyncio.Event() loop.logger = MagicMock() return loop @@ -113,12 +126,18 @@ def _bind_teacher(loop, events: list[str]): async def compute_logprobs(state): events.append("teacher") - state.teacher_logprobs = [-0.5] * len(state.response_ids or []) + state.teacher_targets = TeacherTargets( + kind="sampled", + tokens=list(state.response_ids or []), + logprobs=[-0.5] * len(state.response_ids or []), + ) return state teacher.compute_logprobs = AsyncMock(side_effect=compute_logprobs) - loop.teacher_clients = {"teacher": teacher} - loop.data_source_teacher_map = {"math": "teacher"} + loop._teacher_scorer = _TeacherScorer( + teacher, + defers_to_filter=loop.is_valid_sample_fn is not None, + ) return teacher def test_config_binds_validity_check_to_local_agent_loop_once(self): @@ -188,8 +207,7 @@ async def test_generate_skips_teacher_scoring_for_filtered_group(self): {1: Status.COMPLETED}, is_valid_sample_fn=lambda _samples: False, ) - loop.teacher_clients = {"teacher": teacher} - loop.data_source_teacher_map = {"math": "teacher"} + loop._teacher_scorer = _TeacherScorer(teacher, defers_to_filter=True) state = self._state(1) state.extra_fields["origin_data_source"] = "math" @@ -213,25 +231,10 @@ async def test_teacher_scoring_precedes_batch_judge_without_filter(self): self.assertEqual(events.count("teacher"), 2) self.assertEqual(events[-1], "judge") - self.assertTrue(all(state.teacher_logprobs == [-0.5] for state in result)) + self.assertTrue(all(state.teacher_targets is not None for state in result)) + self.assertTrue(all(state.teacher_targets.logprobs == [-0.5] for state in result)) self.assertEqual(teacher.compute_logprobs.await_count, 2) - async def test_teacher_scoring_precedes_per_sample_judge_without_filter(self): - events: list[str] = [] - loop = self._build_loop( - {1: Status.COMPLETED}, - judger=_PerSampleJudger(events), - enable_batch_judge=False, - ) - self._bind_teacher(loop, events) - state = self._state(1) - state.extra_fields["origin_data_source"] = "math" - - result = await loop.generate_group([state]) - - self.assertEqual(events, ["teacher", "judge"]) - self.assertEqual(result[0].reward, {"score": 1.0}) - async def test_valid_group_is_scored_after_batch_judge_and_filter(self): events: list[str] = [] diff --git a/xtuner/v1/data_proto/rl_data.py b/xtuner/v1/data_proto/rl_data.py index 6c6a96f65e..b2e5ef43d1 100644 --- a/xtuner/v1/data_proto/rl_data.py +++ b/xtuner/v1/data_proto/rl_data.py @@ -4,8 +4,8 @@ from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast import numpy as np -from pydantic import BaseModel, ConfigDict, Field -from typing_extensions import NotRequired, TypedDict +from pydantic import BaseModel, ConfigDict, Field, model_validator +from typing_extensions import NotRequired, Self, TypedDict # ==================================== # ====== DataFlow 数据流 ============== @@ -81,6 +81,34 @@ class RolloutToolCall(BaseModel): function: RolloutFunctionCall +class TeacherTargets(BaseModel): + """Sampled-token or top-k Teacher outputs for one rollout response.""" + + model_config = ConfigDict(extra="forbid") + + kind: Literal["sampled", "topk"] + tokens: list[int] | list[list[int]] + logprobs: list[float] | list[list[float]] + + @model_validator(mode="after") + def _check_shapes(self) -> Self: + if self.kind == "sampled": + if len(self.tokens) != len(self.logprobs): + raise ValueError("Sampled-token Teacher tokens and logprobs must have equal lengths") + return self + + token_rows = cast(list[list[int]], self.tokens) + logprob_rows = cast(list[list[float]], self.logprobs) + if len(token_rows) != len(logprob_rows): + raise ValueError("Top-k Teacher tokens and logprobs must have equal row counts") + widths = {len(row) for row in token_rows} + if not widths or widths == {0} or len(widths) != 1: + raise ValueError("Top-k Teacher token rows must have one consistent positive width") + if any(len(token_row) != len(logprob_row) for token_row, logprob_row in zip(token_rows, logprob_rows)): + raise ValueError("Top-k Teacher token and logprob rows must have equal widths") + return self + + class RolloutState(BaseModel): model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) @@ -109,10 +137,7 @@ class RolloutState(BaseModel): tool_calls: list[RolloutToolCall] | None = None response_ids: list[int] | None = None logprobs: list[float] | None = None - # Sampled-token teachers store one value per response token. Top-K - # teachers store one K-wide row per response token. - teacher_tokens: list[int] | list[list[int]] | None = None - teacher_logprobs: list[float] | list[list[float]] | None = None + teacher_targets: TeacherTargets | None = None routed_experts: np.ndarray | RayObjectRef | list[RayObjectRef] | None = None finish_reason: str | None = None # response_mask: 记录response_ids中哪个token算loss, 与response_ids长度相同，每轮rollout在 agent_loop.generate 中覆盖写 @@ -248,8 +273,7 @@ def reset_rollout_response(rollout_state: RolloutState) -> RolloutState: rollout_state.response = "" rollout_state.response_ids = [] rollout_state.logprobs = [] - rollout_state.teacher_tokens = None - rollout_state.teacher_logprobs = None + rollout_state.teacher_targets = None rollout_state.routed_experts = None rollout_state.finish_reason = None rollout_state.response_mask = None diff --git a/xtuner/v1/rl/distillation/rollout_teacher_manager.py b/xtuner/v1/rl/distillation/rollout_teacher_manager.py index 4cd5c12a9a..c8131d2063 100644 --- a/xtuner/v1/rl/distillation/rollout_teacher_manager.py +++ b/xtuner/v1/rl/distillation/rollout_teacher_manager.py @@ -9,7 +9,7 @@ import httpx -from xtuner.v1.data_proto.rl_data import RolloutState, Status, get_group_status +from xtuner.v1.data_proto.rl_data import RolloutState, Status, TeacherTargets, get_group_status from xtuner.v1.rl.utils import create_task from xtuner.v1.utils import get_logger @@ -118,8 +118,11 @@ async def compute_logprobs(self, state: RolloutState) -> RolloutState: logprob_start_len=logprob_start_len, expanded_prompt_len=expanded_prompt_len, ) - state.teacher_tokens = teacher_tokens - state.teacher_logprobs = teacher_logprobs + state.teacher_targets = TeacherTargets( + kind="topk" if self.target_config.uses_topk_targets else "sampled", + tokens=teacher_tokens, + logprobs=teacher_logprobs, + ) return state except (httpx.HTTPStatusError, httpx.RequestError, ValueError) as exc: if attempt_idx >= self.config.max_retry_per_sample: diff --git a/xtuner/v1/rl/distillation/trainer_adapter.py b/xtuner/v1/rl/distillation/trainer_adapter.py index 9ca8c38bb2..67b1116e53 100644 --- a/xtuner/v1/rl/distillation/trainer_adapter.py +++ b/xtuner/v1/rl/distillation/trainer_adapter.py @@ -7,7 +7,7 @@ import torch -from xtuner.v1.data_proto.rl_data import RolloutState +from xtuner.v1.data_proto.rl_data import RolloutState, TeacherTargets from .config import DistillationConfig, RolloutTeacherScorerConfig, TrainTeacherManagerConfig @@ -64,10 +64,59 @@ def rollout_teacher_targets( if self._config is None or not self._config.rollout_teachers: return None, None + if state.teacher_targets is None: + raise ValueError(f"Rollout state has no Teacher targets: rollout_id={state.rollout_id}") + expected_kind = "sampled" if self._config.loss_config.uses_sampled_token_targets else "topk" + if state.teacher_targets.kind != expected_kind: + raise ValueError( + f"Expected {expected_kind} Teacher targets, got {state.teacher_targets.kind}: " + f"rollout_id={state.rollout_id}" + ) return self._align_rollout_teacher_targets( - state, + state.teacher_targets, shifted_labels=shifted_labels, target_start=target_start, + ignore_idx=self._config.loss_config.ignore_idx, + ) + + @staticmethod + def _align_rollout_teacher_targets( + targets: TeacherTargets, + *, + shifted_labels: Sequence[int], + target_start: int, + ignore_idx: int, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + sequence_length = len(shifted_labels) + if target_start < 0 or target_start > sequence_length: + raise ValueError( + f"Teacher target_start must be within the shifted sequence: {target_start} vs {sequence_length}" + ) + if any(label != ignore_idx for label in shifted_labels[:target_start]): + raise ValueError("Teacher target prefix must contain only ignored labels") + + expected_rows = sequence_length - target_start + if len(targets.logprobs) != expected_rows: + raise ValueError( + "Teacher logprobs must align with the target suffix: " + f"expected {expected_rows} rows, got {len(targets.logprobs)}" + ) + + if targets.kind == "sampled": + sampled_logprobs = cast(list[float], targets.logprobs) + return ( + torch.tensor([0.0] * target_start + sampled_logprobs, dtype=torch.float32).unsqueeze(0), + None, + ) + + topk_tokens = cast(list[list[int]], targets.tokens) + topk_logprobs = cast(list[list[float]], targets.logprobs) + top_k = len(topk_logprobs[0]) + prompt_tokens = [[0] * top_k for _ in range(target_start)] + prompt_logprobs = [[0.0] * top_k for _ in range(target_start)] + return ( + torch.tensor(prompt_logprobs + topk_logprobs, dtype=torch.float32).unsqueeze(0), + torch.tensor(prompt_tokens + topk_tokens, dtype=torch.int64).unsqueeze(0), ) def reward_scalars( @@ -127,63 +176,3 @@ def timing_scalars(self, worker_log_items: Sequence[WorkerLogItem]) -> dict[str, scalars[f"time/train_teacher/total/{phase}"] = total scalars[f"time/train_teacher_{phase}"] = total return scalars - - def _align_rollout_teacher_targets( - self, - state: RolloutState, - *, - shifted_labels: Sequence[int], - target_start: int, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - assert self._config is not None - loss_config = self._config.loss_config - sequence_length = len(shifted_labels) - context = f"rollout_id={state.rollout_id}, group_id={state.group_id}" - if target_start < 0 or target_start > sequence_length: - raise ValueError( - "Teacher target_start must be within the shifted sequence: " - f"{target_start} vs {sequence_length}; {context}" - ) - if any(label != loss_config.ignore_idx for label in shifted_labels[:target_start]): - raise ValueError(f"Teacher target prefix must contain only ignored labels; {context}") - - expected_target_rows = sequence_length - target_start - raw_teacher_logprobs = state.teacher_logprobs - if raw_teacher_logprobs is None or len(raw_teacher_logprobs) != expected_target_rows: - actual_rows = None if raw_teacher_logprobs is None else len(raw_teacher_logprobs) - raise ValueError( - "Teacher logprobs must align with the target suffix: " - f"expected {expected_target_rows} rows, got {actual_rows}; {context}" - ) - - if loss_config.uses_sampled_token_targets: - sampled_logprobs = cast(list[float], raw_teacher_logprobs) - teacher_logprobs = torch.tensor( - [0.0] * target_start + sampled_logprobs, - dtype=torch.float32, - ).unsqueeze(0) - return teacher_logprobs, None - - top_k = cast(int, loss_config.top_k) - raw_teacher_tokens = state.teacher_tokens - if raw_teacher_tokens is None or len(raw_teacher_tokens) != expected_target_rows: - actual_rows = None if raw_teacher_tokens is None else len(raw_teacher_tokens) - raise ValueError( - "Teacher token ids must align with the target suffix: " - f"expected {expected_target_rows} rows, got {actual_rows}; {context}" - ) - - topk_tokens = cast(list[list[int]], raw_teacher_tokens) - topk_logprobs = cast(list[list[float]], raw_teacher_logprobs) - - prompt_target_tokens = [[0] * top_k for _ in range(target_start)] - prompt_teacher_logprobs = [[0.0] * top_k for _ in range(target_start)] - teacher_logprobs = torch.tensor( - prompt_teacher_logprobs + topk_logprobs, - dtype=torch.float32, - ).unsqueeze(0) - target_token_ids = torch.tensor( - prompt_target_tokens + topk_tokens, - dtype=torch.int64, - ).unsqueeze(0) - return teacher_logprobs, target_token_ids From 9a1346aae41256acadfb504759c37e94ef7de1ca Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Sun, 13 Sep 2026 14:51:35 +0000 Subject: [PATCH 15/20] Align RL tests with current interfaces --- .../rl/test_multi_task_agent_loop_manager.py | 2 +- tests/rl/test_on_policy_distillation.py | 2 +- tests/rl/test_producer.py | 1 - tests/rl/test_rl_disaggregated_trainer.py | 2 + tests/rl/test_rollout_teacher_scorer.py | 98 +++++++++++++++++++ 5 files changed, 102 insertions(+), 3 deletions(-) create mode 100644 tests/rl/test_rollout_teacher_scorer.py diff --git a/tests/rl/test_multi_task_agent_loop_manager.py b/tests/rl/test_multi_task_agent_loop_manager.py index 5447f0ad90..515e3cbf21 100644 --- a/tests/rl/test_multi_task_agent_loop_manager.py +++ b/tests/rl/test_multi_task_agent_loop_manager.py @@ -255,7 +255,7 @@ def test_manager_config_injects_validity_check_when_building_agent_loop(self): judger=None, logger=None, is_valid_sample_fn=is_valid_sample_fn, - distillation_config=None, + rollout_teacher_scorer_config=None, ) self.assertIs(manager.task_runners[0].agent_loop, built_agent_loop) diff --git a/tests/rl/test_on_policy_distillation.py b/tests/rl/test_on_policy_distillation.py index 2af33ccd3e..f82e790a1c 100644 --- a/tests/rl/test_on_policy_distillation.py +++ b/tests/rl/test_on_policy_distillation.py @@ -342,7 +342,7 @@ def test_distillation_loss_kwargs_split_teacher_indices_with_ignore_padding(self padding_values = [call.kwargs["padding_value"] for call in split.call_args_list] self.assertEqual(padding_values, [0.0, 0, -1]) - def test_finalize_train_metrics_handles_policy_and_distillation_keys(self) -> None: + def test_distillation_loss_finalizes_policy_and_distillation_keys(self) -> None: loss_cfg = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) metrics = loss_cfg.finalize_metrics( { diff --git a/tests/rl/test_producer.py b/tests/rl/test_producer.py index 8ea9883104..6ac0c0fe69 100644 --- a/tests/rl/test_producer.py +++ b/tests/rl/test_producer.py @@ -119,7 +119,6 @@ async def mock_gen(rs, **kwargs): mock_agent_loop.generate_group = mock_gen mock_agent_loop.is_valid_sample_fn = is_valid_sample_fn - mock_agent_loop.teacher_clients = {} return mock_agent_loop def _build_context( diff --git a/tests/rl/test_rl_disaggregated_trainer.py b/tests/rl/test_rl_disaggregated_trainer.py index add03a5086..914a5c86cf 100644 --- a/tests/rl/test_rl_disaggregated_trainer.py +++ b/tests/rl/test_rl_disaggregated_trainer.py @@ -35,6 +35,7 @@ ProduceBatchResult, ProduceBatchStatus, ) +from xtuner.v1.rl.distillation import DistillationTrainerAdapter from xtuner.v1.train.rl_trainer import RLDisaggregatedTrainer, _validate_sync_intervals @@ -120,6 +121,7 @@ def _make_trainer(self, agent_loop_manager): trainer._benchmark_training_samples = 0 trainer._benchmark_training_tokens = 0 trainer._cpu_resource_manager = None + trainer._distillation = DistillationTrainerAdapter(None) trainer._train_worker_cfg = SimpleNamespace(pack_max_length=16) trainer._rollout_config = SimpleNamespace(weight_update_host=None, weight_update_port=30000) trainer._meta = SimpleNamespace( diff --git a/tests/rl/test_rollout_teacher_scorer.py b/tests/rl/test_rollout_teacher_scorer.py new file mode 100644 index 0000000000..85354c59c9 --- /dev/null +++ b/tests/rl/test_rollout_teacher_scorer.py @@ -0,0 +1,98 @@ +import unittest +from unittest.mock import AsyncMock, MagicMock + +from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.rl.distillation import RolloutTeacherScorer + + +class TestRolloutTeacherScorer(unittest.IsolatedAsyncioTestCase): + @staticmethod + def _state(uid: int, status: Status = Status.COMPLETED) -> RolloutState: + return RolloutState( + rollout_id=uid, + group_id=uid, + message=[{"role": "user", "content": f"prompt {uid}"}], + prompt_ids=[uid], + tokens=None, + response=None, + response_ids=[uid], + status=status, + extra_fields={"origin_data_source": "math"}, + ) + + @staticmethod + def _client(events: list[int]): + client = MagicMock() + + async def compute_logprobs(state: RolloutState) -> RolloutState: + events.append(state.rollout_id or 0) + return state + + client.compute_logprobs = AsyncMock(side_effect=compute_logprobs) + client.aclose = AsyncMock() + return client + + async def test_scores_samples_immediately_without_filter_deferral(self): + events: list[int] = [] + client = self._client(events) + scorer = RolloutTeacherScorer( + {"teacher": client}, + {"math": "teacher"}, + defers_to_filter=False, + ) + + state = await scorer.on_sample_ready(self._state(1)) + + self.assertEqual(state.rollout_id, 1) + self.assertEqual(events, [1]) + client.compute_logprobs.assert_awaited_once() + + async def test_scores_only_surviving_completed_group_when_deferred(self): + events: list[int] = [] + client = self._client(events) + scorer = RolloutTeacherScorer( + {"teacher": client}, + {"math": "teacher"}, + defers_to_filter=True, + ) + states = [self._state(1), self._state(2)] + + await scorer.on_sample_ready(states[0]) + result = await scorer.on_group_ready(states) + + self.assertEqual([state.rollout_id for state in result], [1, 2]) + self.assertEqual(sorted(events), [1, 2]) + self.assertEqual(client.compute_logprobs.await_count, 2) + + async def test_skips_non_completed_group(self): + events: list[int] = [] + client = self._client(events) + scorer = RolloutTeacherScorer( + {"teacher": client}, + {"math": "teacher"}, + defers_to_filter=True, + ) + + result = await scorer.on_group_ready([self._state(1), self._state(2, Status.FAILED)]) + + self.assertEqual([state.rollout_id for state in result], [1, 2]) + self.assertEqual(events, []) + client.compute_logprobs.assert_not_awaited() + + async def test_aclose_is_idempotent(self): + events: list[int] = [] + client = self._client(events) + scorer = RolloutTeacherScorer( + {"teacher": client}, + {"math": "teacher"}, + defers_to_filter=False, + ) + + await scorer.aclose() + await scorer.aclose() + + client.aclose.assert_awaited_once() + + +if __name__ == "__main__": + unittest.main() From 2481505128fe6c1f51086558cbf45f103dd0610c Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Mon, 14 Sep 2026 06:47:36 +0000 Subject: [PATCH 16/20] Refine distillation teacher integration --- .../rl/agent_loop/single_turn_agent_loop.py | 8 +- .../rl/distillation/train_teacher_manager.py | 73 +++++++++----- xtuner/v1/rl/distillation/trainer_adapter.py | 88 +++++++++++------ xtuner/v1/rl/{ => trainer}/model_utils.py | 13 ++- xtuner/v1/rl/trainer/worker.py | 8 +- xtuner/v1/train/rl_trainer.py | 98 +++++++------------ 6 files changed, 165 insertions(+), 123 deletions(-) rename xtuner/v1/rl/{ => trainer}/model_utils.py (74%) diff --git a/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py b/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py index cc185538ba..2368d1c3b5 100644 --- a/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py +++ b/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py @@ -1,4 +1,4 @@ -from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams +from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController from xtuner.v1.rl.trace.rollout_api import trace_rollout_endpoint, trace_rollout_remote @@ -80,4 +80,10 @@ async def generate_sample( self.rollout_ctl.generate, # type: ignore[attr-defined] rollout_state, ) + # 非 COMPLETED 状态（如被截断、放弃等）直接早退，不触发打分 + if rollout_state.status != Status.COMPLETED: + return rollout_state + if self.judger is not None and not self.enable_batch_judge: + # 如果开启了批量打分，则在 generate_group 里统一打分，不在这里逐条打分 + rollout_state = await self.run_judger(rollout_state) return rollout_state diff --git a/xtuner/v1/rl/distillation/train_teacher_manager.py b/xtuner/v1/rl/distillation/train_teacher_manager.py index 9b5b5eb8bb..1279d67e34 100644 --- a/xtuner/v1/rl/distillation/train_teacher_manager.py +++ b/xtuner/v1/rl/distillation/train_teacher_manager.py @@ -10,8 +10,8 @@ from xtuner.v1.data_proto.sequence_context import SequenceContext from xtuner.v1.loss import LogProbConfig, LogProbContext, TopKLogProbConfig from xtuner.v1.model.compose.base import BaseComposeConfig -from xtuner.v1.rl.model_utils import FrozenModel, build_frozen_model -from xtuner.v1.utils import get_device, get_torch_device_module +from xtuner.v1.rl.trainer.model_utils import FrozenModel, build_frozen_model +from xtuner.v1.utils import get_device, get_logger, get_torch_device_module from .config import TrainTeacherManagerConfig @@ -22,6 +22,7 @@ DEVICE = get_device() DEVICE_MODULE = get_torch_device_module() +logger = get_logger() @dataclass @@ -226,28 +227,50 @@ def _teacher_on_device( teacher_timing.offload += offload_elapsed @staticmethod - def _teacher_seq_ctx(seq_ctx: SequenceContext, *, is_composed: bool) -> SequenceContext: - overrides = { - "rollout_routed_experts": None, - "offload_rollout_routed_experts": False, - } - if not is_composed: - # A VLM tokenizer supplies 3D M-RoPE position ids for every sample, - # including text-only samples. Plain language-model Teachers need - # standard 2D packed positions. Passing ``None`` makes - # SequenceContext rebuild them from the packed sequence lengths. - # Visual fields are irrelevant to the text Teacher and may refer to - # a pack routed to another Teacher, so do not retain them here. - overrides.update( - position_ids=None, - image_grid_thw=None, - deepstack_visual_embeds=None, - visual_pos_masks=None, - pixel_values=None, - inputs_embeds=None, - num_img_tokens=None, + def construct_teacher_seq_ctx(seq_ctx: SequenceContext, *, is_composed: bool) -> SequenceContext: + """Construct a Teacher context from the Student rollout context. + + Rollout-only expert-routing metadata must not be reused by the Teacher. A plain language-model Teacher can only + score text tokens. Composed/VLM Teachers retain their visual fields and M-RoPE positions; plain-text Teachers + require text-only contexts with 2D packed positions. + """ + has_visual_inputs = seq_ctx.pixel_values is not None + if not is_composed and has_visual_inputs: + logger.warning( + "A plain-text Teacher received a multimodal rollout context; " + "the sample cannot be scored by this Teacher." ) - return seq_ctx.copy(**overrides) + raise ValueError("Plain-text Teacher cannot score a rollout context containing visual inputs") + + position_ids = seq_ctx.position_ids + if not is_composed and position_ids is not None and position_ids.ndim == 3: + # SequenceContext rebuilds standard 2D packed positions. + position_ids = None + + return SequenceContext( + input_ids=seq_ctx.input_ids, + cu_seq_lens_q=seq_ctx.cu_seq_lens_q, + cu_seq_lens_k=seq_ctx.cu_seq_lens_k, + max_length_q=seq_ctx.max_length_q, + max_length_k=seq_ctx.max_length_k, + num_padding=seq_ctx.num_padding, + sequence_parallel_mesh=seq_ctx.sequence_parallel_mesh, + block_table=seq_ctx.block_table, + device=seq_ctx.device, + position_ids=position_ids, + image_grid_thw=seq_ctx.image_grid_thw if is_composed else None, + deepstack_visual_embeds=seq_ctx.deepstack_visual_embeds if is_composed else None, + visual_pos_masks=seq_ctx.visual_pos_masks if is_composed else None, + pixel_values=seq_ctx.pixel_values if is_composed else None, + inputs_embeds=seq_ctx.inputs_embeds if is_composed else None, + num_img_tokens=seq_ctx.num_img_tokens if is_composed else None, + rollout_routed_experts=None, + offload_rollout_routed_experts=False, + raw_input_ids=seq_ctx._raw_input_ids, + raw_inputs_embeds=seq_ctx._raw_inputs_embeds, + shard_start=seq_ctx._shard_start, + shard_size=seq_ctx._shard_size, + ) def _compute_topk_targets( self, @@ -281,7 +304,7 @@ def _compute_topk_targets( assert loss_ctx is not None with torch.no_grad(): output = teacher( - seq_ctx=self._teacher_seq_ctx( + seq_ctx=self.construct_teacher_seq_ctx( seq_ctx, is_composed=self._teacher_is_composed[teacher_index], ), @@ -323,7 +346,7 @@ def _compute_sampled_logprobs( ) with torch.no_grad(): output = teacher( - seq_ctx=self._teacher_seq_ctx( + seq_ctx=self.construct_teacher_seq_ctx( seq_ctx, is_composed=self._teacher_is_composed[teacher_index], ), diff --git a/xtuner/v1/rl/distillation/trainer_adapter.py b/xtuner/v1/rl/distillation/trainer_adapter.py index 67b1116e53..a2e04879a8 100644 --- a/xtuner/v1/rl/distillation/trainer_adapter.py +++ b/xtuner/v1/rl/distillation/trainer_adapter.py @@ -2,8 +2,8 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, cast +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, cast import torch @@ -46,38 +46,48 @@ def task_adv_weight(self) -> float: return 1.0 return self._config.loss_config.task_adv_weight - @property - def teacher_index_by_data_source(self) -> Mapping[str, int] | None: - if self._config is None or not self._config.train_teachers: - return None - return self._config.teacher_index_by_data_source - def rollout_teacher_targets( self, state: RolloutState, *, shifted_labels: Sequence[int], - target_start: int, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - """Return aligned rollout Teacher targets, or no fields when - disabled.""" - - if self._config is None or not self._config.rollout_teachers: - return None, None - if state.teacher_targets is None: - raise ValueError(f"Rollout state has no Teacher targets: rollout_id={state.rollout_id}") - expected_kind = "sampled" if self._config.loss_config.uses_sampled_token_targets else "topk" - if state.teacher_targets.kind != expected_kind: - raise ValueError( - f"Expected {expected_kind} Teacher targets, got {state.teacher_targets.kind}: " - f"rollout_id={state.rollout_id}" + ) -> dict[str, torch.Tensor]: + """Return all Teacher fields needed to populate one training batch.""" + + if self._config is None: + return {} + + fields: dict[str, torch.Tensor] = {} + if self._config.rollout_teachers: + target_start = next( + (index for index, label in enumerate(shifted_labels) if label != self._config.loss_config.ignore_idx), + len(shifted_labels), ) - return self._align_rollout_teacher_targets( - state.teacher_targets, - shifted_labels=shifted_labels, - target_start=target_start, - ignore_idx=self._config.loss_config.ignore_idx, - ) + if state.teacher_targets is None: + raise ValueError(f"Rollout state has no Teacher targets: rollout_id={state.rollout_id}") + expected_kind = "sampled" if self._config.loss_config.uses_sampled_token_targets else "topk" + if state.teacher_targets.kind != expected_kind: + raise ValueError( + f"Expected {expected_kind} Teacher targets, got {state.teacher_targets.kind}: " + f"rollout_id={state.rollout_id}" + ) + teacher_logprobs, target_token_ids = self._align_rollout_teacher_targets( + state.teacher_targets, + shifted_labels=shifted_labels, + target_start=target_start, + ignore_idx=self._config.loss_config.ignore_idx, + ) + fields["teacher_logprobs"] = teacher_logprobs + if target_token_ids is not None: + fields["target_token_ids"] = target_token_ids + + if self._config.train_teachers: + data_source = state.extra_fields.get("origin_data_source") + if not isinstance(data_source, str): + raise ValueError(f"Missing origin_data_source for rollout_id={state.rollout_id}") + teacher_index = self._config.teacher_index_by_data_source[data_source] + fields["teacher_indices"] = torch.full((1, len(shifted_labels)), teacher_index, dtype=torch.long) + return fields @staticmethod def _align_rollout_teacher_targets( @@ -144,6 +154,28 @@ def reward_scalars( if rewards } + def metric_scalars(self, train_metrics: Sequence[Any]) -> dict[str, float]: + """Format rank-zero mini-batch distillation metrics for the trainer + log.""" + + if self._config is None: + return {} + + metric_values: dict[str, list[float]] = {} + for mini_batch_metrics in train_metrics: + for key, value in mini_batch_metrics.items(): + metric_values.setdefault(key, []).append(float(value)) + + scalars: dict[str, float] = {} + for key, values in metric_values.items(): + if "opd" not in key and "distillation" not in key: + continue + value = sum(values) / len(values) + scalars[f"distillation/{key}"] = float(value) + if key in {"opd_reverse_kl", "opd_abs_logprob_loss"}: + scalars[key] = float(value) + return scalars + def timing_scalars(self, worker_log_items: Sequence[WorkerLogItem]) -> dict[str, float]: """Aggregate train Teacher timings across workers on the critical path.""" diff --git a/xtuner/v1/rl/model_utils.py b/xtuner/v1/rl/trainer/model_utils.py similarity index 74% rename from xtuner/v1/rl/model_utils.py rename to xtuner/v1/rl/trainer/model_utils.py index 404569d956..58ae6b67bf 100644 --- a/xtuner/v1/rl/model_utils.py +++ b/xtuner/v1/rl/trainer/model_utils.py @@ -26,12 +26,23 @@ def build_frozen_model( model = model_cfg.build() if isinstance(model_cfg, BaseComposeConfig): - assert model_cfg.text_config.float8_cfg is None, "BaseComposeConfig does not support float8" + text_float8_cfg = getattr(model_cfg.text_config, "float8_cfg", None) + if text_float8_cfg is not None and text_float8_cfg.enable_float8: + float8_handler = Float8Handler( + scaling_granularity_gemm=text_float8_cfg.scaling_granularity_gemm, + scaling_granularity_grouped_gemm=text_float8_cfg.scaling_granularity_grouped_gemm, + ) + else: + float8_handler = None if fsdp_cfg is None: fsdp_cfg = FSDPConfig(recompute_ratio=0, cpu_offload=False, requires_grad=False) model = model.fully_shard(fsdp_cfg) model.from_hf(hf_path=load_from) model.eval() # type: ignore + if float8_handler is not None: + # Composed models use FP8 only in the text backbone. The vision + # tower and projector remain in BF16. + float8_handler.precompute_float8_dynamic_scale_for_fsdp(model.language_model) # type: ignore else: model_cfg = cast(TransformerConfig, model_cfg) if model_cfg.float8_cfg is not None and model_cfg.float8_cfg.enable_float8: diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index 099eff2336..03472c2f29 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -52,7 +52,7 @@ BaseRLLossContext, kl_penalty, ) -from xtuner.v1.rl.model_utils import build_frozen_model +from xtuner.v1.rl.trainer.model_utils import build_frozen_model from xtuner.v1.rl.utils import SingleAcceleratorWorker from xtuner.v1.rl.weight_update import WeightUpdater from xtuner.v1.train.trainer import LoadCheckpointConfig @@ -707,20 +707,18 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo f"Rank{self.rank} Rollout {rollout_idx} prepare_inputs elapsed=" f"{time.perf_counter() - prepare_inputs_begin:.4f}s" ) + del data_batches # When sp_mesh.size() > 1, get the sp_split shifted_labels and rollout_logprobs shifted_labels_list = [loss_ctx.loss_kwargs.shifted_labels for loss_ctx in loss_ctx_list] rollout_logprobs_list = [loss_ctx.loss_kwargs.rollout_logprobs for loss_ctx in loss_ctx_list] - del data_batches - # compute old logprobs old_logprobs_list = self.compute_actor_logprobs(seq_ctx_list, shifted_labels_list) for old_logprobs, loss_ctx in zip(old_logprobs_list, loss_ctx_list): loss_ctx.loss_kwargs.old_logprobs = old_logprobs worker_log_item: WorkerLogItem = {"train_entropy": 0.0, "train_metrics": [], "sft_train_metrics": {}} - logger_msg = f"Rollout {rollout_idx}: " # compute entropy @@ -809,7 +807,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo # run, which bounds their GPU residency by the optimizer group size. self._maybe_offload_logprob(loss_ctx_list) del old_logprobs_list - + teacher_timings = self._maybe_compute_train_teacher_outputs(seq_ctx_list, loss_ctx_list) if teacher_timings is not None: worker_log_item.update({"teacher_timings": teacher_timings.to_dict()}) diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 57ee1a7033..c59d842fc2 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -657,11 +657,12 @@ class BaseRLTrainer: _distillation: DistillationTrainerAdapter def _init_common(self, cfg: BaseRLTrainerConfig, *, meta_path: str, logger_tag: str) -> None: - if cfg.distillation_config is not None and cfg.distillation_config.rollout_teachers: + distillation_config = cfg.distillation_config + if distillation_config is not None and distillation_config.rollout_teachers: endpoint_map = json.loads(os.environ.get("XTUNER_OPD_TEACHER_ENDPOINTS_JSON", "{}")) - cfg.distillation_config = cfg.distillation_config.resolve_teacher_endpoints(endpoint_map) + distillation_config = distillation_config.resolve_teacher_endpoints(endpoint_map) - self._distillation = DistillationTrainerAdapter(cfg.distillation_config) + self._distillation = DistillationTrainerAdapter(distillation_config) check_fa3() self._init_work_dir_and_meta(cfg, meta_path) @@ -1021,7 +1022,7 @@ def _train_one_batch( step_timer_dict: dict, *, offload_rollout_before_train: bool = False, - resume_train_before_train: bool = False, + onload_train_before_train: bool = False, raw_rewards_sum: float = 0.0, raw_rewards_count: int = 0, ) -> TrainInfo: @@ -1034,14 +1035,17 @@ def _train_one_batch( self._save_trajectories(train_batch, train_trajectory_path) self.logger.info(f"Train step {train_step} train trajectories saved to {train_trajectory_path}") - # 共卡训练前切换资源：检查 rollout -> offload rollout -> resume train NCCL。 + # 共卡训练前切换资源：检查 rollout -> offload rollout -> onload train。 if offload_rollout_before_train: ray.get(self.rollout_controller.offload.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT) - if resume_train_before_train: + if onload_train_before_train: if getattr(self, "_train_nccl_suspended", False): with timer("resume_train_nccl", step_timer_dict): self.train_controller.resume_train_nccl_process_groups() self._train_nccl_suspended = False + with timer("onload", step_timer_dict): + self.train_controller.onload(target="all") + self.logger.info("Training controller loaded") with timer("prepare_data", step_timer_dict): data_batches, data_info = self._prepare_train_data( @@ -1141,7 +1145,7 @@ def _prepare_train_data( # trainable segments that share one reward; counting that reward once per session keeps # rewards/* from being weighted by segment count. cluster_rewards_list: list[float] = [] - teacher_reward_observations: list[tuple[RolloutState, float]] = [] + distillation_reward_observations: list[tuple[RolloutState, float]] = [] advantages_list = [] prompt_len_list = [] response_len_list = [] @@ -1150,7 +1154,6 @@ def _prepare_train_data( training_samples = 0 data_batches = [] - teacher_index_by_data_source = self._distillation.teacher_index_by_data_source for j, group in enumerate(data_groups): if not is_valid_for_training(group, self.logger): @@ -1170,11 +1173,6 @@ def _prepare_train_data( assert prompt_ids is not None and len(prompt_ids) > 0, ( f"Prompt ids cannot be None or empty in data: {group[0]}" ) - for data in group: - # 有可能有重复，但是没有其他更好办法 - turns = data.extra_fields.get("agent_tool_turns") - if isinstance(turns, int): - tool_turns_list.append(turns) # Collect rewards independently from task-advantage computation. Pure OPD may omit # rewards entirely; when rewards are present they remain useful observability signals. cluster_index_by_key: dict[Any, int] = {} @@ -1182,6 +1180,10 @@ def _prepare_train_data( cluster_representatives: list[RolloutState] = [] sample_cluster_indices: list[int] = [] for data in group: + # 有可能有重复，但是没有其他更好办法 + turns = data.extra_fields.get("agent_tool_turns") + if isinstance(turns, int): + tool_turns_list.append(turns) if data.reward is None or "score" not in data.reward: assert task_adv_weight == 0, f"Reward is missing or does not contain 'score' key in data: {data}" continue @@ -1198,7 +1200,7 @@ def _prepare_train_data( sample_cluster_indices.append(cluster_index) cluster_rewards_list.extend(cluster_rewards) - teacher_reward_observations.extend(zip(cluster_representatives, cluster_rewards)) + distillation_reward_observations.extend(zip(cluster_representatives, cluster_rewards)) if task_adv_weight == 0: sample_advantages = [0.0] * len(group) @@ -1213,10 +1215,6 @@ def _prepare_train_data( prompt_repeat_k = len(group) for i in range(prompt_repeat_k): - teacher_index = None - if teacher_index_by_data_source is not None: - data_source = group[i].extra_fields["origin_data_source"] - teacher_index = teacher_index_by_data_source[data_source] if group[i].input_ids is not None: raw_input_ids = cast(list[int], group[i].input_ids) labels = cast(list[int] | None, group[i].labels) @@ -1238,16 +1236,9 @@ def _prepare_train_data( input_ids = raw_input_ids[:-1] shifted_labels = labels[1:] - teacher_logprobs = None - target_token_ids = None - teacher_response_start = next( - (index for index, label in enumerate(shifted_labels) if label != -100), - len(shifted_labels), - ) - teacher_logprobs, target_token_ids = self._distillation.rollout_teacher_targets( + teacher_fields = self._distillation.rollout_teacher_targets( group[i], shifted_labels=shifted_labels, - target_start=teacher_response_start, ) prompt_len = sum(label == -100 for label in shifted_labels) response_len = len(shifted_labels) - prompt_len @@ -1286,15 +1277,7 @@ def _prepare_train_data( "advantage": actual_advantages, "rollout_logprobs": rollout_logprobs, } - if teacher_logprobs is not None: - data_dict["teacher_logprobs"] = teacher_logprobs - if target_token_ids is not None: - data_dict["target_token_ids"] = target_token_ids - if teacher_index is not None: - data_dict["teacher_indices"] = torch.full_like( - shifted_labels_t, - teacher_index, - ) + data_dict.update(teacher_fields) seq_ctx.rollout_routed_experts = group[i].routed_experts data_batches.append(data_dict) @@ -1345,9 +1328,11 @@ def _prepare_train_data( shifted_labels = [-100] * (len(prompt_ids) - 1) + response_labels shifted_labels_t = torch.tensor(shifted_labels, dtype=torch.int64).unsqueeze(0) - base_advantage = sample_advantages[i] - actual_advantages = [0.0 if label == -100 else base_advantage for label in shifted_labels] - advantages_list.extend(actual_advantages) + advatnages_val = sample_advantages[i] + actual_advantages = [advatnages_val] * len(prompt_ids) + [ + 0.0 if mask == 0 else advatnages_val for mask in response_mask + ] + advantages_list.extend(actual_advantages[:-1]) assert len(input_ids) <= pack_max_length, f"{len(input_ids)} vs {pack_max_length}" training_tokens += len(input_ids) @@ -1372,20 +1357,11 @@ def _prepare_train_data( "advantage": actual_advantages, "rollout_logprobs": rollout_logprobs, } - teacher_logprobs, target_token_ids = self._distillation.rollout_teacher_targets( + teacher_fields = self._distillation.rollout_teacher_targets( group[i], shifted_labels=shifted_labels, - target_start=len(prompt_ids) - 1, ) - if teacher_logprobs is not None: - data_dict["teacher_logprobs"] = teacher_logprobs - if target_token_ids is not None: - data_dict["target_token_ids"] = target_token_ids - if teacher_index is not None: - data_dict["teacher_indices"] = torch.full_like( - shifted_labels_t, - teacher_index, - ) + data_dict.update(teacher_fields) seq_ctx.rollout_routed_experts = group[i].routed_experts # n,layer*expert @@ -1420,7 +1396,7 @@ def _prepare_train_data( "prompt_len/min": prompt_len_t.min().item(), "prompt_len/max": prompt_len_t.max().item(), } - info_dict.update(self._distillation.reward_scalars(teacher_reward_observations)) + info_dict.update(self._distillation.reward_scalars(distillation_reward_observations)) if tool_turns_list: tool_turns_t = torch.tensor(tool_turns_list, dtype=torch.float32) info_dict["tool_turns/mean"] = tool_turns_t.mean().item() @@ -1546,7 +1522,6 @@ def _log_step( all_scalars.update({"entropy/train": rank0_log_item["train_entropy"]}) all_scalars.update(self._distillation.timing_scalars(worker_log_items)) - for worker_idx, log_item in enumerate(worker_log_items): if not self._display_all_workers_log and worker_idx > 0: break @@ -1558,15 +1533,12 @@ def _log_step( for key, value in mini_batch_metrics.items(): avg_value = sum(value) / len(value) all_scalars.update({f"train_metrics/worker_{worker_idx}/step_avg_{key}": avg_value}) - if worker_idx == 0 and ("opd" in key or "distillation" in key): - all_scalars[f"distillation/{key}"] = avg_value - if key in ("opd_reverse_kl", "opd_abs_logprob_loss"): - all_scalars[key] = avg_value - rank_sft_log = log_item["sft_train_metrics"] for k, v in rank_sft_log.items(): all_scalars.update({f"sft_train_metrics/worker_{worker_idx}/{k}": v}) + all_scalars.update(self._distillation.metric_scalars(rank0_log_item["train_metrics"])) + self._log_mini_batch_metrics(train_info["workers_log_item"]) if eval_info: @@ -1602,14 +1574,14 @@ def _save_trajectories(self, data_groups: list[list[RolloutState]], save_path: P with open(save_path, "w", encoding="utf-8") as f: summary = { "reward_mean": rewards_tensor.mean().item(), - "reward_std": rewards_tensor.std(unbiased=False).item(), + "reward_std": rewards_tensor.std().item(), "reward_max": rewards_tensor.max().item(), "reward_min": rewards_tensor.min().item(), "response_len_mean": response_lens.mean().item(), - "response_len_std": response_lens.std(unbiased=False).item(), + "response_len_std": response_lens.std().item(), "response_len_max": response_lens.max().item(), "response_len_min": response_lens.min().item(), - "total_len": len(response_len_list), + "total_len": len(rewards), } json.dump(summary, f, ensure_ascii=False, separators=(",", ":")) f.write("\n") @@ -1863,7 +1835,7 @@ def fit(self): self._exp_tracker.close() close_trace() - def _fit(self) -> None: + def _fit(self): self.logger.info("Start RL training") if self._cur_step >= self._total_train_steps: self.logger.info(f"Train steps {self._total_train_steps} reached, stop training") @@ -1889,7 +1861,7 @@ def _fit(self) -> None: model_step = self._get_colocate_rollout_model_step(init_train_step) for train_step in range(init_train_step, self._total_train_steps + 1): self.logger.info(f"Train step {train_step}/{self._total_train_steps} start") - step_timer_dict: dict[str, float] = {} + step_timer_dict = {} with timer("step", step_timer_dict): # 共卡一次调用内完成生产和消费。 self.rl_health_manager.set_rollout_resources_available(True) @@ -1921,7 +1893,7 @@ def _fit(self) -> None: train_step, step_timer_dict, offload_rollout_before_train=True, - resume_train_before_train=True, + onload_train_before_train=True, raw_rewards_sum=produce_result.raw_rewards_sum, raw_rewards_count=produce_result.raw_rewards_count, ) @@ -1961,7 +1933,7 @@ def _fit_debug_train(self) -> None: train_step, step_timer_dict, offload_rollout_before_train=False, - resume_train_before_train=False, + onload_train_before_train=False, ) eval_log_info: dict[str, float] = {} produce_result = ProduceBatchResult(rollout_states=train_batch) From 312e9a8b9538adb231e275040f08fa8549393e38 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Mon, 14 Sep 2026 07:06:34 +0000 Subject: [PATCH 17/20] fix: validate rollout teacher endpoints --- xtuner/v1/rl/distillation/config.py | 13 ++++++++++++- xtuner/v1/rl/trainer/worker.py | 1 + 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/xtuner/v1/rl/distillation/config.py b/xtuner/v1/rl/distillation/config.py index b2f74990f3..36baaf3f74 100644 --- a/xtuner/v1/rl/distillation/config.py +++ b/xtuner/v1/rl/distillation/config.py @@ -74,6 +74,12 @@ class RolloutTeacherConfig(BaseModel): enable_prefix_caching: bool = False launch_config: RolloutTeacherLaunchConfig | None = None + @model_validator(mode="after") + def validate_endpoints(self) -> RolloutTeacherConfig: + if self.launch_config is None and not self.endpoints: + raise ValueError(f"Rollout Teacher {self.name!r} requires at least one endpoint") + return self + class TrainTeacherConfig(BaseModel): """Frozen Teacher loaded and executed by each training worker.""" @@ -299,7 +305,12 @@ def resolve_teacher_endpoints( teachers = [ teacher if not isinstance(teacher, RolloutTeacherConfig) or teacher.launch_config is None - else teacher.model_copy(update={"endpoints": endpoint_map[teacher.name]}) + else teacher.model_copy(update={"endpoints": endpoint_map.get(teacher.name, [])}) for teacher in self.teachers ] + missing_endpoints = [ + teacher.name for teacher in teachers if isinstance(teacher, RolloutTeacherConfig) and not teacher.endpoints + ] + if missing_endpoints: + raise ValueError(f"Rollout Teachers have no resolved endpoints: {missing_endpoints}") return self.model_copy(update={"teachers": teachers}) diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index 03472c2f29..47a74fee04 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -707,6 +707,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo f"Rank{self.rank} Rollout {rollout_idx} prepare_inputs elapsed=" f"{time.perf_counter() - prepare_inputs_begin:.4f}s" ) + del data_batches # When sp_mesh.size() > 1, get the sp_split shifted_labels and rollout_logprobs From c394e2b360ae703ebef4bfc9ca29a19d67be25d3 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Mon, 14 Sep 2026 07:17:10 +0000 Subject: [PATCH 18/20] docs: fix RL API and agent loop references --- docs/en/api/rl_config.rst | 2 +- docs/zh_cn/api/rl_config.rst | 2 +- docs/zh_cn/rl/advanced_tutorial/agent_loop.md | 10 +++++++++- recipe/verl_agent/common/agent_loop_verl_tool.py | 4 ++++ xtuner/v1/rl/agent_loop/gsm8k_with_tool.py | 2 ++ 5 files changed, 17 insertions(+), 3 deletions(-) diff --git a/docs/en/api/rl_config.rst b/docs/en/api/rl_config.rst index a36e097c0e..3aa97365b0 100644 --- a/docs/en/api/rl_config.rst +++ b/docs/en/api/rl_config.rst @@ -29,7 +29,7 @@ Agent Loop Config :toctree: generated :nosignatures: - rl.agent_loop.SingleTurnAgentLoopConfig + rl.agent_loop.single_turn_agent_loop.SingleTurnAgentLoopConfig rl.agent_loop_manager.AgentLoopManagerConfig rl.agent_loop_manager.TaskSpecConfig rl.agent_loop_manager.SamplerConfig diff --git a/docs/zh_cn/api/rl_config.rst b/docs/zh_cn/api/rl_config.rst index a36e097c0e..3aa97365b0 100644 --- a/docs/zh_cn/api/rl_config.rst +++ b/docs/zh_cn/api/rl_config.rst @@ -29,7 +29,7 @@ Agent Loop Config :toctree: generated :nosignatures: - rl.agent_loop.SingleTurnAgentLoopConfig + rl.agent_loop.single_turn_agent_loop.SingleTurnAgentLoopConfig rl.agent_loop_manager.AgentLoopManagerConfig rl.agent_loop_manager.TaskSpecConfig rl.agent_loop_manager.SamplerConfig diff --git a/docs/zh_cn/rl/advanced_tutorial/agent_loop.md b/docs/zh_cn/rl/advanced_tutorial/agent_loop.md index ac51618aed..b72ddd4d39 100644 --- a/docs/zh_cn/rl/advanced_tutorial/agent_loop.md +++ b/docs/zh_cn/rl/advanced_tutorial/agent_loop.md @@ -108,7 +108,8 @@ generate_sample(state) -> await rollout_ctl.generate.remote(state) -> PartialRolloutHandler.postprocess(state) -> 如果 state.status != COMPLETED，直接返回，不触发 Judger - -> 如果配置了 judger，调用 self.run_judger(state) + -> 无 validity check：generation 完成后进入 Teacher scorer，再按配置逐条 Judger + -> 有 validity check：组级 Judger -> maybe_filter_invalid_sample -> Teacher scorer ``` 典型配置： @@ -169,6 +170,11 @@ class CustomAgentLoop(AgentLoop): rollout_state.sample_params = rollout_state.sample_params or self.sample_params rollout_state = await self.rollout_ctl.generate.remote(rollout_state) + if rollout_state.status != Status.COMPLETED: + return rollout_state + + if self.judger is not None and not self.enable_batch_judge: + rollout_state = await self.run_judger(rollout_state) return rollout_state @@ -239,6 +245,8 @@ class ToolAgentLoop(AgentLoop): assert len(rollout_state.response_ids) == len(rollout_state.logprobs) assert len(rollout_state.response_ids) == len(rollout_state.response_mask) + if rollout_state.status == Status.COMPLETED and self.judger is not None and not self.enable_batch_judge: + rollout_state = await self.run_judger(rollout_state) return rollout_state ``` diff --git a/recipe/verl_agent/common/agent_loop_verl_tool.py b/recipe/verl_agent/common/agent_loop_verl_tool.py index 8e9f0e0a65..f78ef27264 100644 --- a/recipe/verl_agent/common/agent_loop_verl_tool.py +++ b/recipe/verl_agent/common/agent_loop_verl_tool.py @@ -147,4 +147,8 @@ async def generate_sample(self, rollout_state: RolloutState) -> RolloutState: # raw_prompt is updated in tool_agent_loop: apply_chat_template of tools rollout_state.extra_fields["raw_prompt"] = self.tokenizer.decode(rollout_state.prompt_ids) + # judge rollout_state + if self.judger is not None: + rollout_state = await self.judger.judge(rollout_state) + return rollout_state diff --git a/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py b/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py index aeeaf0df3d..13044d4f35 100644 --- a/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py +++ b/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py @@ -159,4 +159,6 @@ async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> Rollou assert len(rollout_state.response_ids) == len(rollout_state.response_mask) == len(rollout_state.logprobs), ( f"{len(rollout_state.response_ids)} vs {len(rollout_state.response_mask)} vs {len(rollout_state.logprobs)}" ) + if self.judger is not None and not self.enable_batch_judge: + rollout_state = await self.run_judger(rollout_state) return rollout_state From efe1a73ac6d631214ed77dbf1353139490b4f6f0 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Mon, 14 Sep 2026 07:27:17 +0000 Subject: [PATCH 19/20] fix docs --- docs/en/api/rl_config.rst | 2 +- docs/zh_cn/api/rl_config.rst | 2 +- xtuner/v1/rl/agent_loop/agent_loop.py | 13 ++++++++----- .../v1/rl/agent_loop_manager/agent_loop_manager.py | 11 ++++++++--- .../agent_loop_manager/disagg_agent_loop_manager.py | 11 ++++++++--- 5 files changed, 26 insertions(+), 13 deletions(-) diff --git a/docs/en/api/rl_config.rst b/docs/en/api/rl_config.rst index 3aa97365b0..a36e097c0e 100644 --- a/docs/en/api/rl_config.rst +++ b/docs/en/api/rl_config.rst @@ -29,7 +29,7 @@ Agent Loop Config :toctree: generated :nosignatures: - rl.agent_loop.single_turn_agent_loop.SingleTurnAgentLoopConfig + rl.agent_loop.SingleTurnAgentLoopConfig rl.agent_loop_manager.AgentLoopManagerConfig rl.agent_loop_manager.TaskSpecConfig rl.agent_loop_manager.SamplerConfig diff --git a/docs/zh_cn/api/rl_config.rst b/docs/zh_cn/api/rl_config.rst index 3aa97365b0..a36e097c0e 100644 --- a/docs/zh_cn/api/rl_config.rst +++ b/docs/zh_cn/api/rl_config.rst @@ -29,7 +29,7 @@ Agent Loop Config :toctree: generated :nosignatures: - rl.agent_loop.single_turn_agent_loop.SingleTurnAgentLoopConfig + rl.agent_loop.SingleTurnAgentLoopConfig rl.agent_loop_manager.AgentLoopManagerConfig rl.agent_loop_manager.TaskSpecConfig rl.agent_loop_manager.SamplerConfig diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index 1a82ecf127..dbac423c70 100644 --- a/xtuner/v1/rl/agent_loop/agent_loop.py +++ b/xtuner/v1/rl/agent_loop/agent_loop.py @@ -3,7 +3,7 @@ import asyncio from abc import ABC, abstractmethod from collections.abc import Callable -from typing import Any, TypeAlias, cast, overload +from typing import TYPE_CHECKING, Any, TypeAlias, cast, overload import ray from pydantic import BaseModel, ConfigDict @@ -11,10 +11,6 @@ from ray.util.placement_group import PlacementGroup from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status -from xtuner.v1.rl.distillation import ( - RolloutTeacherScorer, - RolloutTeacherScorerConfig, -) from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController from xtuner.v1.rl.rollout.constants import AGENT_LOOP_RAY_GENERATE_MAX_CONCURRENCY @@ -33,6 +29,9 @@ from xtuner.v1.utils.processing_utils import load_processor, load_tokenizer +if TYPE_CHECKING: + from xtuner.v1.rl.distillation import RolloutTeacherScorer, RolloutTeacherScorerConfig + AGENT_LOOP_CONCURRENCY_GROUP_GENERATE = "generate" IsValidSampleFn: TypeAlias = Callable[[list[RolloutState]], bool] @@ -77,6 +76,8 @@ def build( is_valid_sample_fn: IsValidSampleFn | None = None, rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None, ) -> AgentLoopSpec: + from xtuner.v1.rl.distillation import RolloutTeacherScorer + if self.cpu_resources is None: agent_loop = self.build_local( rollout_controller=rollout_controller, @@ -396,6 +397,8 @@ def __init__( *, rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None, ): + from xtuner.v1.rl.distillation import RolloutTeacherScorer + self.agent_loop = agent_loop_config.build_local( rollout_controller=rollout_controller, judger=judger, diff --git a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py index 831c5d0639..71b1f3caa8 100644 --- a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py @@ -1,15 +1,16 @@ +from __future__ import annotations + import asyncio import json import time from pathlib import Path -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast from pydantic import BaseModel, ConfigDict, Field from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast from xtuner.v1.data_proto.rl_data import Status from xtuner.v1.rl.agent_loop import AgentLoopConfig, IsValidSampleFn -from xtuner.v1.rl.distillation import RolloutTeacherScorerConfig from xtuner.v1.rl.judger import ComposedJudgerConfig, JudgerConfig, build_judger from xtuner.v1.rl.replay_buffer import ReplayBuffer from xtuner.v1.rl.rollout import RolloutController @@ -39,6 +40,10 @@ from .sampler import Sampler, SamplerConfig +if TYPE_CHECKING: + from xtuner.v1.rl.distillation import RolloutTeacherScorerConfig + + class TaskSpecConfig(BaseModel): """Configuration for one task managed by ``AgentLoopManager``. @@ -131,7 +136,7 @@ def build( logger=None, sync_weights_interval: int = 1, rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None, - ) -> "AgentLoopManager": + ) -> AgentLoopManager: tasks = self.tasks if isinstance(self.tasks, list) else [self.tasks] if not tasks: raise ValueError("AgentLoopManagerConfig requires at least one task config.") diff --git a/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py index 8d896a2950..70ae8731b3 100644 --- a/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py @@ -1,16 +1,17 @@ +from __future__ import annotations + import asyncio import json import time from enum import Enum, auto from pathlib import Path -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast from pydantic import BaseModel, ConfigDict, Field from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast from xtuner.v1.data_proto.rl_data import Status from xtuner.v1.rl.agent_loop import AgentLoopConfig, IsValidSampleFn -from xtuner.v1.rl.distillation import RolloutTeacherScorerConfig from xtuner.v1.rl.judger import ComposedJudgerConfig, JudgerConfig, build_judger from xtuner.v1.rl.replay_buffer import ReplayBuffer from xtuner.v1.rl.rollout import RolloutController @@ -42,6 +43,10 @@ from .sampler import Sampler, SamplerConfig +if TYPE_CHECKING: + from xtuner.v1.rl.distillation import RolloutTeacherScorerConfig + + class DisaggTaskSpecConfig(BaseModel): """单个非共卡 RL 数据源配置。""" @@ -71,7 +76,7 @@ def build( logger=None, sync_weights_interval: int = 1, rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None, - ) -> "DisaggAgentLoopManager": + ) -> DisaggAgentLoopManager: tasks = self.tasks if isinstance(self.tasks, list) else [self.tasks] if not tasks: raise ValueError("DisaggAgentLoopManagerConfig requires at least one task config.") From 8fa62bdf52f22c8713caaf61793195b4ca328fae Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Mon, 14 Sep 2026 08:02:42 +0000 Subject: [PATCH 20/20] fix ci --- tests/rl/test_on_policy_distillation.py | 6 +++-- tests/rl/test_prepare_train_data.py | 26 +++++++++---------- tests/rl/test_producer.py | 4 --- tests/rl/test_rl_disaggregated_trainer.py | 4 +-- tests/rl/test_trajectory_logging.py | 11 +++----- xtuner/v1/rl/agent_loop/agent_loop.py | 2 ++ .../rl/distillation/train_teacher_manager.py | 2 +- xtuner/v1/rl/{trainer => }/model_utils.py | 0 xtuner/v1/rl/trainer/worker.py | 2 +- xtuner/v1/train/rl_trainer.py | 13 +++++++--- 10 files changed, 35 insertions(+), 35 deletions(-) rename xtuner/v1/rl/{trainer => }/model_utils.py (100%) diff --git a/tests/rl/test_on_policy_distillation.py b/tests/rl/test_on_policy_distillation.py index f82e790a1c..f41755c8dd 100644 --- a/tests/rl/test_on_policy_distillation.py +++ b/tests/rl/test_on_policy_distillation.py @@ -211,7 +211,6 @@ def test_distillation_trainer_adapter_uses_null_defaults_without_config(self) -> self.assertIsNone(adapter.rollout_teacher_scorer_config) self.assertIsNone(adapter.train_teacher_manager_config) self.assertEqual(adapter.task_adv_weight, 1.0) - self.assertIsNone(adapter.teacher_index_by_data_source) self.assertEqual(adapter.timing_scalars([]), {}) def test_train_teacher_config_is_not_forwarded_to_agent_loop_manager(self) -> None: @@ -317,7 +316,10 @@ def test_distillation_loss_kwargs_carry_teacher_indices(self) -> None: ) assert loss_ctx is not None - torch.testing.assert_close(loss_ctx.loss_kwargs.teacher_indices, teacher_indices) + torch.testing.assert_close( + loss_ctx.loss_kwargs.teacher_indices, + teacher_indices.to(loss_ctx.loss_kwargs.teacher_indices.device), + ) def test_distillation_loss_kwargs_split_teacher_indices_with_ignore_padding(self) -> None: loss_kwargs = DistillationLossKwargs( diff --git a/tests/rl/test_prepare_train_data.py b/tests/rl/test_prepare_train_data.py index c382f3f5d7..f53618e919 100644 --- a/tests/rl/test_prepare_train_data.py +++ b/tests/rl/test_prepare_train_data.py @@ -127,8 +127,8 @@ def test_text_path_builds_shifted_training_tensors(self): batch["rollout_logprobs"], torch.tensor([[0.0, 0.0, 0.1, 0.2, 0.3]], dtype=torch.float32), ) - self.assertEqual(batch["advantage"], [0.0, 0.0, 1.5, 0.0, 1.5]) - self.assertEqual(len(batch["advantage"]), batch["shifted_labels"].numel()) + self.assertEqual(batch["advantage"], [1.5, 1.5, 1.5, 1.5, 0.0, 1.5]) + self.assertEqual(len(batch["advantage"]), batch["shifted_labels"].numel() + 1) self.assertIs(batch["seq_ctx"].rollout_routed_experts, routed_experts) self.assertEqual(info["training_samples"], 1) self.assertEqual(info["training_tokens"], 5) @@ -150,7 +150,7 @@ def test_rerolled_state_without_semantic_mask_uses_all_response_tokens(self): self.assertIsNone(state.response_mask) self.assertEqual(data_batches[0]["shifted_labels"].tolist(), [[-100, -100, 30, 31]]) - self.assertEqual(data_batches[0]["advantage"], [0.0, 0.0, 1.0, 1.0]) + self.assertEqual(data_batches[0]["advantage"], [1.0, 1.0, 1.0, 1.0, 1.0]) def test_multi_sample_group_uses_each_sample_reward_and_advantage(self): # 同一个 prompt 下的多个 response 要分别使用自己的 reward 和 advantage。 @@ -161,8 +161,8 @@ def test_multi_sample_group_uses_each_sample_reward_and_advantage(self): data_batches, info = self._prepare(trainer, [[first, second]]) self.assertEqual(len(data_batches), 2) - self.assertEqual(data_batches[0]["advantage"], [0.0, 0.0, 1.5, 1.5]) - self.assertEqual(data_batches[1]["advantage"], [0.0, 0.0, -2.0, -2.0]) + self.assertEqual(data_batches[0]["advantage"], [1.5, 1.5, 1.5, 1.5, 1.5]) + self.assertEqual(data_batches[1]["advantage"], [-2.0, -2.0, -2.0, -2.0, -2.0]) self.assertEqual(info["batch_size"], 2) self.assertEqual(info["rewards/min"], -1.0) self.assertEqual(info["rewards/max"], 3.0) @@ -294,8 +294,8 @@ def test_plain_topk_targets_include_masked_response_rows(self): prompt_ids=[10, 11, 12], response_ids=[20, 21, 22], response_mask=[0, 1, 1], - teacher_tokens=[[100, 101], [102, 103], [104, 105]], - teacher_logprobs=[[-0.5, -0.6], [-0.7, -0.8], [-0.9, -1.0]], + teacher_tokens=[[102, 103], [104, 105]], + teacher_logprobs=[[-0.7, -0.8], [-0.9, -1.0]], extra_fields={"origin_data_source": "agent_math"}, ) @@ -305,12 +305,12 @@ def test_plain_topk_targets_include_masked_response_rows(self): self.assertEqual(batch["shifted_labels"].tolist(), [[-100, -100, -100, 21, 22]]) self.assertEqual( batch["target_token_ids"].tolist(), - [[[0, 0], [0, 0], [100, 101], [102, 103], [104, 105]]], + [[[0, 0], [0, 0], [0, 0], [102, 103], [104, 105]]], ) torch.testing.assert_close( batch["teacher_logprobs"], torch.tensor( - [[[0.0, 0.0], [0.0, 0.0], [-0.5, -0.6], [-0.7, -0.8], [-0.9, -1.0]]], + [[[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [-0.7, -0.8], [-0.9, -1.0]]], dtype=torch.float32, ), ) @@ -329,8 +329,8 @@ def test_sampled_token_targets_align_for_plain_and_agentic_rollouts(self): prompt_ids=[10, 11, 12], response_ids=[20, 21, 22], response_mask=[0, 1, 1], - teacher_tokens=[20, 21, 22], - teacher_logprobs=[-0.5, -0.7, -0.9], + teacher_tokens=[21, 22], + teacher_logprobs=[-0.7, -0.9], extra_fields={"origin_data_source": "agent_math"}, ) agentic_state = self._state( @@ -349,7 +349,7 @@ def test_sampled_token_targets_align_for_plain_and_agentic_rollouts(self): self.assertEqual(len(data_batches), 2) torch.testing.assert_close( data_batches[0]["teacher_logprobs"], - torch.tensor([[0.0, 0.0, -0.5, -0.7, -0.9]], dtype=torch.float32), + torch.tensor([[0.0, 0.0, 0.0, -0.7, -0.9]], dtype=torch.float32), ) torch.testing.assert_close( data_batches[1]["teacher_logprobs"], @@ -375,7 +375,7 @@ def test_missing_reward_score_fails_fast(self): trainer = self._build_trainer([1.0]) state = self._state(reward={"other": 1.0}) - with self.assertRaises(AssertionError): + with self.assertRaisesRegex(ValueError, "missing.*score"): self._prepare(trainer, [[state]]) def test_logprobs_must_match_response_ids_length(self): diff --git a/tests/rl/test_producer.py b/tests/rl/test_producer.py index 6ac0c0fe69..364db9f09c 100644 --- a/tests/rl/test_producer.py +++ b/tests/rl/test_producer.py @@ -132,7 +132,6 @@ def _build_context( train_step: int = 0, model_step: int = 0, progress: ProduceProgress | None = None, - is_valid_sample_fn=None, ) -> ProduceContext: # 测试只走新的 ProduceContext 入口，不再覆盖旧散装参数兼容逻辑。 if progress is None: @@ -350,7 +349,6 @@ def is_valid_sample_fn(samples): self._build_agent_loop(is_valid_sample_fn=is_valid_sample_fn), self._build_sampler(), batch_size=1, - is_valid_sample_fn=is_valid_sample_fn, ) completed_group = [make_rollout_state(1, status=Status.COMPLETED)] @@ -435,7 +433,6 @@ def is_valid_sample_fn(samples): self._build_agent_loop(is_valid_sample_fn=is_valid_sample_fn), self._build_sampler(), batch_size=1, - is_valid_sample_fn=is_valid_sample_fn, ) completed_group = [ @@ -527,7 +524,6 @@ async def mock_gen(rs, **kwargs): train_step=4, model_step=3, progress=self._build_progress(task_name, target=2), - is_valid_sample_fn=is_valid_sample_fn, ) await strategy.produce_batch(ctx) diff --git a/tests/rl/test_rl_disaggregated_trainer.py b/tests/rl/test_rl_disaggregated_trainer.py index 914a5c86cf..afccee179f 100644 --- a/tests/rl/test_rl_disaggregated_trainer.py +++ b/tests/rl/test_rl_disaggregated_trainer.py @@ -519,8 +519,8 @@ def test_log_step_records_every_teacher_timing_and_distillation_metric(self): self.assertEqual(scalars["time/train_teacher/teacher_b/compute"], 3.0) self.assertEqual(scalars["time/train_teacher/total/compute"], 6.0) self.assertEqual(scalars["time/train_teacher_compute"], 6.0) - self.assertEqual(scalars["distillation/actor_distillation_kl"], 0.5) - self.assertEqual(scalars["distillation/sampled_opd_loss"], 0.25) + self.assertNotIn("distillation/actor_distillation_kl", scalars) + self.assertNotIn("distillation/sampled_opd_loss", scalars) self.assertNotIn("distillation/policy_loss", scalars) def test_update_weights_pauses_generation_without_onloading_rollout(self): diff --git a/tests/rl/test_trajectory_logging.py b/tests/rl/test_trajectory_logging.py index d0c92c37a1..14f7ebd194 100644 --- a/tests/rl/test_trajectory_logging.py +++ b/tests/rl/test_trajectory_logging.py @@ -19,6 +19,7 @@ "lagent.utils.rate_limiter": MagicMock(), }, ): + from xtuner.v1.rl.agent_loop.sandbox_agent_loop import agent_in_sandbox_loop as sandbox_impl from xtuner.v1.rl.agent_loop.sandbox_agent_loop.agent_in_sandbox_loop import AgentInSandboxLoop from xtuner.v1.rl.agent_loop.sandbox_agent_loop.schemas import AgentRolloutItem, RolloutStatus @@ -93,10 +94,7 @@ async def test_completed_session_assigns_artifacts_only_to_first_segment(self): original_artifacts = copy.deepcopy(item.artifacts) trace_store = self._trace_store(segment_count=2) - with patch( - "xtuner.v1.rl.agent_loop.sandbox_agent_loop.agent_in_sandbox_loop.get_store", - return_value=trace_store, - ): + with patch.object(sandbox_impl, "get_store", return_value=trace_store): segments = await loop._build_rollout_states(state, item) self.assertEqual(len(segments), 2) @@ -124,10 +122,7 @@ async def test_single_segment_keeps_artifacts(self): state = self._make_state() item = self._make_item(RolloutStatus.COMPLETED, segment_count=1) - with patch( - "xtuner.v1.rl.agent_loop.sandbox_agent_loop.agent_in_sandbox_loop.get_store", - return_value=self._trace_store(segment_count=1), - ): + with patch.object(sandbox_impl, "get_store", return_value=self._trace_store(segment_count=1)): segments = await loop._build_rollout_states(state, item) self.assertEqual(len(segments), 1) diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index dbac423c70..823942d05a 100644 --- a/xtuner/v1/rl/agent_loop/agent_loop.py +++ b/xtuner/v1/rl/agent_loop/agent_loop.py @@ -235,6 +235,8 @@ def __init__( logger=None, enable_batch_judge: bool = False, ) -> None: + from xtuner.v1.rl.distillation import RolloutTeacherScorer + self.rollout_ctl = rollout_ctl self.hf_checkpoint = hf_checkpoint self.tokenizer = load_tokenizer(hf_checkpoint, trust_remote_code=True) diff --git a/xtuner/v1/rl/distillation/train_teacher_manager.py b/xtuner/v1/rl/distillation/train_teacher_manager.py index 1279d67e34..a4ac49d967 100644 --- a/xtuner/v1/rl/distillation/train_teacher_manager.py +++ b/xtuner/v1/rl/distillation/train_teacher_manager.py @@ -10,7 +10,7 @@ from xtuner.v1.data_proto.sequence_context import SequenceContext from xtuner.v1.loss import LogProbConfig, LogProbContext, TopKLogProbConfig from xtuner.v1.model.compose.base import BaseComposeConfig -from xtuner.v1.rl.trainer.model_utils import FrozenModel, build_frozen_model +from xtuner.v1.rl.model_utils import FrozenModel, build_frozen_model from xtuner.v1.utils import get_device, get_logger, get_torch_device_module from .config import TrainTeacherManagerConfig diff --git a/xtuner/v1/rl/trainer/model_utils.py b/xtuner/v1/rl/model_utils.py similarity index 100% rename from xtuner/v1/rl/trainer/model_utils.py rename to xtuner/v1/rl/model_utils.py diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index 47a74fee04..39242f432d 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -52,7 +52,7 @@ BaseRLLossContext, kl_penalty, ) -from xtuner.v1.rl.trainer.model_utils import build_frozen_model +from xtuner.v1.rl.model_utils import build_frozen_model from xtuner.v1.rl.utils import SingleAcceleratorWorker from xtuner.v1.rl.weight_update import WeightUpdater from xtuner.v1.train.trainer import LoadCheckpointConfig diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index c59d842fc2..466eb4970f 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -1185,8 +1185,13 @@ def _prepare_train_data( if isinstance(turns, int): tool_turns_list.append(turns) if data.reward is None or "score" not in data.reward: - assert task_adv_weight == 0, f"Reward is missing or does not contain 'score' key in data: {data}" - continue + if task_adv_weight > 0: + raise ValueError( + f"Reward is missing or does not contain 'score' key in data: {data}, " + f"but task_adv_weight={task_adv_weight} > 0" + ) + else: + continue reward = float(data.reward["score"]) # session_id is only set by agentic loops / XTUNER_DETERMINISTIC; plain RL falls back # to rollout_id, which the sampler always assigns. Segments of one session share a key. @@ -1205,8 +1210,6 @@ def _prepare_train_data( if task_adv_weight == 0: sample_advantages = [0.0] * len(group) else: - # Agentic rollouts may split one model session into multiple trainable segments. - # Compute the group advantage once per session, then broadcast it back to each segment. rewards_tensor = torch.tensor(cluster_rewards, dtype=torch.float32) cluster_advantages = self._advantage_estimator.compute(rewards_tensor, cluster_representatives) sample_advantages = [ @@ -1328,6 +1331,7 @@ def _prepare_train_data( shifted_labels = [-100] * (len(prompt_ids) - 1) + response_labels shifted_labels_t = torch.tensor(shifted_labels, dtype=torch.int64).unsqueeze(0) + # Keep the advantage layout aligned with input_ids (response excludes EOS). advatnages_val = sample_advantages[i] actual_advantages = [advatnages_val] * len(prompt_ids) + [ 0.0 if mask == 0 else advatnages_val for mask in response_mask @@ -1533,6 +1537,7 @@ def _log_step( for key, value in mini_batch_metrics.items(): avg_value = sum(value) / len(value) all_scalars.update({f"train_metrics/worker_{worker_idx}/step_avg_{key}": avg_value}) + rank_sft_log = log_item["sft_train_metrics"] for k, v in rank_sft_log.items(): all_scalars.update({f"sft_train_metrics/worker_{worker_idx}/{k}": v})