From c294d9e80e790ccf118195d7b5ef9935f35d31d1 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 27 Aug 2026 22:50:30 +0800 Subject: [PATCH 1/3] chore: remove megatron processor service --- cookbook/client/server/megatron/run.sh | 101 ++++++++---------- .../client/server/megatron/server_config.yaml | 25 ----- 2 files changed, 47 insertions(+), 79 deletions(-) diff --git a/cookbook/client/server/megatron/run.sh b/cookbook/client/server/megatron/run.sh index 3c017b93d..34abf52bd 100644 --- a/cookbook/client/server/megatron/run.sh +++ b/cookbook/client/server/megatron/run.sh @@ -3,7 +3,8 @@ # ============================================ # Twinkle Megatron 服务启动脚本 # ============================================ -# 功能：启动 Ray 集群（支持多 GPU/CPU 节点）、LGTM 观测栈和 Twinkle 服务器 +# 功能：启动 Ray 集群（支持 GPU 节点）、LGTM 观测栈和 Twinkle 服务器。 +# 设置 TWINKLE_DASHSERVING_ADAPTER=1 时，同时启动 ASI Native HTTP Adapter。 # # 用法：./run.sh [选项] # @@ -11,15 +12,15 @@ # --restart 如果已有 run.sh 实例正在运行，请求其退出并由 entrypoint 重启服务 # --head NODE Head 节点 GPU 设备列表，逗号分隔 (默认: 0,1,2,3) # --gpu-workers LIST GPU Worker 列表，分号分隔多个节点 (默认: 4,5,6,7) -# --cpu-workers N CPU Worker 数量 (默认: 1) -# --temp-dir DIR Ray 临时目录 (默认: /dashscope/caches/application/ray_logs) -# --save-dir DIR Twinkle 模型保存目录 (默认: /dashscope/caches/application/save) +# --temp-dir DIR Ray 临时目录 (默认: /twinkle/runtime/ray_logs) +# --save-dir DIR Twinkle 模型保存目录 (默认: /twinkle/runtime/save) # --server-config FILE Twinkle 服务器配置文件路径 (默认: /twinkle/cookbook/client/server/megatron/server_config.yaml) # --help 显示帮助信息 # # 环境变量： -# MODELSCOPE_CACHE 默认 /dashscope/caches/application/.cache -# TWINKLE_WORK_DIR 默认 /dashscope/caches/application/twinkle +# MODELSCOPE_CACHE 默认 /twinkle/runtime/.cache +# TWINKLE_WORK_DIR 默认 /twinkle/runtime/work +# TWINKLE_DASHSERVING_ADAPTER 设为 1 时启动端口 9000 的 Adapter # TWINKLE_RUN_EXISTING_ACTION 已有 run.sh 进程运行时的行为：exit 或 restart（默认 exit） # TWINKLE_RUN_RESTART_TIMEOUT_SECONDS --restart 等待已有实例接收请求秒数（默认 120） # @@ -30,9 +31,6 @@ # # 直接启动服务，前台等待 server 子进程 # bash /twinkle/cookbook/client/server/megatron/run.sh --restart # # 更新代码后请求已有 run.sh 退出并由 entrypoint 重启 -# ./run.sh --head "0,1,2,3" --gpu-workers "4,5,6,7" --cpu-workers 1 -# ./run.sh --head "0,1,2,3" --gpu-workers "" --cpu-workers 0 -# ./run.sh --head "" --cpu-workers 4 # 纯 CPU 模式 # ./run.sh --temp-dir /tmp/my_ray_logs # 自定义临时目录 # ============================================ @@ -53,20 +51,18 @@ set -e # 遇到错误立即退出 # 示例："4,5,6,7" 或 "4,5,6,7;8,9,10,11" # 可通过命令行参数 $2 传入 -# CPU Worker 数量 -# 可通过命令行参数 $3 传入 - # --- 网络配置 --- RAY_PORT=6379 RAY_ADDRESS="127.0.0.1:$RAY_PORT" # --- 路径配置 --- -export MODELSCOPE_CACHE="${MODELSCOPE_CACHE:-/dashscope/caches/application/.cache}" -TWINKLE_WORK_DIR="${TWINKLE_WORK_DIR:-/dashscope/caches/application/twinkle}" -DEFAULT_TEMP_DIR="/dashscope/caches/application/ray_logs" +TWINKLE_RUNTIME_DIR="${TWINKLE_RUNTIME_DIR:-/twinkle/runtime}" +export MODELSCOPE_CACHE="${MODELSCOPE_CACHE:-$TWINKLE_RUNTIME_DIR/.cache}" +TWINKLE_WORK_DIR="${TWINKLE_WORK_DIR:-$TWINKLE_RUNTIME_DIR/work}" +DEFAULT_TEMP_DIR="${TWINKLE_TEMP_DIR:-$TWINKLE_RUNTIME_DIR/ray_logs}" LOG_FILE="run.log" REDIS_LOG_FILE="/twinkle/redis.log" -DEFAULT_SAVE_DIR="/dashscope/caches/application/save" +DEFAULT_SAVE_DIR="${TWINKLE_SAVE_DIR:-$TWINKLE_RUNTIME_DIR/save}" DEFAULT_SERVER_CONFIG_FILE="/twinkle/cookbook/client/server/megatron/server_config.yaml" # --- LGTM 版本配置（与 grafana/otel-lgtm:0.28.0 保持一致） --- @@ -87,8 +83,10 @@ TWINKLE_RUN_RESTART_REQUEST_FILE="${TWINKLE_RUN_RESTART_REQUEST_FILE:-/tmp/twink TWINKLE_RUN_EXISTING_ACTION="${TWINKLE_RUN_EXISTING_ACTION:-exit}" TWINKLE_RUN_RESTART_TIMEOUT_SECONDS="${TWINKLE_RUN_RESTART_TIMEOUT_SECONDS:-120}" SERVER_PID="" +ADAPTER_PID="" TAIL_PID="" RESTART_REQUESTED_BY_SIGNAL=0 +TWINKLE_DASHSERVING_ADAPTER="${TWINKLE_DASHSERVING_ADAPTER:-0}" # ============================================ # 参数解析（支持 --key=value 或 --key value 格式） @@ -97,7 +95,6 @@ RESTART_REQUESTED_BY_SIGNAL=0 # 默认值 HEAD_NODE="0,1,2,3" GPU_WORKERS_INPUT="4,5,6,7" -CPU_WORKER_COUNT="1" TEMP_DIR="$DEFAULT_TEMP_DIR" SAVE_DIR="$DEFAULT_SAVE_DIR" SERVER_CONFIG_FILE="$DEFAULT_SERVER_CONFIG_FILE" @@ -110,15 +107,15 @@ print_usage() { --restart 如果已有 run.sh 实例正在运行，请求其退出并由 entrypoint 重启服务 --head NODE Head 节点 GPU 设备列表，逗号分隔 (默认: 0,1,2,3) --gpu-workers LIST GPU Worker 列表，分号分隔多个节点 (默认: 4,5,6,7) - --cpu-workers N CPU Worker 数量 (默认: 1) --temp-dir DIR Ray 临时目录 --save-dir DIR Twinkle 模型保存目录 (默认: $DEFAULT_SAVE_DIR) --server-config FILE Twinkle 服务器配置文件路径 (默认: $DEFAULT_SERVER_CONFIG_FILE) --help, -h 显示帮助信息 环境变量: - MODELSCOPE_CACHE 默认: /dashscope/caches/application/.cache - TWINKLE_WORK_DIR 默认: /dashscope/caches/application/twinkle + MODELSCOPE_CACHE 默认: /twinkle/runtime/.cache + TWINKLE_WORK_DIR 默认: /twinkle/runtime/work + TWINKLE_DASHSERVING_ADAPTER 设为 1 时启动端口 9000 的 Adapter TWINKLE_RUN_EXISTING_ACTION 已有 run.sh 进程运行时的行为：exit 或 restart (默认: exit) TWINKLE_RUN_RESTART_TIMEOUT_SECONDS --restart 等待已有实例接收请求秒数 (默认: 120) @@ -134,7 +131,6 @@ print_usage() { ./run.sh --head '0,1,2,3' --gpu-workers '4,5,6,7' ./run.sh --head '0,1,2,3,4,5,6,7' # 单机 8 卡 ./run.sh --gpu-workers '4,5,6,7;8,9,10,11' # 多 GPU Worker - ./run.sh --cpu-workers 4 --head '' # 纯 CPU 模式 EOF } @@ -161,14 +157,6 @@ while [[ $# -gt 0 ]]; do GPU_WORKERS_INPUT="${1#*=}" shift ;; - --cpu-workers) - CPU_WORKER_COUNT="$2" - shift 2 - ;; - --cpu-workers=*) - CPU_WORKER_COUNT="${1#*=}" - shift - ;; --temp-dir) TEMP_DIR="$2" shift 2 @@ -367,15 +355,6 @@ cleanup_pid_file() { fi } -require_non_negative_int() { - local name="$1" - local value="$2" - if ! [[ "$value" =~ ^[0-9]+$ ]]; then - print_error "$name 必须是非负整数，当前值: $value" - exit 1 - fi -} - require_positive_int() { local name="$1" local value="$2" @@ -396,7 +375,13 @@ validate_runtime_config() { esac require_positive_int "TWINKLE_RUN_RESTART_TIMEOUT_SECONDS" "$TWINKLE_RUN_RESTART_TIMEOUT_SECONDS" - require_non_negative_int "CPU_WORKER_COUNT" "$CPU_WORKER_COUNT" + case "$TWINKLE_DASHSERVING_ADAPTER" in + 0|1) ;; + *) + print_error "TWINKLE_DASHSERVING_ADAPTER 只能是 0 或 1，当前值: $TWINKLE_DASHSERVING_ADAPTER" + exit 1 + ;; + esac } require_command() { @@ -493,8 +478,10 @@ stop_pid() { cleanup_existing_runtime() { stop_pid "$TAIL_PID" "日志 tail" + stop_pid "$ADAPTER_PID" "DashServing Adapter" stop_pid "$SERVER_PID" "Twinkle Server" TAIL_PID="" + ADAPTER_PID="" SERVER_PID="" print_info "停止已有的 Twinkle Server..." @@ -609,11 +596,6 @@ print_runtime_config() { done fi - if [ "$CPU_WORKER_COUNT" -gt 0 ]; then - echo "" - echo " [CPU Worker 节点] $CPU_WORKER_COUNT 个" - fi - echo "" print_info "运行参数：" echo " - Ray 地址: $RAY_ADDRESS" @@ -683,16 +665,6 @@ start_ray_cluster() { print_success "GPU Worker $((i+1)) 启动成功！" done - if [ "$CPU_WORKER_COUNT" -gt 0 ]; then - print_info "启动 $CPU_WORKER_COUNT 个 CPU Worker..." - for ((i=1; i<=CPU_WORKER_COUNT; i++)); do - CUDA_VISIBLE_DEVICES="" ray start \ - --address=$RAY_ADDRESS \ - --num-gpus=0 - done - print_success "CPU Worker 启动成功！" - fi - echo "" print_info "集群状态：" ray status 2>/dev/null || true @@ -742,6 +714,20 @@ start_twinkle_server() { start_log_tail } +start_dashserving_adapter() { + if [ "$TWINKLE_DASHSERVING_ADAPTER" != "1" ]; then + return + fi + + print_header "启动 DashServing Adapter" + export PORT="${PORT:-9000}" + export TWINKLE_INTERNAL_URL="${TWINKLE_INTERNAL_URL:-http://127.0.0.1:8000}" + export TWINKLE_DS_TIMEOUT_SECONDS="${TWINKLE_DS_TIMEOUT_SECONDS:-600}" + nohup python -m twinkle.server.dashserving >> "$LOG_FILE" 2>&1 & + ADAPTER_PID=$! + print_success "DashServing Adapter 已启动 (PID: $ADAPTER_PID, port: $PORT)" +} + wait_runtime() { print_info "Twinkle runtime 已启动，等待 server 进程..." while true; do @@ -754,6 +740,12 @@ wait_runtime() { return 1 fi + if [ "$TWINKLE_DASHSERVING_ADAPTER" = "1" ] \ + && ! kill -0 "$ADAPTER_PID" 2>/dev/null; then + print_error "DashServing Adapter 进程已退出 (PID: $ADAPTER_PID)" + return 1 + fi + if ! kill -0 "$TAIL_PID" 2>/dev/null; then print_warning "日志 tail 进程已退出，重新启动..." start_log_tail @@ -772,6 +764,7 @@ run_service_once() { start_ray_cluster start_lgtm start_twinkle_server + start_dashserving_adapter wait_runtime } diff --git a/cookbook/client/server/megatron/server_config.yaml b/cookbook/client/server/megatron/server_config.yaml index 759601985..24132a07e 100644 --- a/cookbook/client/server/megatron/server_config.yaml +++ b/cookbook/client/server/megatron/server_config.yaml @@ -140,28 +140,3 @@ applications: TWINKLE_TRUST_REMOTE_CODE: "0" TWINKLE_LONG_POLL_TIMEOUT: "120" TWINKLE_FAIL_FAST: "0" - - # 4. Processor Service - - name: processor - route_prefix: /api/v1/processor - import_path: processor - args: - ncpu_proc_per_node: 2 - device_group: - name: model - ranks: 2 - device_type: CPU - device_mesh: - device_type: CPU - dp_size: 2 - deployments: - - name: ProcessorManagement - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 128 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_FAIL_FAST: "0" From 4893c6c2d9f0b853ca999fbefc3faed4d3fed23c Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 28 Aug 2026 10:39:01 +0800 Subject: [PATCH 2/3] fix: store lora optimizer checkpoints safely --- .../model/megatron/multi_lora_megatron.py | 7 +- src/twinkle/utils/safetensors.py | 72 ++++++++++++++++++- tests/utils/test_safetensors_state_dict.py | 38 ++++++++++ 3 files changed, 113 insertions(+), 4 deletions(-) create mode 100644 tests/utils/test_safetensors_state_dict.py diff --git a/src/twinkle/model/megatron/multi_lora_megatron.py b/src/twinkle/model/megatron/multi_lora_megatron.py index 3bd197158..35c971c7e 100644 --- a/src/twinkle/model/megatron/multi_lora_megatron.py +++ b/src/twinkle/model/megatron/multi_lora_megatron.py @@ -23,6 +23,7 @@ from twinkle.metric import Metric from twinkle.processor import InputProcessor from twinkle.utils import get_logger +from twinkle.utils.safetensors import load_state_dict, save_state_dict from ..multi_lora import MultiLora from ._mindspeed_runtime import ensure_mindspeed_adaptor_patched from .megatron import MegatronModel @@ -213,7 +214,7 @@ def set_lr_scheduler(self, scheduler_cls: Union[LRScheduler, Type[LRScheduler], @staticmethod def _rank_local_optimizer_path(checkpoint_dir: str) -> str: rank = dist.get_rank() if dist.is_initialized() else 0 - return os.path.join(checkpoint_dir, f'optimizer_rank_{rank}.pt') + return os.path.join(checkpoint_dir, f'optimizer_rank_{rank}.safetensors') @staticmethod def _save_local_training_rng_state(): @@ -257,7 +258,7 @@ def _save_multi_lora_optimizer(self, checkpoint_dir: str, optimizer_config, **kw if optimizer_config.lr_scheduler is not None: state_dict['opt_param_scheduler'] = optimizer_config.lr_scheduler.state_dict() - torch.save(state_dict, self._rank_local_optimizer_path(checkpoint_dir)) + save_state_dict(state_dict, self._rank_local_optimizer_path(checkpoint_dir)) if dist.is_initialized(): dist.barrier() @@ -266,7 +267,7 @@ def _load_multi_lora_optimizer(self, checkpoint_dir: str, adapter_name: str = '' no_load_optim = kwargs.pop('no_load_optim', False) no_load_rng = kwargs.pop('no_load_rng', True) optimizer_config = self.optimizer_group.get(adapter_name) - state_dict = torch.load(self._rank_local_optimizer_path(checkpoint_dir), map_location='cpu', weights_only=False) + state_dict = load_state_dict(self._rank_local_optimizer_path(checkpoint_dir)) if not no_load_optim and optimizer_config is not None: if optimizer_config.optimizer is not None and 'optimizer' in state_dict: diff --git a/src/twinkle/utils/safetensors.py b/src/twinkle/utils/safetensors.py index 42619c54d..e7df6c2fe 100644 --- a/src/twinkle/utils/safetensors.py +++ b/src/twinkle/utils/safetensors.py @@ -1,7 +1,9 @@ import json +import numpy as np import os +import torch from functools import partial -from typing import Literal +from typing import Any, Literal from .device_mesh import is_last_rank, is_master @@ -19,6 +21,74 @@ def load(self): return self.tensor +_STATE_TYPE = '__twinkle_state_type__' + + +def save_state_dict(state: Any, path: str) -> None: + """Safely persist a nested checkpoint state without pickle serialization.""" + from safetensors.torch import save_file + + tensors = {} + manifest = _encode_state(state, tensors) + save_file(tensors, path) + with open(f'{path}.json', 'w') as f: + json.dump(manifest, f, separators=(',', ':')) + + +def load_state_dict(path: str) -> Any: + """Load a checkpoint written by :func:`save_state_dict`.""" + from safetensors.torch import load_file + + with open(f'{path}.json') as f: + manifest = json.load(f) + return _decode_state(manifest, load_file(path, device='cpu')) + + +def _encode_state(value: Any, tensors: dict[str, torch.Tensor]) -> Any: + if isinstance(value, torch.Tensor): + name = f'tensor_{len(tensors)}' + tensors[name] = value.detach().cpu().contiguous().clone() + return {_STATE_TYPE: 'tensor', 'name': name} + if isinstance(value, np.ndarray): + name = f'tensor_{len(tensors)}' + tensors[name] = torch.from_numpy(value).contiguous().clone() + return {_STATE_TYPE: 'ndarray', 'name': name, 'dtype': value.dtype.str} + if isinstance(value, np.generic): + return {_STATE_TYPE: 'numpy_scalar', 'dtype': value.dtype.str, 'value': value.item()} + if isinstance(value, dict): + return { + _STATE_TYPE: 'dict', + 'items': [[_encode_state(k, tensors), _encode_state(v, tensors)] for k, v in value.items()] + } + if isinstance(value, tuple): + return {_STATE_TYPE: 'tuple', 'items': [_encode_state(v, tensors) for v in value]} + if isinstance(value, list): + return {_STATE_TYPE: 'list', 'items': [_encode_state(v, tensors) for v in value]} + if value is None or isinstance(value, (bool, int, float, str)): + return value + raise TypeError(f'Unsupported checkpoint value type: {type(value).__name__}') + + +def _decode_state(value: Any, tensors: dict[str, torch.Tensor]) -> Any: + if not isinstance(value, dict) or _STATE_TYPE not in value: + return value + + value_type = value[_STATE_TYPE] + if value_type == 'tensor': + return tensors[value['name']] + if value_type == 'ndarray': + return tensors[value['name']].numpy().astype(np.dtype(value['dtype']), copy=False) + if value_type == 'numpy_scalar': + return np.dtype(value['dtype']).type(value['value']) + if value_type == 'dict': + return {_decode_state(k, tensors): _decode_state(v, tensors) for k, v in value['items']} + if value_type == 'tuple': + return tuple(_decode_state(v, tensors) for v in value['items']) + if value_type == 'list': + return [_decode_state(v, tensors) for v in value['items']] + raise ValueError(f'Unknown checkpoint value type: {value_type}') + + class SafetensorLazyLoader: def __init__(self, hf_model_dir: str, is_peft_format: bool = False): diff --git a/tests/utils/test_safetensors_state_dict.py b/tests/utils/test_safetensors_state_dict.py new file mode 100644 index 000000000..39c345ad4 --- /dev/null +++ b/tests/utils/test_safetensors_state_dict.py @@ -0,0 +1,38 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import random + +import numpy as np +import torch + +from twinkle.utils.safetensors import load_state_dict, save_state_dict + + +def test_state_dict_safetensors_round_trip(tmp_path, monkeypatch): + state = { + 'optimizer': { + 'state': {0: {'exp_avg': torch.tensor([1.0, 2.0]), 'step': torch.tensor(3)}}, + 'param_groups': [{'lr': 1e-4, 'betas': (0.9, 0.999)}], + }, + 'rng_state': { + 'random_rng_state': random.getstate(), + 'np_rng_state': np.random.get_state(), + 'torch_rng_state': torch.tensor([1, 2], dtype=torch.uint8), + }, + 'iteration': 3, + } + + path = tmp_path / 'optimizer_rank_0.safetensors' + save_state_dict(state, str(path)) + + def fail_if_called(*args, **kwargs): + raise AssertionError('torch.load must not be used for optimizer checkpoints') + + monkeypatch.setattr(torch, 'load', fail_if_called) + loaded = load_state_dict(str(path)) + + assert path.exists() + assert (tmp_path / 'optimizer_rank_0.safetensors.json').exists() + assert torch.equal(loaded['optimizer']['state'][0]['exp_avg'], state['optimizer']['state'][0]['exp_avg']) + assert loaded['optimizer']['param_groups'] == state['optimizer']['param_groups'] + assert loaded['rng_state']['random_rng_state'] == state['rng_state']['random_rng_state'] + assert np.array_equal(loaded['rng_state']['np_rng_state'][1], state['rng_state']['np_rng_state'][1]) From 59a769f5779e88c870565cc150fae78165917e2e Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 28 Aug 2026 15:58:22 +0800 Subject: [PATCH 3/3] fix: restore checkpoint and DPO compatibility --- .../model/megatron/multi_lora_megatron.py | 2 +- src/twinkle/server/model/backends/common.py | 25 +++++++++++++++ .../server/model/backends/megatron_model.py | 15 +-------- .../model/backends/transformers_model.py | 1 + .../server/model/test_tinker_compat_output.py | 31 +++++++++++++++++++ 5 files changed, 59 insertions(+), 15 deletions(-) diff --git a/src/twinkle/model/megatron/multi_lora_megatron.py b/src/twinkle/model/megatron/multi_lora_megatron.py index 35c971c7e..ebda91501 100644 --- a/src/twinkle/model/megatron/multi_lora_megatron.py +++ b/src/twinkle/model/megatron/multi_lora_megatron.py @@ -273,7 +273,7 @@ def _load_multi_lora_optimizer(self, checkpoint_dir: str, adapter_name: str = '' if optimizer_config.optimizer is not None and 'optimizer' in state_dict: optimizer_config.optimizer.load_state_dict(state_dict['optimizer']) device = Platform.get_local_device() - for group_state in optimizer_config.optimizer.state.values(): + for _, group_state in optimizer_config.optimizer.state.items(): if not isinstance(group_state, dict): continue for k, v in group_state.items(): diff --git a/src/twinkle/server/model/backends/common.py b/src/twinkle/server/model/backends/common.py index 7c90a4ed0..81f3146cc 100644 --- a/src/twinkle/server/model/backends/common.py +++ b/src/twinkle/server/model/backends/common.py @@ -125,6 +125,31 @@ def _to_float(v): class TwinkleCompatModelBase: """Base class containing common logic for Twinkle compatibility wrappers.""" + @staticmethod + def _normalize_ref_outputs(kwargs: dict) -> None: + """Convert HTTP/Ray-serialized reference logps to one padded tensor.""" + ref_outputs = kwargs.get('ref_outputs') + if not isinstance(ref_outputs, dict): + return + + logps = ref_outputs.get('logps') + if not isinstance(logps, (list, tuple)) or not logps or isinstance(logps[0], torch.Tensor): + return + + rows = [] + for item in logps: + if isinstance(item, (list, tuple)) and item and isinstance(item[0], (list, tuple)): + rows.extend(item) + else: + rows.append(item) + + from twinkle.utils import pad_and_stack_tensors + ref_outputs['logps'] = pad_and_stack_tensors( + [torch.as_tensor(row, dtype=torch.float32) for row in rows], + pad_value=0.0, + concat=False, + ) + def get_template(self, adapter_name: str) -> Template: return self.optimizer_group[adapter_name].template diff --git a/src/twinkle/server/model/backends/megatron_model.py b/src/twinkle/server/model/backends/megatron_model.py index f945f1efc..56f9acc98 100644 --- a/src/twinkle/server/model/backends/megatron_model.py +++ b/src/twinkle/server/model/backends/megatron_model.py @@ -131,20 +131,7 @@ def forward_only(self, *, inputs: InputFeature | list[InputFeature] | Trajectory @nccl_safe_megatron def forward_backward(self, *, inputs: InputFeature | list[InputFeature] | Trajectory | list[Trajectory], **kwargs): """Forward+backward for twinkle-native clients (InputFeature/Trajectory I/O).""" - # Normalize ragged ref_outputs logps into a regular 2D tensor. - # After HTTP + collect_tensor_dict, logps is a nested list grouped - # by microbatch with varying seq_lens across DP ranks. Flatten to - # per-sample 1D lists and pad_and_stack — same as datum.py L84-88. - ref_outputs = kwargs.get('ref_outputs') - if isinstance(ref_outputs, dict) and 'logps' in ref_outputs: - logps = ref_outputs['logps'] - if isinstance(logps, (list, tuple)) and logps and not isinstance(logps[0], torch.Tensor): - # Flatten [[mb0_sample0, mb0_sample1], [mb1_sample0, ...]] → [sample0, sample1, ...] - flat = [s for item in logps for s in (item if isinstance(item[0], (list, tuple)) else [item])] - from twinkle.utils import pad_and_stack_tensors - ref_outputs['logps'] = pad_and_stack_tensors([torch.tensor(s, dtype=torch.float32) for s in flat], - pad_value=0.0, - concat=False) + self._normalize_ref_outputs(kwargs) output = super().forward_backward(inputs=inputs, **kwargs) return to_cpu_safe_output(output) diff --git a/src/twinkle/server/model/backends/transformers_model.py b/src/twinkle/server/model/backends/transformers_model.py index 1a382805b..8dc503bb0 100644 --- a/src/twinkle/server/model/backends/transformers_model.py +++ b/src/twinkle/server/model/backends/transformers_model.py @@ -110,6 +110,7 @@ def forward_only(self, *, inputs: InputFeature | list[InputFeature] | Trajectory @nccl_safe def forward_backward(self, *, inputs: InputFeature | list[InputFeature] | Trajectory | list[Trajectory], **kwargs): """Forward+backward for twinkle-native clients (InputFeature/Trajectory I/O).""" + self._normalize_ref_outputs(kwargs) output = super().forward_backward(inputs=inputs, **kwargs) return to_cpu_safe_output(output) diff --git a/tests/server/model/test_tinker_compat_output.py b/tests/server/model/test_tinker_compat_output.py index cb863b49f..2041aefa9 100644 --- a/tests/server/model/test_tinker_compat_output.py +++ b/tests/server/model/test_tinker_compat_output.py @@ -149,3 +149,34 @@ def test_extract_rl_features_keeps_grpo_logps_and_advantages_ragged(): assert result['old_logps'] == [[1.0, 2.0, 3.0], [4.0, 5.0]] assert result['advantages'] == [[0.5, 0.5, 0.5], [-0.5, -0.5]] + + +def test_normalize_ref_outputs_flattens_and_pads_serialized_logps(): + kwargs = { + 'ref_outputs': { + 'logps': [ + [[1.0, 2.0, 3.0], [4.0, 5.0]], + [[6.0, 7.0]], + ] + } + } + + TwinkleCompatModelBase._normalize_ref_outputs(kwargs) + + torch.testing.assert_close( + kwargs['ref_outputs']['logps'], + torch.tensor([ + [1.0, 2.0, 3.0], + [4.0, 5.0, 0.0], + [6.0, 7.0, 0.0], + ]), + ) + + +def test_normalize_ref_outputs_keeps_tensor_unchanged(): + logps = torch.tensor([[1.0, 2.0]]) + kwargs = {'ref_outputs': {'logps': logps}} + + TwinkleCompatModelBase._normalize_ref_outputs(kwargs) + + assert kwargs['ref_outputs']['logps'] is logps