diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 00000000..ee1dfce7 --- /dev/null +++ b/THIRD_PARTY_NOTICES.md @@ -0,0 +1,40 @@ +# Third-Party Notices + +This repository includes or adapts code from the projects listed below. Source +files that already carry upstream copyright and license headers keep those +headers in place. + +## Apache-2.0 Components + +The full Apache License 2.0 text is provided in `LICENSE`. + +- `scripts/normalize.py` adapts code or structure from + `Physical-Intelligence/openpi`, licensed under Apache-2.0. +- `wall_x/model/qact/qwen2_5/configuration_qwen2_5_vl.py` and + `wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py` include Qwen / + Hugging Face derived code, licensed under Apache-2.0. + +## MIT Components + +LIBERO evaluation helpers in `wall_x/_vendor/harrix/envs/libero*.py` interface +with LIBERO components, licensed under the MIT License: + +Copyright (c) 2023 Lifelong Robot Learning + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/requirements.txt b/requirements.txt index b7e53678..4a214661 100644 --- a/requirements.txt +++ b/requirements.txt @@ -24,3 +24,8 @@ tqdm>=4.67,<5 tyro>=0.9.0,<1.0 wandb>=0.23,<0.25 matplotlib>=3.10,<4 +imageio>=2.34,<3 +imageio-ffmpeg>=0.5,<0.7 +# DMuon optimizer (recommended for LIBERO finetuning; see the *_dmuon.yml +# configs). Public, installed from git (not published on PyPI). +dmuon @ git+https://github.com/X-Square-Robot/dmuon.git diff --git a/scripts/draw_openloop_plot.py b/scripts/draw_openloop_plot.py index cbff1ef9..7128a96c 100644 --- a/scripts/draw_openloop_plot.py +++ b/scripts/draw_openloop_plot.py @@ -7,16 +7,13 @@ Reference: ``infer_openloop_websocket.py`` (websocket client + open-loop loop). -For single-arm LIBERO checkpoints, start the server with raw model output:: - - python -m wall_x._vendor.harrix.serving.launch_serving \\ - --env X2ROBOT --port 32194 \\ - --no-serialize-actions \\ - model-config:server-model-config \\ - --model-config.checkpoint-path /path/to/ckpt \\ - --model-config.train-config-path /path/to/libero.yml \\ - --model-config.action-horizon 10 \\ - --model-config.robot-type desktop +For single-arm LIBERO checkpoints, start the vanilla server (raw model output):: + + bash scripts/run_serving.sh \\ + --checkpoint-path /path/to/ckpt \\ + --train-config-path /path/to/libero.yml \\ + --action-horizon 10 \\ + --port 32194 """ from __future__ import annotations @@ -361,15 +358,9 @@ def decode_predict_action( if parts: return np.concatenate(parts, axis=1).astype(np.float32) - try: - from wall_x._vendor.harrix.envs.libero_common import decode_chunk - - chunk = decode_chunk(pa, train_config) - if chunk.ndim == 2 and chunk.shape[1] > 0: - return chunk.astype(np.float32) - except Exception: - pass - + # No dof_config available: fall back to the raw model action rows. (The + # primary path above already drops the virtual action_padding tail, which is + # the OSS-correct decode.) return pa.astype(np.float32) diff --git a/scripts/fake_inference.py b/scripts/fake_inference.py index d431d2c5..5649467a 100755 --- a/scripts/fake_inference.py +++ b/scripts/fake_inference.py @@ -13,7 +13,6 @@ from pathlib import Path import numpy as np -import torch def _ensure_local_harrix_on_path() -> None: @@ -56,35 +55,13 @@ def _build_fake_observation(seed: int, image_size: int, gripper_dim: int) -> dic def _predict_raw_action(adapter, payload: dict) -> np.ndarray: - from wall_x._vendor.harrix.envs.libero_common import encode_proprio + """Run one flow-inference pass via the adapter's batched predict path. - observation = encode_proprio( - payload["observation"], adapter._train_config, adapter._action_horizon - ) - prefix, postfix = adapter._get_flow_prompt(payload["instruction"]) - batch_inputs = adapter._construct_model_input([observation], [prefix], [postfix]) - - dataset_names = batch_inputs["dataset_names"] - padding = torch.zeros( - (len(dataset_names), 1, adapter._action_dim), - dtype=torch.float32, - ) - padding_action = adapter._normalizer_action.normalize_data( - padding, dataset_names - ).to(batch_inputs["input_ids"].device) - - model_output = adapter._model.generate_flow_action( - action_horizon=adapter._action_horizon, - action_dim=adapter._action_dim, - num_inference_timesteps=adapter._num_inference_timesteps, - padding_action=padding_action, - noise=payload.get("noise"), - **batch_inputs, - ) - predict_action = model_output["predict_action"] - if isinstance(predict_action, torch.Tensor): - predict_action = predict_action.detach().cpu().numpy() - return np.asarray(predict_action[0]) + ``predict_batch`` is the same path the LIBERO evaluation uses; it handles + proprio encoding, prompt construction, and flow sampling internally. + """ + actions = adapter.predict_batch([payload]) + return np.asarray(actions[0]) def parse_args() -> argparse.Namespace: @@ -117,9 +94,9 @@ def main() -> int: args = parse_args() _ensure_local_harrix_on_path() - import wall_x._vendor.harrix.adapters # noqa: F401 register model adapters - from wall_x._vendor.harrix.adapters.registry import build_adapter - from wall_x._vendor.harrix.eval_config import ( + import wall_x._vendor.harrix.e2e_infer.adapters # noqa: F401 register model adapters + from wall_x._vendor.harrix.e2e_infer.adapters.registry import build_adapter + from wall_x._vendor.harrix.e2e_infer.config import ( EvalConfig, LiberoEnvParams, autofill_from_checkpoint, @@ -137,7 +114,9 @@ def main() -> int: cfg = autofill_from_checkpoint(cfg) adapter = build_adapter(cfg) - gripper_dim = _right_gripper_dim_from_config(getattr(adapter, "_train_config", {})) + gripper_dim = _right_gripper_dim_from_config( + {"agent_pos_config": getattr(adapter._typed_cfg.task, "agent_pos_config", {})} + ) payload = { "observation": _build_fake_observation(args.seed, args.image_size, gripper_dim), "instruction": args.instruction, diff --git a/scripts/infer_libero.py b/scripts/infer_libero.py index 8d224386..e8b3412c 100755 --- a/scripts/infer_libero.py +++ b/scripts/infer_libero.py @@ -15,6 +15,7 @@ import yaml + LIBERO_DEFAULT_MAX_INFER_TIMES = { "libero_spatial": 22, "libero_object": 28, @@ -211,10 +212,7 @@ def main() -> int: args = parse_args() _ensure_local_harrix_on_path() - from wall_x._vendor.harrix.eval_config import ( - autofill_from_checkpoint, - load_eval_config, - ) + from wall_x._vendor.harrix.e2e_infer.config import autofill_from_checkpoint, load_eval_config raw = _load_or_build_raw_config(args) @@ -225,7 +223,7 @@ def main() -> int: cfg = autofill_from_checkpoint(load_eval_config(tmp_config)) if cfg.runtime.driver_mode != "in_process": raise ValueError("Only driver_mode='in_process' is supported") - from wall_x._vendor.harrix.drivers.inproc import run + from wall_x._vendor.harrix.e2e_infer.drivers.inproc import run run(cfg) return 0 diff --git a/scripts/normalize.py b/scripts/normalize.py new file mode 100644 index 00000000..9b1ef008 --- /dev/null +++ b/scripts/normalize.py @@ -0,0 +1,163 @@ +# Adapted from Physical-Intelligence/openpi: +# https://github.com/Physical-Intelligence/openpi/blob/650c5b0283a49c42784fb5055a0507da2c6d347d/src/openpi/shared/normalize.py +# Upstream license: Apache-2.0. +import json +import pathlib + +import numpy as np +import numpydantic +import pydantic + + +@pydantic.dataclasses.dataclass +class NormStats: + mean: numpydantic.NDArray + std: numpydantic.NDArray + q01: numpydantic.NDArray | None = None # 1st quantile + q99: numpydantic.NDArray | None = None # 99th quantile + + +class RunningStats: + """Compute running statistics of a batch of vectors.""" + + def __init__(self): + self._count = 0 + self._mean = None + self._mean_of_squares = None + self._min = None + self._max = None + self._histograms = None + self._bin_edges = None + self._num_quantile_bins = 5000 # for computing quantiles on the fly + + def update(self, batch: np.ndarray) -> None: + """ + Update the running statistics with a batch of vectors. + + Args: + vectors (np.ndarray): An array where all dimensions except the last are batch dimensions. + """ + batch = batch.reshape(-1, batch.shape[-1]) + num_elements, vector_length = batch.shape + if self._count == 0: + self._mean = np.mean(batch, axis=0) + self._mean_of_squares = np.mean(batch**2, axis=0) + self._min = np.min(batch, axis=0) + self._max = np.max(batch, axis=0) + self._histograms = [ + np.zeros(self._num_quantile_bins) for _ in range(vector_length) + ] + self._bin_edges = [ + np.linspace( + self._min[i] - 1e-10, + self._max[i] + 1e-10, + self._num_quantile_bins + 1, + ) + for i in range(vector_length) + ] + else: + if vector_length != self._mean.size: + raise ValueError( + "The length of new vectors does not match the initialized vector length." + ) + new_max = np.max(batch, axis=0) + new_min = np.min(batch, axis=0) + max_changed = np.any(new_max > self._max) + min_changed = np.any(new_min < self._min) + self._max = np.maximum(self._max, new_max) + self._min = np.minimum(self._min, new_min) + + if max_changed or min_changed: + self._adjust_histograms() + + self._count += num_elements + + batch_mean = np.mean(batch, axis=0) + batch_mean_of_squares = np.mean(batch**2, axis=0) + + # Update running mean and mean of squares. + self._mean += (batch_mean - self._mean) * (num_elements / self._count) + self._mean_of_squares += (batch_mean_of_squares - self._mean_of_squares) * ( + num_elements / self._count + ) + + self._update_histograms(batch) + + def get_statistics(self) -> NormStats: + """ + Compute and return the statistics of the vectors processed so far. + + Returns: + dict: A dictionary containing the computed statistics. + """ + if self._count < 2: + raise ValueError("Cannot compute statistics for less than 2 vectors.") + + variance = self._mean_of_squares - self._mean**2 + stddev = np.sqrt(np.maximum(0, variance)) + q01, q99 = self._compute_quantiles([0.01, 0.99]) + return NormStats(mean=self._mean, std=stddev, q01=q01, q99=q99) + + def _adjust_histograms(self): + """Adjust histograms when min or max changes.""" + for i in range(len(self._histograms)): + old_edges = self._bin_edges[i] + new_edges = np.linspace( + self._min[i], self._max[i], self._num_quantile_bins + 1 + ) + + # Redistribute the existing histogram counts to the new bins + new_hist, _ = np.histogram( + old_edges[:-1], bins=new_edges, weights=self._histograms[i] + ) + + self._histograms[i] = new_hist + self._bin_edges[i] = new_edges + + def _update_histograms(self, batch: np.ndarray) -> None: + """Update histograms with new vectors.""" + for i in range(batch.shape[1]): + hist, _ = np.histogram(batch[:, i], bins=self._bin_edges[i]) + self._histograms[i] += hist + + def _compute_quantiles(self, quantiles): + """Compute quantiles based on histograms.""" + results = [] + for q in quantiles: + target_count = q * self._count + q_values = [] + for hist, edges in zip(self._histograms, self._bin_edges, strict=True): + cumsum = np.cumsum(hist) + idx = np.searchsorted(cumsum, target_count) + q_values.append(edges[idx]) + results.append(np.array(q_values)) + return results + + +class _NormStatsDict(pydantic.BaseModel): + norm_stats: dict[str, NormStats] + + +def serialize_json(norm_stats: dict[str, NormStats]) -> str: + """Serialize the running statistics to a JSON string.""" + return _NormStatsDict(norm_stats=norm_stats).model_dump_json(indent=2) + + +def deserialize_json(data: str) -> dict[str, NormStats]: + """Deserialize the running statistics from a JSON string.""" + return _NormStatsDict(**json.loads(data)).norm_stats + + +def save(directory: pathlib.Path | str, norm_stats: dict[str, NormStats]) -> None: + """Save the normalization stats to a directory.""" + path = pathlib.Path(directory) / "norm_stats.json" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(serialize_json(norm_stats)) + + +def load(directory: pathlib.Path | str) -> dict[str, NormStats]: + """Load the normalization stats from a directory.""" + path = pathlib.Path(directory) / "norm_stats.json" + if not path.exists(): + raise FileNotFoundError(f"Norm stats file not found at: {path}") + return deserialize_json(path.read_text()) diff --git a/scripts/run_serving.sh b/scripts/run_serving.sh index 14bfeebf..07891046 100755 --- a/scripts/run_serving.sh +++ b/scripts/run_serving.sh @@ -1,4 +1,18 @@ #!/usr/bin/env bash +# Launch the open-source (vanilla) Wall-X websocket policy server. +# +# Serves a checkpoint over a websocket so a client can stream observations and +# receive action chunks. It uses the vanilla model path (identical to the LIBERO +# evaluation), with no closed-source inference engine. +# +# The vanilla cut returns RAW model action chunks (shape action_horizon x +# action_dim); robot-side serialization/interpolation is left to the client, +# because the real-robot preprocessors are not part of the open-source release. +# Flags that require the closed inference engine (CUDA graph / experimental +# engine) or the excluded robot preprocessors (action serialization / robot +# types / non-LIBERO envs) are recognized and rejected with a clear message, so +# a command copied from the full (internal) serving stack fails loudly instead +# of silently behaving differently. set -euo pipefail usage() { @@ -7,38 +21,47 @@ Usage: bash scripts/run_serving.sh --checkpoint-path /path/to/checkpoint [options] Required: - --checkpoint-path PATH Checkpoint directory or checkpoint file. - -Common options: - --train-config-path PATH Training config used by the checkpoint. - --port PORT WebSocket port. Default: 32195. - --host HOST Bind host. Default: 0.0.0.0. - --env X2ROBOT|LIBERO Serving environment. Default: X2ROBOT. - --cuda-id ID Sets CUDA_VISIBLE_DEVICES. Default: 0. - --image-passing-mode MODE base64 or numpy. Default: base64. - --action-horizon N Model action horizon. Default: 32. - --robot-type TYPE desktop, turtle, or ex001. Default: desktop. - --raw-actions Alias for --no-serialize-actions. - --serialize-actions Return robot-serialized actions. - --no-serialize-actions Return raw model action chunks. Default. - --max-batch-size N Enable dynamic batching. - --enable-cuda-graph Enable CUDA graph in the serving runtime. - --enable-experimental-engine Enable the experimental inference engine. - --debug Enable debug logging. - --dry-run Print the command without running it. - -Additional arguments after "--" are forwarded to launch_serving.py, for example: - bash scripts/run_serving.sh --checkpoint-path /ckpt -- \ - --model-config.norm-key libero_all - -Environment variables can also be used, e.g. CHECKPOINT_PATH, -TRAIN_CONFIG_PATH, PORT, CUDA_ID, ACTION_HORIZON, WALLX_ENV. + --checkpoint-path PATH Checkpoint directory or checkpoint file. + +Options (vanilla serving): + --train-config-path PATH Training config used by the checkpoint + (default: /config.yml). + --norm-key KEY Normalizer key. Default: libero_all. + --cam-names "V1 V2 ..." Camera view names in model input order. + Default: "face_view right_wrist_view". + --action-horizon N Model action horizon (default: from train config). + --host HOST Bind host. Default: 0.0.0.0. + --port PORT WebSocket port. Default: 32195. + --cuda-id ID Sets CUDA_VISIBLE_DEVICES. Default: 0. + --image-passing-mode MODE base64 or numpy. Default: base64. + --max-batch-size N Enable server-side dynamic batching (0 = disabled). + --env ENV Serving environment. Only LIBERO is supported by the + vanilla obs assembly. Default: LIBERO. + --debug Enable debug logging. + --dry-run Print the command without running it. + +Not available in the open-source (vanilla) cut — rejected with guidance: + --serialize-actions / --robot-type (need the excluded robot preprocessors; + vanilla returns raw action chunks) + --enable-cuda-graph / --enable-experimental-engine (closed opt inference engine) + +The server returns raw model action chunks (shape action_horizon x action_dim); +robot-side action serialization is left to the client. EOF } +# Reject a flag that only the full (internal) serving stack supports, pointing at +# the reason, instead of accepting it as a silent no-op. +reject_full_stack_flag() { + echo "error: '$1' is not available in the open-source (vanilla) serving." >&2 + echo " $2" >&2 + echo " The vanilla server returns raw model action chunks; serialize" >&2 + echo " and post-process on the client side." >&2 + exit 2 +} + SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" SOURCE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" - if [[ -d "${SOURCE_ROOT}/wall_x" ]]; then export PYTHONPATH="${SOURCE_ROOT}:${PYTHONPATH:-}" fi @@ -46,179 +69,72 @@ fi PYTHON_BIN="${PYTHON_BIN:-python}" CHECKPOINT_PATH="${CHECKPOINT_PATH:-}" TRAIN_CONFIG_PATH="${TRAIN_CONFIG_PATH:-}" -PORT="${PORT:-32195}" +NORM_KEY="${NORM_KEY:-libero_all}" +CAM_NAMES="${CAM_NAMES:-face_view right_wrist_view}" +ACTION_HORIZON="${ACTION_HORIZON:-}" HOST="${HOST:-0.0.0.0}" -WALLX_ENV="${WALLX_ENV:-X2ROBOT}" +PORT="${PORT:-32195}" CUDA_ID="${CUDA_ID:-0}" IMAGE_PASSING_MODE="${IMAGE_PASSING_MODE:-base64}" -ACTION_HORIZON="${ACTION_HORIZON:-32}" -ROBOT_TYPE="${ROBOT_TYPE:-desktop}" -ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${ROBOT_ACTION_INTERPOLATE_MULTIPLIER:-1}" -ROBOT_ACTION_END_RATIO="${ROBOT_ACTION_END_RATIO:-1.0}" -MODEL_DEVICE="${MODEL_DEVICE:-cuda}" MAX_BATCH_SIZE="${MAX_BATCH_SIZE:-}" -DEFAULT_PROMPT="${DEFAULT_PROMPT:-}" -SERIALIZE_ACTIONS="${SERIALIZE_ACTIONS:-0}" -ENABLE_CUDA_GRAPH="${ENABLE_CUDA_GRAPH:-0}" -ENABLE_EXPERIMENTAL_ENGINE="${ENABLE_EXPERIMENTAL_ENGINE:-0}" -DEBUG="${DEBUG:-0}" +DEBUG=0 DRY_RUN=0 -EXTRA_ARGS=() while [[ $# -gt 0 ]]; do case "$1" in - -h|--help) - usage - exit 0 - ;; - --checkpoint-path) - CHECKPOINT_PATH="${2:?missing value for --checkpoint-path}" - shift 2 - ;; - --checkpoint-path=*) - CHECKPOINT_PATH="${1#*=}" - shift - ;; - --train-config-path) - TRAIN_CONFIG_PATH="${2:?missing value for --train-config-path}" - shift 2 - ;; - --train-config-path=*) - TRAIN_CONFIG_PATH="${1#*=}" - shift - ;; - --port) - PORT="${2:?missing value for --port}" - shift 2 - ;; - --port=*) - PORT="${1#*=}" - shift - ;; - --host) - HOST="${2:?missing value for --host}" - shift 2 - ;; - --host=*) - HOST="${1#*=}" - shift - ;; - --env) - WALLX_ENV="${2:?missing value for --env}" - shift 2 - ;; - --env=*) - WALLX_ENV="${1#*=}" - shift - ;; - --cuda-id) - CUDA_ID="${2:?missing value for --cuda-id}" - shift 2 - ;; - --cuda-id=*) - CUDA_ID="${1#*=}" - shift - ;; - --image-passing-mode) - IMAGE_PASSING_MODE="${2:?missing value for --image-passing-mode}" - shift 2 - ;; - --image-passing-mode=*) - IMAGE_PASSING_MODE="${1#*=}" - shift - ;; - --action-horizon) - ACTION_HORIZON="${2:?missing value for --action-horizon}" - shift 2 - ;; - --action-horizon=*) - ACTION_HORIZON="${1#*=}" - shift - ;; - --robot-type) - ROBOT_TYPE="${2:?missing value for --robot-type}" - shift 2 - ;; - --robot-type=*) - ROBOT_TYPE="${1#*=}" - shift - ;; - --robot-action-interpolate-multiplier) - ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${2:?missing value for --robot-action-interpolate-multiplier}" - shift 2 - ;; - --robot-action-interpolate-multiplier=*) - ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${1#*=}" - shift - ;; - --robot-action-end-ratio) - ROBOT_ACTION_END_RATIO="${2:?missing value for --robot-action-end-ratio}" - shift 2 - ;; - --robot-action-end-ratio=*) - ROBOT_ACTION_END_RATIO="${1#*=}" - shift - ;; - --model-device) - MODEL_DEVICE="${2:?missing value for --model-device}" - shift 2 - ;; - --model-device=*) - MODEL_DEVICE="${1#*=}" - shift - ;; - --max-batch-size) - MAX_BATCH_SIZE="${2:?missing value for --max-batch-size}" - shift 2 - ;; - --max-batch-size=*) - MAX_BATCH_SIZE="${1#*=}" - shift - ;; - --default-prompt) - DEFAULT_PROMPT="${2:?missing value for --default-prompt}" - shift 2 - ;; - --default-prompt=*) - DEFAULT_PROMPT="${1#*=}" - shift - ;; + -h|--help) usage; exit 0 ;; + --checkpoint-path) CHECKPOINT_PATH="${2:?}"; shift 2 ;; + --checkpoint-path=*) CHECKPOINT_PATH="${1#*=}"; shift ;; + --train-config-path) TRAIN_CONFIG_PATH="${2:?}"; shift 2 ;; + --train-config-path=*) TRAIN_CONFIG_PATH="${1#*=}"; shift ;; + --norm-key) NORM_KEY="${2:?}"; shift 2 ;; + --norm-key=*) NORM_KEY="${1#*=}"; shift ;; + --cam-names) CAM_NAMES="${2:?}"; shift 2 ;; + --cam-names=*) CAM_NAMES="${1#*=}"; shift ;; + --action-horizon) ACTION_HORIZON="${2:?}"; shift 2 ;; + --action-horizon=*) ACTION_HORIZON="${1#*=}"; shift ;; + --host) HOST="${2:?}"; shift 2 ;; + --host=*) HOST="${1#*=}"; shift ;; + --port) PORT="${2:?}"; shift 2 ;; + --port=*) PORT="${1#*=}"; shift ;; + --cuda-id) CUDA_ID="${2:?}"; shift 2 ;; + --cuda-id=*) CUDA_ID="${1#*=}"; shift ;; + --image-passing-mode) IMAGE_PASSING_MODE="${2:?}"; shift 2 ;; + --image-passing-mode=*) IMAGE_PASSING_MODE="${1#*=}"; shift ;; + --max-batch-size) MAX_BATCH_SIZE="${2:?}"; shift 2 ;; + --max-batch-size=*) MAX_BATCH_SIZE="${1#*=}"; shift ;; + --env) ENV_ARG="${2:?}"; shift 2 ;; + --env=*) ENV_ARG="${1#*=}"; shift ;; + --raw-actions|--no-serialize-actions) shift ;; # raw is the vanilla default --serialize-actions) - SERIALIZE_ACTIONS=1 - shift - ;; - --no-serialize-actions|--raw-actions) - SERIALIZE_ACTIONS=0 - shift - ;; - --enable-cuda-graph) - ENABLE_CUDA_GRAPH=1 - shift - ;; - --enable-experimental-engine) - ENABLE_EXPERIMENTAL_ENGINE=1 - shift - ;; - --debug) - DEBUG=1 - shift - ;; - --dry-run) - DRY_RUN=1 - shift - ;; - --) - shift - EXTRA_ARGS+=("$@") - break - ;; - *) - EXTRA_ARGS+=("$1") - shift - ;; + reject_full_stack_flag "$1" \ + "It needs the robot preprocessors, which are not in the OSS cut." ;; + --robot-type|--robot-type=*) + reject_full_stack_flag "${1%%=*}" \ + "It needs the robot preprocessors, which are not in the OSS cut." ;; + --enable-cuda-graph|--enable-experimental-engine) + reject_full_stack_flag "$1" \ + "It needs the closed opt inference engine, not in the OSS cut." ;; + --debug) DEBUG=1; shift ;; + --dry-run) DRY_RUN=1; shift ;; + *) echo "error: unknown argument '$1'" >&2; echo >&2; usage >&2; exit 2 ;; esac done +# --env is a convenience: only the LIBERO-style observation assembly ships in the +# vanilla cut. Accept LIBERO (the default); reject anything else with guidance. +ENV_ARG="${ENV_ARG:-LIBERO}" +case "${ENV_ARG^^}" in + LIBERO) ;; + *) + echo "error: --env '${ENV_ARG}' is not supported by the vanilla serving." >&2 + echo " Only LIBERO-style observations (eef_pos / eef_axisangle /" >&2 + echo " gripper + face_view / right_wrist_view) are assembled in the" >&2 + echo " open-source cut." >&2 + exit 2 + ;; +esac + if [[ -z "${CHECKPOINT_PATH}" ]]; then echo "error: --checkpoint-path is required." >&2 echo >&2 @@ -226,55 +142,35 @@ if [[ -z "${CHECKPOINT_PATH}" ]]; then exit 2 fi -export CUDA_VISIBLE_DEVICES="${CUDA_ID}" -export ENABLE_FAST_PREPROCESS="${ENABLE_FAST_PREPROCESS:-true}" - CMD=( - "${PYTHON_BIN}" -m wall_x._vendor.harrix.serving.launch_serving - --env "${WALLX_ENV}" + "${PYTHON_BIN}" -m wall_x._vendor.harrix.serving.launch_vanilla_serving + --checkpoint-path "${CHECKPOINT_PATH}" + --norm-key "${NORM_KEY}" --host "${HOST}" --port "${PORT}" + --cuda-id "${CUDA_ID}" --image-passing-mode "${IMAGE_PASSING_MODE}" ) - -if [[ "${SERIALIZE_ACTIONS}" == "1" ]]; then - CMD+=(--serialize-actions) -else - CMD+=(--no-serialize-actions) +# --cam-names is nargs="+" on the launcher; pass each token unquoted-split. +if [[ -n "${CAM_NAMES}" ]]; then + # shellcheck disable=SC2206 + CAM_ARR=(${CAM_NAMES}) + CMD+=(--cam-names "${CAM_ARR[@]}") fi -if [[ -n "${MAX_BATCH_SIZE}" ]]; then - CMD+=(--max-batch-size "${MAX_BATCH_SIZE}") -fi -if [[ -n "${DEFAULT_PROMPT}" ]]; then - CMD+=(--default-prompt "${DEFAULT_PROMPT}") +if [[ -n "${TRAIN_CONFIG_PATH}" ]]; then + CMD+=(--train-config-path "${TRAIN_CONFIG_PATH}") fi -if [[ "${ENABLE_CUDA_GRAPH}" == "1" ]]; then - CMD+=(--enable-cuda-graph) +if [[ -n "${ACTION_HORIZON}" ]]; then + CMD+=(--action-horizon "${ACTION_HORIZON}") fi -if [[ "${ENABLE_EXPERIMENTAL_ENGINE}" == "1" ]]; then - CMD+=(--enable-experimental-engine) +if [[ -n "${MAX_BATCH_SIZE}" ]]; then + CMD+=(--max-batch-size "${MAX_BATCH_SIZE}") fi if [[ "${DEBUG}" == "1" ]]; then CMD+=(--debug) fi -CMD+=( - model-config:server-model-config - --model-config.checkpoint-path "${CHECKPOINT_PATH}" - --model-config.action-horizon "${ACTION_HORIZON}" - --model-config.robot-type "${ROBOT_TYPE}" - --model-config.robot-action-interpolate-multiplier "${ROBOT_ACTION_INTERPOLATE_MULTIPLIER}" - --model-config.robot-action-end-ratio "${ROBOT_ACTION_END_RATIO}" - --model-config.model-device "${MODEL_DEVICE}" -) - -if [[ -n "${TRAIN_CONFIG_PATH}" ]]; then - CMD+=(--model-config.train-config-path "${TRAIN_CONFIG_PATH}") -fi - -CMD+=("${EXTRA_ARGS[@]}") - -printf 'Launching Wall-X serving:\n' +printf 'Launching Wall-X vanilla serving:\n' printf ' %q' "${CMD[@]}" printf '\n' diff --git a/wall_x/_vendor/harrix/adapters/__init__.py b/wall_x/_vendor/harrix/adapters/__init__.py deleted file mode 100644 index 5791ac80..00000000 --- a/wall_x/_vendor/harrix/adapters/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Adapter package entry point. - -Importing ``variants`` triggers variant registration side effects. -""" - -from wall_x._vendor.harrix.adapters import variants # noqa: F401 diff --git a/wall_x/_vendor/harrix/adapters/variants/qwen2_5.py b/wall_x/_vendor/harrix/adapters/variants/qwen2_5.py deleted file mode 100644 index 036409b5..00000000 --- a/wall_x/_vendor/harrix/adapters/variants/qwen2_5.py +++ /dev/null @@ -1,17 +0,0 @@ -"""Qwen2.5 inference adapter registration. - -The shared inference logic lives in ``harrix.adapters.qwen_vlact``. This module -binds the Qwen2.5 variant key to Wall-X's training-side adapter class. -""" - -from wall_x._vendor.harrix.adapters.qwen_vlact import QwenVLActInferAdapter -from wall_x._vendor.harrix.adapters.registry import register_adapter - - -@register_adapter("qwen2_5") -class Qwen2_5InferAdapter(QwenVLActInferAdapter): - @classmethod - def _training_adapter(cls): - from wall_x.model.qact.qwen2_5.adapter import Qwen2_5Adapter - - return Qwen2_5Adapter diff --git a/wall_x/_vendor/harrix/drivers/inproc/__init__.py b/wall_x/_vendor/harrix/drivers/inproc/__init__.py deleted file mode 100644 index 5813d0cb..00000000 --- a/wall_x/_vendor/harrix/drivers/inproc/__init__.py +++ /dev/null @@ -1,197 +0,0 @@ -"""In-process driver: DirectModelHandle + N x SubprocEnvHandle. - -The model lives in the driver process while simulator envs live in subprocesses. -Evaluation proceeds in chunk-level lockstep: - 1. Seed the driver process and construct the model handle. - 2. Start one subprocess env handle per worker. - 3. Claim a task-local frame from JobState. - 4. Reset envs, batch active observations, run model.predict_batch, then - fan out action chunks to env subprocesses. - 5. Complete all episodes in the frame and move to the next frame. -""" - -from __future__ import annotations - -import json -import logging -import os -import time - -import numpy as np - -from wall_x._vendor.harrix.eval_config import EvalConfig -from wall_x._vendor.harrix.drivers.inproc.env_handle import SubprocEnvHandle -from wall_x._vendor.harrix.drivers.inproc.model_handle import DirectModelHandle -from wall_x._vendor.harrix.drivers.job_state import JobState - -logger = logging.getLogger(__name__) - - -def run(cfg: EvalConfig) -> None: - import wall_x._vendor.harrix.envs # noqa: F401 trigger env register - import wall_x._vendor.harrix.adapters # noqa: F401 trigger adapter register - from wall_x._vendor.harrix.envs.registry import enumerate_episodes_for - from wall_x._vendor.harrix.utils.seed import set_seed_everywhere - - # 1) Seed the driver process. - set_seed_everywhere(cfg.env.seed) - - # 2) Build DirectModelHandle; the model is loaded in the driver process. - logger.info("Constructing DirectModelHandle in the driver process") - t_model = time.time() - model_handle = DirectModelHandle(cfg) - logger.info("Model loaded in %.1fs", time.time() - t_model) - - # 3) Start env subprocesses. - logger.info("Starting %s SubprocEnvHandle(s)", cfg.runtime.num_workers) - env_handles = [ - SubprocEnvHandle(cfg, worker_id=i) for i in range(cfg.runtime.num_workers) - ] - - # 4) JobState runs in frame-sync mode for lockstep evaluation. - os.makedirs(cfg.runtime.log_dir, exist_ok=True) - log_path = os.path.join(cfg.runtime.log_dir, "state.jsonl") - report_path = os.path.join(cfg.runtime.log_dir, "report.json") - episodes = enumerate_episodes_for(cfg) - logger.info("env.type=%r; scheduled %s episodes", cfg.env.type, len(episodes)) - state = JobState( - episodes, - log_path, - batch_sync_mode=True, - batch_size=cfg.runtime.num_workers, - ) - - # 5) Main loop, one frame at a time. - t0 = time.time() - frame_idx = 0 - while not state.is_drained(): - frame_eps = state.claim_frame() - if not frame_eps: - # Previous frame still has in-flight episodes. - time.sleep(0.05) - continue - - results = _run_frame(frame_eps, model_handle, env_handles, cfg, frame_idx) - for ep, res in zip(frame_eps, results): - if "_error" in res: - state.fail(ep, str(res["_error"])) - else: - state.complete(ep, res) - frame_idx += 1 - - elapsed = time.time() - t0 - logger.info("All episodes finished in %.1fs (%.1f min)", elapsed, elapsed / 60) - - state.dump_final(report_path) - with open(report_path) as f: - report = json.load(f) - overall = report["overall"] - logger.info( - "attempted=%s, successes=%s, success_rate=%.2f%%, failed=%s", - overall["attempted"], - overall["successes"], - overall["success_rate"] * 100, - overall["failed"], - ) - logger.info("Report: %s", report_path) - logger.info("State log: %s", log_path) - - for h in env_handles: - h.shutdown() - model_handle.shutdown() - - -def _run_frame( - frame_eps: list, - model_handle: DirectModelHandle, - env_handles: list, - cfg: EvalConfig, - frame_idx: int, -) -> list[dict]: - """Run one task-local frame with chunk-level lockstep. - - Active workers are batched together at each chunk boundary. Workers that - already finished no longer participate in later forwards. - """ - from wall_x._vendor.harrix.envs.libero_common import encode_raw_obs - - n = len(frame_eps) - t_frame_start = time.time() - - # ---- a) reset: fan out, then gather ---- - for i in range(n): - env_handles[i].submit_reset(tuple(frame_eps[i])) - initials = [env_handles[i].wait_reset() for i in range(n)] - - obs_list = [r["obs"] for r in initials] - instr_list = [r["instruction"] for r in initials] - status = [ - { - "done": False, - "success": False, - "steps": 0, - "task_desc": initials[i].get("task_desc", ""), - } - for i in range(n) - ] - - max_rounds = cfg.env.libero.max_infer_times - - # ---- b) chunk lockstep ---- - for round_idx in range(max_rounds): - active = [i for i in range(n) if not status[i]["done"]] - if not active: - break - - payloads = [ - { - "observation": encode_raw_obs(obs_list[i]), - "instruction": instr_list[i], - "noise": None, - } - for i in active - ] - chunks = model_handle.predict_batch(payloads) - - # fan-out submit - for k, i in enumerate(active): - env_handles[i].submit_execute_chunk(chunks[k]) - # gather - for k, i in enumerate(active): - try: - r = env_handles[i].wait_execute_chunk() - except Exception as e: - status[i]["_error"] = str(e) - status[i]["done"] = True - continue - obs_list[i] = r["obs"] - status[i]["steps"] += r["steps"] - if r["done"]: - status[i]["done"] = True - status[i]["success"] = True - - for i in range(n): - env_handles[i].submit_finalize_episode(status[i]["success"]) - for i in range(n): - env_handles[i].wait_finalize_episode() - - elapsed_frame = time.time() - t_frame_start - logger.info( - "frame=%s n=%s succ=%s/%s elapsed=%.1fs", - frame_idx, - n, - sum(s["success"] for s in status), - n, - elapsed_frame, - ) - - return [ - { - "success": s["success"], - "steps": s["steps"], - "elapsed_sec": round(elapsed_frame / max(1, n), 3), - "task_desc": s["task_desc"], - **({"_error": s["_error"]} if "_error" in s else {}), - } - for s in status - ] diff --git a/wall_x/_vendor/harrix/e2e_infer/__init__.py b/wall_x/_vendor/harrix/e2e_infer/__init__.py new file mode 100644 index 00000000..dc5abf3a --- /dev/null +++ b/wall_x/_vendor/harrix/e2e_infer/__init__.py @@ -0,0 +1 @@ +"""End-to-end inference drivers, environments and model adapters.""" diff --git a/wall_x/_vendor/harrix/e2e_infer/adapters/__init__.py b/wall_x/_vendor/harrix/e2e_infer/adapters/__init__.py new file mode 100644 index 00000000..88ed9c10 --- /dev/null +++ b/wall_x/_vendor/harrix/e2e_infer/adapters/__init__.py @@ -0,0 +1,11 @@ +"""Adapter package entry point. + +Importing ``variants`` triggers variant registration side effects. +""" + +import logging + +from wall_x._vendor.harrix.e2e_infer.adapters import variants # noqa: F401 + +logger = logging.getLogger(__name__) + diff --git a/wall_x/_vendor/harrix/adapters/base.py b/wall_x/_vendor/harrix/e2e_infer/adapters/base.py similarity index 71% rename from wall_x/_vendor/harrix/adapters/base.py rename to wall_x/_vendor/harrix/e2e_infer/adapters/base.py index 68b7f91b..4ff8a54e 100644 --- a/wall_x/_vendor/harrix/adapters/base.py +++ b/wall_x/_vendor/harrix/e2e_infer/adapters/base.py @@ -1,9 +1,4 @@ -"""Common inference adapter abstraction. - -The adapter owns all architecture-specific model setup and exposes a single -``predict_batch`` entry point. Environment drivers pass payloads through without -interpreting the env-adapter schema. -""" +"""Thin environment-to-executor adapter abstraction.""" from __future__ import annotations @@ -11,17 +6,13 @@ import numpy as np -from wall_x._vendor.harrix.eval_config import EvalConfig +from wall_x._vendor.harrix.e2e_infer.config import EvalConfig class BaseInferAdapter(ABC): @abstractmethod def __init__(self, cfg: EvalConfig) -> None: - """Load processors, model config, checkpoints, and normalizers. - - Subclasses should reject unsupported ``cfg.model.action_mode`` values - during construction. - """ + """Build the owned executor and environment conversion state.""" @property @abstractmethod diff --git a/wall_x/_vendor/harrix/e2e_infer/adapters/qwen_vlact.py b/wall_x/_vendor/harrix/e2e_infer/adapters/qwen_vlact.py new file mode 100644 index 00000000..06bbf099 --- /dev/null +++ b/wall_x/_vendor/harrix/e2e_infer/adapters/qwen_vlact.py @@ -0,0 +1,49 @@ +"""Thin environment adapter for the Qwen2.5 evaluation executor.""" + +import numpy as np + +from wall_x._vendor.harrix.e2e_infer.adapters.base import BaseInferAdapter +from wall_x._vendor.harrix.e2e_infer.envs.libero_common import decode_chunk, encode_proprio +from wall_x._vendor.harrix.model_executors.qwen2_5_eval import Qwen2_5EvalExecutor + + +class QwenVLActInferAdapter(BaseInferAdapter): + """Select the model variant and delegate execution to model_executors.""" + + @classmethod + def _training_adapter(cls): + raise NotImplementedError( + f"{cls.__name__} must override _training_adapter()" + ) + + def __init__(self, cfg) -> None: + self._executor = Qwen2_5EvalExecutor(cfg, self._training_adapter()) + + @property + def chunk_horizon(self) -> int: + return self._executor.chunk_horizon + + @property + def action_mode(self) -> str: + return self._executor.action_mode + + @property + def action_dim(self) -> int: + return self._executor.action_dim + + def predict_batch(self, payloads): + task = self._executor.task_config + requests = [ + { + **payload, + "observation": encode_proprio( + payload["observation"], task, self.chunk_horizon + ), + } + for payload in payloads + ] + outputs = self._executor.infer_batch(requests) + return [ + np.asarray(decode_chunk(output, task), dtype=np.float32) + for output in outputs + ] diff --git a/wall_x/_vendor/harrix/adapters/registry.py b/wall_x/_vendor/harrix/e2e_infer/adapters/registry.py similarity index 87% rename from wall_x/_vendor/harrix/adapters/registry.py rename to wall_x/_vendor/harrix/e2e_infer/adapters/registry.py index ddc0ec4d..ef14acdf 100644 --- a/wall_x/_vendor/harrix/adapters/registry.py +++ b/wall_x/_vendor/harrix/e2e_infer/adapters/registry.py @@ -4,8 +4,8 @@ from typing import Type -from wall_x._vendor.harrix.adapters.base import BaseInferAdapter -from wall_x._vendor.harrix.eval_config import EvalConfig +from wall_x._vendor.harrix.e2e_infer.adapters.base import BaseInferAdapter +from wall_x._vendor.harrix.e2e_infer.config import EvalConfig ADAPTER_REGISTRY: dict[str, Type[BaseInferAdapter]] = {} diff --git a/wall_x/_vendor/harrix/adapters/variants/__init__.py b/wall_x/_vendor/harrix/e2e_infer/adapters/variants/__init__.py similarity index 100% rename from wall_x/_vendor/harrix/adapters/variants/__init__.py rename to wall_x/_vendor/harrix/e2e_infer/adapters/variants/__init__.py diff --git a/wall_x/_vendor/harrix/e2e_infer/adapters/variants/qwen2_5.py b/wall_x/_vendor/harrix/e2e_infer/adapters/variants/qwen2_5.py new file mode 100644 index 00000000..0e51ff2d --- /dev/null +++ b/wall_x/_vendor/harrix/e2e_infer/adapters/variants/qwen2_5.py @@ -0,0 +1,16 @@ +"""Qwen2.5 inference adapter registration. + +The shared inference logic lives in ``harrix.e2e_infer.adapters.qwen_vlact``. This module +binds the Qwen2.5 variant key to Wall-X's training-side adapter class. +""" + +from wall_x._vendor.harrix.e2e_infer.adapters.qwen_vlact import QwenVLActInferAdapter +from wall_x._vendor.harrix.e2e_infer.adapters.registry import register_adapter +from wall_x.model.qact.qwen2_5.adapter import Qwen2_5Adapter + + +@register_adapter("qwen2_5") +class Qwen2_5InferAdapter(QwenVLActInferAdapter): + @classmethod + def _training_adapter(cls): + return Qwen2_5Adapter diff --git a/wall_x/_vendor/harrix/eval_config.py b/wall_x/_vendor/harrix/e2e_infer/config.py similarity index 79% rename from wall_x/_vendor/harrix/eval_config.py rename to wall_x/_vendor/harrix/e2e_infer/config.py index 1ab59d61..e843a4a8 100644 --- a/wall_x/_vendor/harrix/eval_config.py +++ b/wall_x/_vendor/harrix/e2e_infer/config.py @@ -25,13 +25,11 @@ force_render_task_indices: null # e.g. [5] always render that task chunk_granular_render_toggle: false rebuild_env_per_episode: false - rollout_dir: null # save third-person MP4 replays per episode - rollout_fps: 30 runtime: num_workers: 1 max_batch_size: 1 ws_port: 8765 - log_dir: /path/to/wallx_log + log_dir: /tmp/wallx_v5_log batch_sync_mode: false debug: deterministic_model: false @@ -56,9 +54,14 @@ class ModelSection: checkpoint_path: str = "" train_config_path: Optional[str] = None norm_key: str = "libero_all" - cam_names: list = field(default_factory=lambda: ["face_view", "right_wrist_view"]) action_horizon: Optional[int] = None - # Adapter implementation key in harrix.adapters.registry.ADAPTER_REGISTRY. + num_inference_steps: Optional[int] = None + cfg_scale: Optional[float] = None + seed: Optional[int] = None + cam_names: list = field( + default_factory=lambda: ["face_view", "right_wrist_view"] + ) + # Adapter implementation key in harrix.e2e_infer.adapters.registry.ADAPTER_REGISTRY. architecture: str = "qwen2_5" # Inference algorithm. Each adapter validates its supported subset. action_mode: str = "flow" @@ -83,15 +86,13 @@ class LiberoEnvParams: # Rebuild the simulator for every episode instead of only on task changes. # This is mainly a debugging option because it changes the simulator RNG path. rebuild_env_per_episode: bool = False - # When set, save a third-person MP4 replay for each episode under this directory. - # Falls back to the ``WALLX_ROLLOUT_DIR`` environment variable when null. rollout_dir: Optional[str] = None rollout_fps: int = 30 @dataclass class EnvSection: - # Env implementation key in harrix.envs.registry. + # Env implementation key in harrix.e2e_infer.envs.registry. type: str = "libero" seed: int = 42 libero: LiberoEnvParams = field(default_factory=LiberoEnvParams) @@ -102,13 +103,16 @@ class RuntimeSection: num_workers: int = 1 max_batch_size: int = 1 ws_port: int = 8765 - log_dir: str = "/path/to/wallx_log" + log_dir: str = "/tmp/wallx_v5_log" # Run fixed task-local frames instead of dynamic work stealing. This improves # reproducibility at the cost of possible idle workers inside a frame. batch_sync_mode: bool = False # Public Wall-X evaluation uses the in-process driver. Other drivers may be # enabled by downstream/internal integrations. driver_mode: str = "in_process" + # In-process driver only: run each env inline in the driver process + # (DirectEnvHandle) instead of a spawned subprocess (SubprocEnvHandle). + inline_env: bool = False @dataclass @@ -130,7 +134,9 @@ def _build_dataclass(cls, raw): if raw is None: return cls() if not isinstance(raw, dict): - raise ValueError(f"Expected dict for {cls.__name__}, got {type(raw).__name__}") + raise ValueError( + f"Expected dict for {cls.__name__}, got {type(raw).__name__}" + ) field_names = {f.name for f in dataclasses.fields(cls)} field_types = {f.name: f.type for f in dataclasses.fields(cls)} unknown = set(raw.keys()) - field_names @@ -164,7 +170,7 @@ def load_eval_config(yaml_path: str) -> EvalConfig: if not cfg.model.checkpoint_path: raise ValueError(f"model.checkpoint_path is required in {yaml_path}") - if not os.path.isdir(cfg.model.checkpoint_path): + if not _checkpoint_path_exists(cfg.model.checkpoint_path): raise FileNotFoundError( f"model.checkpoint_path does not exist: {cfg.model.checkpoint_path}" ) @@ -172,25 +178,9 @@ def load_eval_config(yaml_path: str) -> EvalConfig: return cfg -def _load_train_config_yaml(path: str) -> dict: - """Load a checkpoint-side train YAML. - - Training checkpoints may contain PyYAML-specific tags such as - ``!!python/tuple``; ``safe_load`` cannot parse those. - """ - with open(path, "r") as f: - return yaml.load(f, Loader=yaml.FullLoader) or {} - - -def _read_action_horizon_flow(train_yml: dict) -> int: - task = train_yml.get("task") or {} - data = train_yml.get("data") or {} - return int( - train_yml.get("action_horizon_flow") - or task.get("action_horizon_flow") - or data.get("action_horizon_flow") - or 32 - ) +def _checkpoint_path_exists(path: str) -> bool: + parts = [p.strip() for p in path.split(",")] if "," in path else [path] + return all(os.path.exists(p) for p in parts) def autofill_from_checkpoint(cfg: EvalConfig) -> EvalConfig: @@ -199,6 +189,12 @@ def autofill_from_checkpoint(cfg: EvalConfig) -> EvalConfig: This resolves ``model.train_config_path`` and ``model.action_horizon``. """ if cfg.model.train_config_path is None: + if "," in cfg.model.checkpoint_path or not os.path.isdir( + cfg.model.checkpoint_path + ): + raise FileNotFoundError( + "model.train_config_path is required when checkpoint_path is not a checkpoint directory" + ) for fname in ("config.yml", "config.yaml"): cand = os.path.join(cfg.model.checkpoint_path, fname) if os.path.exists(cand): @@ -211,7 +207,12 @@ def autofill_from_checkpoint(cfg: EvalConfig) -> EvalConfig: ) if cfg.model.action_horizon is None: - train_yml = _load_train_config_yaml(cfg.model.train_config_path) - cfg.model.action_horizon = _read_action_horizon_flow(train_yml) + # The qwen flow path uses task.action_horizon_flow. + from wall_x.config.loader import load_train_config_from_checkpoint + + train_cfg = load_train_config_from_checkpoint( + cfg.model.checkpoint_path, cfg.model.train_config_path + ) + cfg.model.action_horizon = int(train_cfg.task.action_horizon_flow) return cfg diff --git a/wall_x/_vendor/harrix/drivers/__init__.py b/wall_x/_vendor/harrix/e2e_infer/drivers/__init__.py similarity index 100% rename from wall_x/_vendor/harrix/drivers/__init__.py rename to wall_x/_vendor/harrix/e2e_infer/drivers/__init__.py diff --git a/wall_x/_vendor/harrix/e2e_infer/drivers/inproc/__init__.py b/wall_x/_vendor/harrix/e2e_infer/drivers/inproc/__init__.py new file mode 100644 index 00000000..f94c6f3f --- /dev/null +++ b/wall_x/_vendor/harrix/e2e_infer/drivers/inproc/__init__.py @@ -0,0 +1,227 @@ +"""In-process driver: DirectModelHandle + N x SubprocEnvHandle. + +The model lives in the driver process while simulator envs live in subprocesses. +Evaluation proceeds in chunk-level lockstep: + 1. Seed the driver process and construct the model handle. + 2. Start one subprocess env handle per worker. + 3. Claim a task-local frame from JobState. + 4. Reset envs, batch active observations, run model.predict_batch, then + fan out action chunks to env subprocesses. + 5. Complete all episodes in the frame and move to the next frame. +""" + +from __future__ import annotations + +import json +import logging +import os +import sys +import time + +from wall_x._vendor.harrix.e2e_infer.config import EvalConfig +from wall_x._vendor.harrix.e2e_infer.drivers.inproc.env_handle import DirectEnvHandle, SubprocEnvHandle +from wall_x._vendor.harrix.e2e_infer.drivers.inproc.model_handle import DirectModelHandle +from wall_x._vendor.harrix.e2e_infer.drivers.job_state import JobState + +logger = logging.getLogger(__name__) + + +def run(cfg: EvalConfig) -> None: + from wall_x._vendor.harrix.utils.seed import set_seed_everywhere + + cfg.runtime.log_dir = os.path.abspath(cfg.runtime.log_dir) + + import wall_x._vendor.harrix.e2e_infer.envs # noqa: F401 trigger env register + import wall_x._vendor.harrix.e2e_infer.adapters # noqa: F401 trigger adapter register + from wall_x._vendor.harrix.e2e_infer.envs.registry import enumerate_episodes_for + + # 1) Seed the driver process. + set_seed_everywhere(cfg.env.seed) + + # 2) Build DirectModelHandle; the model is loaded in the driver process. + logger.info("Constructing DirectModelHandle in the driver process") + t_model = time.time() + model_handle = DirectModelHandle(cfg) + logger.info("Model loaded in %.1fs", time.time() - t_model) + + # 3) Start env handles: inline in the driver process (DirectEnvHandle) or in + # spawned subprocesses (SubprocEnvHandle). + env_handle_cls = DirectEnvHandle if cfg.runtime.inline_env else SubprocEnvHandle + logger.info( + "Starting %s %s(s)", cfg.runtime.num_workers, env_handle_cls.__name__ + ) + env_handles = [ + env_handle_cls(cfg, worker_id=i) for i in range(cfg.runtime.num_workers) + ] + completed_ok = False + + # 4) JobState in FIFO mode: a worker claims the next pending episode the + # instant it finishes, so the inference batch stays full (continuous + # batching) instead of collapsing onto the slowest episode of a lockstep + # frame. This is what makes batched inference actually pay off. + os.makedirs(cfg.runtime.log_dir, exist_ok=True) + log_path = os.path.join(cfg.runtime.log_dir, "state.jsonl") + report_path = os.path.join(cfg.runtime.log_dir, "report.json") + episodes = enumerate_episodes_for(cfg) + logger.info("env.type=%r; scheduled %s episodes", cfg.env.type, len(episodes)) + state = JobState(episodes, log_path, batch_sync_mode=False) + + try: + # 5) Continuous-batching main loop. + t0 = time.time() + _run_continuous(state, model_handle, env_handles, cfg) + + elapsed = time.time() - t0 + logger.info( + "All episodes finished in %.1fs (%.1f min)", elapsed, elapsed / 60 + ) + + state.dump_final(report_path) + with open(report_path) as f: + report = json.load(f) + overall = report["overall"] + logger.info( + "attempted=%s, successes=%s, success_rate=%.2f%%, failed=%s", + overall["attempted"], + overall["successes"], + overall["success_rate"] * 100, + overall["failed"], + ) + logger.info("Report: %s", report_path) + logger.info("State log: %s", log_path) + # Exit non-zero when any episode errored (crash/infra), so an + # orchestrator can distinguish a degraded shard from a genuinely low + # score. Plain task failures (completed, success=False) keep exit 0. + exit_code = 0 if overall["failed"] == 0 else 1 + completed_ok = True + finally: + for h in env_handles: + h.shutdown() + model_handle.shutdown() + + +def _run_continuous( + state: JobState, + model_handle: DirectModelHandle, + env_handles: list, + cfg: EvalConfig, +) -> None: + """Continuous-batching scheduler. + + Every env worker is kept busy: the instant a worker finishes its episode it + claims the next pending one, so the inference batch stays ~num_workers wide + until the queue drains. This is the key difference from the old lockstep + frame path, whose batch collapsed to 1 whenever one episode outran the rest. + + A just-finished worker's reset (submit_reset) is issued *before* the active + batch and its completion is collected *after* the batch's predict+execute, so + the expensive expert-filter reset overlaps inference instead of stalling it. + + Requires the env to accept episodes of different tasks across resets when the + schedule mixes tasks. + """ + from wall_x._vendor.harrix.e2e_infer.envs.registry import encode_observation_for, max_infer_rounds_for + + n = len(env_handles) + max_rounds = max_infer_rounds_for(cfg) + + def _blank(): + return { + "ep": None, "obs": None, "instr": "", "steps": 0, "rounds": 0, + "task_desc": "", "phase": "idle", + } + + W = [_blank() for _ in range(n)] + t_last = time.time() + done_count = 0 + + def _finish(i: int, success: bool, error: str | None = None) -> None: + nonlocal done_count + w = W[i] + if error is None: + # Let the env finalize the just-finished episode (e.g. save a + # rollout MP4). No-op unless rollout capture is enabled. + try: + env_handles[i].submit_finalize_episode(success) + env_handles[i].wait_finalize_episode() + except Exception as exc: + logger.warning("finalize_episode failed for worker %d: %s", i, exc) + res = {"success": bool(success), "steps": w["steps"], "task_desc": w["task_desc"]} + if error is not None: + res["_error"] = error + state.fail(w["ep"], str(error)) + else: + state.complete(w["ep"], res) + W[i] = _blank() + done_count += 1 + + while True: + # (a) Refill idle workers from the queue; submit resets async so they + # overlap the active batch below. + for i in range(n): + if W[i]["phase"] != "idle": + continue + ep = state.claim(i) + if ep is None: + continue + W[i]["ep"] = ep + try: + env_handles[i].submit_reset(tuple(ep)) + W[i]["phase"] = "resetting" + except Exception as e: + _finish(i, False, str(e)) + + # (b) Batch all active workers and advance each one chunk. + active = [i for i in range(n) if W[i]["phase"] == "active"] + if active: + payloads = [ + { + "observation": encode_observation_for(cfg, W[i]["obs"]), + "instruction": W[i]["instr"], + "noise": None, + } + for i in active + ] + chunks = model_handle.predict_batch(payloads) + for k, i in enumerate(active): + env_handles[i].submit_execute_chunk(chunks[k]) + for k, i in enumerate(active): + try: + r = env_handles[i].wait_execute_chunk() + except Exception as e: + _finish(i, False, str(e)) + continue + W[i]["obs"] = r["obs"] + W[i]["steps"] += r["steps"] + W[i]["rounds"] += 1 + if r["done"]: + _finish(i, True) + elif W[i]["rounds"] >= max_rounds: + _finish(i, False) # hit the chunk cap: unsuccessful + + # (c) Collect the resets from (a); they overlapped the batch in (b). + for i in range(n): + if W[i]["phase"] != "resetting": + continue + try: + init = env_handles[i].wait_reset() + W[i].update( + obs=init["obs"], instr=init["instruction"], + task_desc=init.get("task_desc", ""), + phase="active", steps=0, rounds=0, + ) + except Exception as e: + _finish(i, False, str(e)) + + now = time.time() + if now - t_last >= 30.0: + in_flight = sum(1 for i in range(n) if W[i]["phase"] != "idle") + logger.info( + "continuous: batch=%s in_flight=%s done=%s %s", + len(active), in_flight, done_count, state.progress(), + ) + t_last = now + + # (d) Done when nothing is in flight and the queue is drained. + if not any(W[i]["phase"] != "idle" for i in range(n)) and state.is_drained(): + break diff --git a/wall_x/_vendor/harrix/drivers/inproc/env_handle.py b/wall_x/_vendor/harrix/e2e_infer/drivers/inproc/env_handle.py similarity index 70% rename from wall_x/_vendor/harrix/drivers/inproc/env_handle.py rename to wall_x/_vendor/harrix/e2e_infer/drivers/inproc/env_handle.py index 64bbec55..e8e57d43 100644 --- a/wall_x/_vendor/harrix/drivers/inproc/env_handle.py +++ b/wall_x/_vendor/harrix/e2e_infer/drivers/inproc/env_handle.py @@ -13,25 +13,58 @@ import numpy as np +class DirectEnvHandle: + def __init__(self, cfg, worker_id: int): + import wall_x._vendor.harrix.e2e_infer.envs # noqa: F401 + from wall_x._vendor.harrix.e2e_infer.envs.registry import build_env + + self._env = build_env(cfg, worker_id) + self._pending = None + + def submit_reset(self, ep_id): + self._pending = self._env.reset_episode(tuple(ep_id)) + + def wait_reset(self): + result = self._pending + self._pending = None + return result + + def submit_execute_chunk(self, actions): + self._pending = self._env.execute_chunk(np.asarray(actions, dtype=np.float32)) + + def wait_execute_chunk(self): + result = self._pending + self._pending = None + return result + + def submit_finalize_episode(self, success: bool): + # Inline handle: run the (infrequent) episode finalize immediately. + self._env.finalize_episode(bool(success)) + + def wait_finalize_episode(self): + return None + + def shutdown(self) -> None: + self._env.shutdown() + + def _subproc_main(cfg, worker_id, child_conn): """Env subprocess entry point.""" - # Spawned children inherit env vars, but set these explicitly for launchers - # that did not configure them. Robosuite validates the EGL id against the - # CUDA_VISIBLE_DEVICES environment string, so keep the same visible id here. - cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") or "0" - os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices - # EGL device index is always 0 after CUDA_VISIBLE_DEVICES remapping. - os.environ["MUJOCO_EGL_DEVICE_ID"] = os.environ.get("MUJOCO_EGL_DEVICE_ID") or "0" - - # Env workers run robosuite only. Do not import torch here; otherwise many - # subprocesses may initialize CUDA contexts and compete with the driver model. + # LIBERO uses offscreen mujoco-EGL only (no CUDA). Drop the inherited + # ``CUDA_VISIBLE_DEVICES``: with a non-zero CVD robosuite's pre-import + # ``MUJOCO_EGL_DEVICE_ID in CUDA_VISIBLE_DEVICES`` substring check and EGL's + # own ``0..len(devices)-1`` range check are mutually exclusive. Unsetting + # lets EGL default to device 0. + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + + # LIBERO workers stay CUDA-free. seed = cfg.env.seed + worker_id np.random.seed(seed) random.seed(seed) os.environ["PYTHONHASHSEED"] = str(seed) - import wall_x._vendor.harrix.envs # noqa: F401 trigger register - from wall_x._vendor.harrix.envs.registry import build_env + import wall_x._vendor.harrix.e2e_infer.envs # noqa: F401 trigger register + from wall_x._vendor.harrix.e2e_infer.envs.registry import build_env env = build_env(cfg, worker_id) try: @@ -119,16 +152,11 @@ def wait_execute_chunk(self): return self._recv() def submit_finalize_episode(self, success: bool): - self._send("finalize_episode", success) + self._send("finalize_episode", bool(success)) def wait_finalize_episode(self): return self._recv() - def finalize_episode(self, success: bool): - """Blocking convenience wrapper.""" - self.submit_finalize_episode(success) - return self.wait_finalize_episode() - def reset_episode(self, ep_id): """Blocking convenience wrapper.""" self.submit_reset(ep_id) diff --git a/wall_x/_vendor/harrix/drivers/inproc/model_handle.py b/wall_x/_vendor/harrix/e2e_infer/drivers/inproc/model_handle.py similarity index 77% rename from wall_x/_vendor/harrix/drivers/inproc/model_handle.py rename to wall_x/_vendor/harrix/e2e_infer/drivers/inproc/model_handle.py index ada020b6..7451a1e4 100644 --- a/wall_x/_vendor/harrix/drivers/inproc/model_handle.py +++ b/wall_x/_vendor/harrix/e2e_infer/drivers/inproc/model_handle.py @@ -9,8 +9,8 @@ class DirectModelHandle: def __init__(self, cfg): # Trigger adapter registration. - import wall_x._vendor.harrix.adapters # noqa: F401 - from wall_x._vendor.harrix.adapters.registry import build_adapter + import wall_x._vendor.harrix.e2e_infer.adapters # noqa: F401 + from wall_x._vendor.harrix.e2e_infer.adapters.registry import build_adapter self._adapter = build_adapter(cfg) diff --git a/wall_x/_vendor/harrix/drivers/job_state.py b/wall_x/_vendor/harrix/e2e_infer/drivers/job_state.py similarity index 94% rename from wall_x/_vendor/harrix/drivers/job_state.py rename to wall_x/_vendor/harrix/e2e_infer/drivers/job_state.py index 0a9b11cf..91e2065d 100644 --- a/wall_x/_vendor/harrix/drivers/job_state.py +++ b/wall_x/_vendor/harrix/e2e_infer/drivers/job_state.py @@ -40,15 +40,16 @@ def __init__( self._batch_sync_mode = bool(batch_sync_mode) self._batch_size = int(batch_size) - completed_set = ( - self._load_completed(log_path) if os.path.exists(log_path) else set() + loaded_completed = ( + self._load_completed_results(log_path) if os.path.exists(log_path) else {} ) + completed_set = set(loaded_completed) remaining: list[tuple] = [ tuple(ep) for ep in all_episodes if tuple(ep) not in completed_set ] self._in_progress: dict[tuple, int] = {} - self._completed: dict[tuple, dict] = {} + self._completed: dict[tuple, dict] = dict(loaded_completed) self._failed: dict[tuple, str] = {} if self._batch_sync_mode: @@ -101,12 +102,14 @@ def _build_frames(remaining: list[tuple], batch_size: int) -> list[dict]: cur_bucket.append(ep) if cur_bucket: for i in range(0, len(cur_bucket), batch_size): - frames.append({"eps": cur_bucket[i : i + batch_size], "next_slot": 0}) + frames.append( + {"eps": cur_bucket[i : i + batch_size], "next_slot": 0} + ) return frames @staticmethod - def _load_completed(log_path: str) -> set[tuple]: - completed = set() + def _load_completed_results(log_path: str) -> dict[tuple, dict]: + completed = {} with open(log_path, "r") as f: for line in f: line = line.strip() @@ -117,7 +120,12 @@ def _load_completed(log_path: str) -> set[tuple]: except json.JSONDecodeError: continue if rec.get("event") == "completed" and "ep_id" in rec: - completed.add(tuple(rec["ep_id"])) + ep = tuple(rec["ep_id"]) + completed[ep] = { + k: v + for k, v in rec.items() + if k not in ("event", "ep_id", "ts", "worker", "frame") + } return completed def _append_log(self, extra: dict): diff --git a/wall_x/_vendor/harrix/envs/__init__.py b/wall_x/_vendor/harrix/e2e_infer/envs/__init__.py similarity index 52% rename from wall_x/_vendor/harrix/envs/__init__.py rename to wall_x/_vendor/harrix/e2e_infer/envs/__init__.py index 605263f5..0786679a 100644 --- a/wall_x/_vendor/harrix/envs/__init__.py +++ b/wall_x/_vendor/harrix/e2e_infer/envs/__init__.py @@ -3,6 +3,6 @@ _LIBERO_IMPORT_ERROR = None try: - from wall_x._vendor.harrix.envs import libero # noqa: F401 -except ModuleNotFoundError as exc: + from wall_x._vendor.harrix.e2e_infer.envs import libero # noqa: F401 +except ImportError as exc: _LIBERO_IMPORT_ERROR = exc diff --git a/wall_x/_vendor/harrix/envs/base.py b/wall_x/_vendor/harrix/e2e_infer/envs/base.py similarity index 79% rename from wall_x/_vendor/harrix/envs/base.py rename to wall_x/_vendor/harrix/e2e_infer/envs/base.py index 5c3aef33..31dce68e 100644 --- a/wall_x/_vendor/harrix/envs/base.py +++ b/wall_x/_vendor/harrix/e2e_infer/envs/base.py @@ -18,7 +18,7 @@ import numpy as np -from wall_x._vendor.harrix.eval_config import EvalConfig +from wall_x._vendor.harrix.e2e_infer.config import EvalConfig class BaseEnv(ABC): @@ -46,6 +46,25 @@ def robot_spec(self) -> dict: - norm_key: str """ + @classmethod + def encode_observation(cls, raw_obs: dict) -> dict: + """Encode a raw env observation into the adapter payload bundle. + + Defaults to identity. Envs may override either to build the bundle here + (LIBERO) or to keep identity because the bundle was already packed during + reset/execute. + """ + return raw_obs + + @classmethod + def max_infer_rounds(cls, cfg: EvalConfig) -> int: + """Max model chunks per episode for the chunk-level (in-process) driver. + + Distinct from the instance-level ``_max_infer_rounds`` used by the + ``run_episode`` loop; envs targeting the in-process driver override this. + """ + return 1 + # ---- chunk-level primitives ---- @abstractmethod @@ -62,6 +81,9 @@ def execute_chunk(self, actions: np.ndarray) -> dict: Returns {"obs": dict, "done": bool, "steps": int}. """ + def finalize_episode(self, success: bool) -> None: + """Run env-specific cleanup after an episode.""" + def shutdown(self) -> None: """Release resources. Subclasses may override.""" @@ -79,8 +101,6 @@ def run_episode( """ import time - from wall_x._vendor.harrix.envs.libero_common import encode_raw_obs - t_ep_start = time.time() initial = self.reset_episode(ep_id) obs = initial["obs"] @@ -91,7 +111,7 @@ def run_episode( success = False steps_total = 0 for round_idx in range(max_rounds): - encoded = encode_raw_obs(obs) + encoded = self.encode_observation(obs) chunk = predict(encoded, instruction, round_idx) result = self.execute_chunk(chunk) obs = result["obs"] @@ -108,9 +128,6 @@ def run_episode( "task_desc": task_desc, } - def finalize_episode(self, success: bool) -> None: - """Hook for env-specific cleanup after an episode (e.g. save rollouts).""" - def _max_infer_rounds(self) -> int: """Return the maximum number of model chunks for one episode.""" raise NotImplementedError( diff --git a/wall_x/_vendor/harrix/envs/libero.py b/wall_x/_vendor/harrix/e2e_infer/envs/libero.py similarity index 71% rename from wall_x/_vendor/harrix/envs/libero.py rename to wall_x/_vendor/harrix/e2e_infer/envs/libero.py index 885d51ee..d6a78d1a 100644 --- a/wall_x/_vendor/harrix/envs/libero.py +++ b/wall_x/_vendor/harrix/e2e_infer/envs/libero.py @@ -11,17 +11,17 @@ import gc import logging import os +import time from typing import Optional import numpy as np -from wall_x._vendor.harrix.envs.base import BaseEnv -from wall_x._vendor.harrix.envs.libero_common import ( +from wall_x._vendor.harrix.e2e_infer.envs.base import BaseEnv +from wall_x._vendor.harrix.e2e_infer.envs.libero_common import ( get_rollout_frame, - model_action_to_libero_env, save_rollout_video, ) -from wall_x._vendor.harrix.envs.libero_sim import ( +from wall_x._vendor.harrix.e2e_infer.envs.libero_sim import ( create_libero_engine, find_image_observables, get_instruction, @@ -32,8 +32,8 @@ resolve_task_info, set_render_enabled, ) -from wall_x._vendor.harrix.envs.registry import register_env -from wall_x._vendor.harrix.eval_config import EvalConfig +from wall_x._vendor.harrix.e2e_infer.envs.registry import register_env +from wall_x._vendor.harrix.e2e_infer.config import EvalConfig logger = logging.getLogger(__name__) @@ -77,6 +77,13 @@ def __init__(self, cfg: EvalConfig, worker_id: int) -> None: if self._bit_dump_dir: os.makedirs(self._bit_dump_dir, exist_ok=True) + # Per-episode state consumed by execute_chunk. + self._current_ep: Optional[tuple] = None + self._current_task_desc: str = "" + self._current_instruction: str = "" + self._chunk_counter_in_ep: int = 0 + self._last_obs_for_dump: Optional[dict] = None + rollout_dir = (self._libero_cfg.rollout_dir or "").strip() if not rollout_dir: rollout_dir = os.environ.get("WALLX_ROLLOUT_DIR", "").strip() @@ -84,25 +91,18 @@ def __init__(self, cfg: EvalConfig, worker_id: int) -> None: rollout_dir = "" self._rollout_dir = rollout_dir or None self._rollout_fps = int(self._libero_cfg.rollout_fps) + self._replay_images: list = [] + self._rollout_saved = False if self._rollout_dir: - worker_subdir = f"worker{worker_id}" if cfg.runtime.num_workers > 1 else "" + worker_subdir = ( + f"worker{worker_id}" if cfg.runtime.num_workers > 1 else "" + ) self._rollout_dir = os.path.join( - self._rollout_dir, - self._task_suite_name, - worker_subdir, + self._rollout_dir, self._task_suite_name, worker_subdir ) os.makedirs(self._rollout_dir, exist_ok=True) logger.info("Rollout MP4 saving enabled: %s", self._rollout_dir) - # Per-episode state consumed by execute_chunk. - self._current_ep: Optional[tuple] = None - self._current_task_desc: str = "" - self._current_instruction: str = "" - self._chunk_counter_in_ep: int = 0 - self._last_obs_for_dump: Optional[dict] = None - self._replay_images: list[np.ndarray] = [] - self._rollout_saved = False - # ---- BaseEnv API ---- @classmethod @@ -122,6 +122,16 @@ def enumerate_episodes(cls, cfg: EvalConfig) -> list[tuple]: eps.append((suite, tid, epi)) return eps + @classmethod + def encode_observation(cls, raw_obs: dict) -> dict: + from wall_x._vendor.harrix.e2e_infer.envs.libero_common import encode_raw_obs + + return encode_raw_obs(raw_obs) + + @classmethod + def max_infer_rounds(cls, cfg: EvalConfig) -> int: + return cfg.env.libero.max_infer_times + @property def robot_spec(self) -> dict: """Return robot spec for driver/adapter validation.""" @@ -129,12 +139,12 @@ def robot_spec(self) -> dict: load_train_config_with_ckpt_overlay, ) - train_cfg = load_train_config_with_ckpt_overlay( + typed_cfg, _ = load_train_config_with_ckpt_overlay( self._cfg.model.train_config_path, self._cfg.model.checkpoint_path, ) return { - "dof_layout": train_cfg.get("dof_config", {}), + "dof_layout": typed_cfg.task.dof_config, "cam_names": list(self._cfg.model.cam_names), "norm_key": self._cfg.model.norm_key, } @@ -149,7 +159,9 @@ def reset_episode(self, ep_id: tuple) -> dict: f"env bound to suite={self._task_suite_name}, got ep with suite={suite}" ) - need_rebuild = self._rebuild_env_per_episode or self._current_task_id != task_id + need_rebuild = ( + self._rebuild_env_per_episode or self._current_task_id != task_id + ) if need_rebuild: self._rebuild_env(task_id) @@ -202,6 +214,7 @@ def execute_chunk(self, actions: np.ndarray) -> dict: self._chunk_counter_in_ep += 1 skip_render = self._chunk_skip_render() + # Disable rendering at the chunk start when render-skip is enabled. if skip_render and self._image_obs: if self._render_enabled_state: @@ -222,7 +235,7 @@ def execute_chunk(self, actions: np.ndarray) -> dict: set_render_enabled(self._image_obs, True) self._render_enabled_state = True - action = model_action_to_libero_env(actions[step_idx].reshape(-1)) + action = actions[step_idx].reshape(-1) obs, _, done_flag, _ = self._libero_env.step(action) last_obs = obs steps += 1 @@ -231,7 +244,11 @@ def execute_chunk(self, actions: np.ndarray) -> dict: done = True # Do not add an extra simulator step on early success; just make # sure future rendering is enabled. - if skip_render and self._image_obs and not self._render_enabled_state: + if ( + skip_render + and self._image_obs + and not self._render_enabled_state + ): set_render_enabled(self._image_obs, True) self._render_enabled_state = True break @@ -239,9 +256,6 @@ def execute_chunk(self, actions: np.ndarray) -> dict: self._last_obs_for_dump = last_obs return {"obs": last_obs, "done": done, "steps": steps} - def finalize_episode(self, success: bool) -> None: - self._save_episode_rollout(success) - def shutdown(self) -> None: if self._libero_env is not None: try: @@ -250,13 +264,13 @@ def shutdown(self) -> None: pass self._libero_env = None - # ---- internals ---- + def finalize_episode(self, success: bool) -> None: + self._save_episode_rollout(success) def _chunk_skip_render(self) -> bool: - """Skip intermediate renders unless rollout MP4 saving needs every frame.""" return self._effective_skip_render and self._rollout_dir is None - def _begin_rollout_capture(self, obs: dict | None) -> None: + def _begin_rollout_capture(self, obs: Optional[dict]) -> None: self._replay_images = [] self._rollout_saved = False if self._rollout_dir and obs is not None: @@ -265,7 +279,7 @@ def _begin_rollout_capture(self, obs: dict | None) -> None: self._render_enabled_state = True self._replay_images.append(get_rollout_frame(obs)) - def _append_rollout_frame(self, obs: dict | None) -> None: + def _append_rollout_frame(self, obs: Optional[dict]) -> None: if self._rollout_dir and obs is not None: self._replay_images.append(get_rollout_frame(obs)) @@ -299,6 +313,93 @@ def _save_episode_rollout(self, success: bool) -> None: exc, ) + # ---- episode-level rollout override ---- + + def run_episode(self, ep_id, predict): + """Run a LIBERO episode with client-side action buffering.""" + from collections import deque + + from wall_x._vendor.harrix.e2e_infer.envs.libero_common import encode_raw_obs + + suite, task_id, ep_idx = ep_id + if suite != self._task_suite_name: + raise ValueError( + f"env bound to suite={self._task_suite_name}, got ep with suite={suite}" + ) + + need_rebuild = ( + self._rebuild_env_per_episode or self._current_task_id != task_id + ) + if need_rebuild: + self._rebuild_env(task_id) + + task_desc, default_states = resolve_task_info(self._task_suite, task_id) + init_state = pick_initial_state( + self._libero_cfg.initial_states_path, + self._custom_initial_states, + task_desc, + default_states, + ep_idx, + ) + + t_ep_start = time.time() + if not need_rebuild: + self._libero_env.reset() + obs = self._libero_env.set_init_state(init_state) + if obs is None: + raise RuntimeError("set_init_state returned None") + + set_render_enabled(self._image_obs, True) + render_enabled_state = True + + dummy_action = get_libero_dummy_action() + for _ in range(10): + obs, _, _, _ = self._libero_env.step(dummy_action) + + chunk_buffer: deque = deque() + instruction = get_instruction(task_desc) + max_steps = ( + self._libero_cfg.max_infer_times * self._cfg.model.action_horizon + ) + success = False + step_count = 0 + chunk_counter_in_ep = 0 + for step in range(max_steps): + if not chunk_buffer: + encoded = encode_raw_obs(obs) + chunk = predict(encoded, instruction, step) + if self._bit_dump_dir: + self._dump_chunk_npz(ep_id, chunk_counter_in_ep, obs, chunk) + chunk_counter_in_ep += 1 + for a in chunk: + chunk_buffer.append(np.asarray(a, dtype=np.float32)) + + action = chunk_buffer.popleft().reshape(-1) + + if self._effective_skip_render and self._image_obs: + next_needs_forward = len(chunk_buffer) == 0 + if (not self._chunk_granular_render_toggle) or ( + next_needs_forward != render_enabled_state + ): + set_render_enabled(self._image_obs, next_needs_forward) + render_enabled_state = next_needs_forward + + obs, _, done, _ = self._libero_env.step(action) + step_count = step + 1 + if bool(done): + success = True + break + + elapsed = time.time() - t_ep_start + return { + "success": bool(success), + "steps": step_count, + "elapsed_sec": round(elapsed, 3), + "task_desc": task_desc, + } + + # ---- internals ---- + def _rebuild_env(self, task_id: int) -> None: if self._libero_env is not None: try: diff --git a/wall_x/_vendor/harrix/envs/libero_common.py b/wall_x/_vendor/harrix/e2e_infer/envs/libero_common.py similarity index 50% rename from wall_x/_vendor/harrix/envs/libero_common.py rename to wall_x/_vendor/harrix/e2e_infer/envs/libero_common.py index 2bfa2f35..8829a636 100644 --- a/wall_x/_vendor/harrix/envs/libero_common.py +++ b/wall_x/_vendor/harrix/e2e_infer/envs/libero_common.py @@ -14,93 +14,6 @@ import numpy as np -# Fallback single-arm LIBERO dof_config when train_config does not provide one. -_LIBERO_FALLBACK_DOF_CONFIG = { - "follow_right_ee_cartesian_pos": 3, - "follow_right_ee_rotation": 3, - "follow_right_gripper": 1, -} - -_LIBERO_FALLBACK_AGENT_POS_CONFIG = dict(_LIBERO_FALLBACK_DOF_CONFIG) -_VIRTUAL_TAIL_KEYS = frozenset(("action_padding",)) - - -def _resolve_dof_config(train_config: dict) -> dict: - return ( - train_config.get("dof_config") - or train_config.get("task", {}).get("dof_config") - or _LIBERO_FALLBACK_DOF_CONFIG - ) - - -def _resolve_agent_pos_config(train_config: dict) -> dict: - return ( - train_config.get("agent_pos_config") - or train_config.get("task", {}).get("agent_pos_config") - or _LIBERO_FALLBACK_AGENT_POS_CONFIG - ) - - -def _move_virtual_keys_to_tail(layout: dict) -> dict: - head = {k: v for k, v in layout.items() if k not in _VIRTUAL_TAIL_KEYS} - tail = {k: v for k, v in layout.items() if k in _VIRTUAL_TAIL_KEYS} - return {**head, **tail} - - -def _effective_agent_pos_config(train_config: dict, state_values: dict) -> dict: - config = _move_virtual_keys_to_tail(dict(_resolve_agent_pos_config(train_config))) - gripper_key = next( - ( - key - for key in config - if key.replace("follow_", "").replace("master_", "") == "right_gripper" - ), - None, - ) - if gripper_key is None: - return config - - old_dim = int(config[gripper_key]) - target_dim = None - override = os.environ.get("WALLX_LIBERO_STATE_GRIPPER_DIM") - if override: - target_dim = int(override) - elif os.environ.get("WALLX_LIBERO_AUTO_STATE_GRIPPER_DIM", "1") != "0": - norm_dim = int(train_config.get("_libero_proprio_norm_dim") or 0) - real_dim = sum(v for k, v in config.items() if k not in _VIRTUAL_TAIL_KEYS) - if norm_dim == real_dim + 1: - target_dim = old_dim + 1 - - available_dim = state_values["right_gripper"].shape[1] - if target_dim is None or target_dim == old_dim or target_dim > available_dim: - return config - - config[gripper_key] = target_dim - delta = target_dim - old_dim - if "action_padding" in config: - config["action_padding"] = max(0, int(config["action_padding"]) - delta) - return config - - -def _build_right_arm_state_values(obs_ndarrays: dict) -> dict[str, np.ndarray]: - """Bare-key state tensors for single-arm LIBERO proprio construction.""" - rot3 = np.asarray(obs_ndarrays["eef_axisangle"], dtype=np.float32).reshape(1, 3) - values: dict[str, np.ndarray] = { - "right_ee_cartesian_pos": np.asarray( - obs_ndarrays["eef_pos"], dtype=np.float32 - ).reshape(1, 3), - "right_ee_rotation": rot3, - "right_gripper": np.asarray(obs_ndarrays["gripper"], dtype=np.float32).reshape( - 1, -1 - ), - } - from wall_x._vendor.x2robot_utils.geometry import euler_to_matrix_zyx_6d_nb - - rot6d = euler_to_matrix_zyx_6d_nb(rot3.astype(np.float64)).reshape(1, 6) - values["right_ee_rotation_6D"] = rot6d.astype(np.float32) - return values - - # Auxiliary action keys that should be masked out for single-arm LIBERO. _DOF_MASK_ZERO_KEYS = frozenset( ( @@ -126,14 +39,54 @@ def _get_libero_image(obs: dict) -> np.ndarray: return obs["agentview_image"][::-1, ::-1] +def _get_libero_wrist_image(obs: dict) -> np.ndarray: + """Return the wrist camera image, rotated to match preprocessing.""" + return obs["robot0_eye_in_hand_image"][::-1, ::-1] + + def get_rollout_frame(obs: dict) -> np.ndarray: - """Return one RGB frame for rollout MP4 saving.""" + """Return one third-person RGB frame for rollout MP4 saving.""" return np.asarray(_get_libero_image(obs), dtype=np.uint8) -def _get_libero_wrist_image(obs: dict) -> np.ndarray: - """Return the wrist camera image, rotated to match preprocessing.""" - return obs["robot0_eye_in_hand_image"][::-1, ::-1] +def _sanitize_task_description(task_description: str, max_len: int = 50) -> str: + return ( + task_description.lower() + .replace(" ", "_") + .replace("\n", "_") + .replace(".", "_")[:max_len] + ) + + +def save_rollout_video( + rollout_dir: str, + rollout_images: list[np.ndarray], + *, + task_id: int, + episode_idx: int, + success: bool, + task_description: str, + fps: int = 30, +) -> "str | None": + """Save an MP4 replay of one LIBERO episode; returns the path (or None).""" + if not rollout_images: + return None + + import imageio + + os.makedirs(rollout_dir, exist_ok=True) + task_slug = _sanitize_task_description(task_description) + mp4_path = os.path.join( + rollout_dir, + f"task{task_id}_ep{episode_idx}--success={int(success)}--{task_slug}.mp4", + ) + writer = imageio.get_writer(mp4_path, fps=fps, macro_block_size=1) + try: + for img in rollout_images: + writer.append_data(np.asarray(img, dtype=np.uint8)) + finally: + writer.close() + return mp4_path def _quat2axisangle(quat) -> np.ndarray: @@ -170,7 +123,9 @@ def encode_raw_obs(raw_obs: dict) -> dict: "eef_axisangle": np.asarray( _quat2axisangle(raw_obs["robot0_eef_quat"]), dtype=np.float32 ), - "gripper": np.asarray(raw_obs["robot0_gripper_qpos"], dtype=np.float32), + "gripper": np.asarray( + raw_obs["robot0_gripper_qpos"][0:1], dtype=np.float32 + ), "face_view": _get_libero_image(raw_obs), "wrist_view": _get_libero_wrist_image(raw_obs), } @@ -183,17 +138,32 @@ def encode_raw_obs(raw_obs: dict) -> dict: def encode_proprio( obs_ndarrays: dict, - train_config: dict, + task, action_horizon: int, ) -> dict: """Convert a single-arm LIBERO ndarray payload into model input fields. - Returned fields include proprioception, agent_pos_mask, dof_mask, - face_view, and right_wrist_view. + ``task`` is the typed ``TaskConfig`` (single source of the action + contract). Returned fields include proprioception, agent_pos_mask, + dof_mask, face_view, and right_wrist_view. """ - state_values = _build_right_arm_state_values(obs_ndarrays) - agent_pos_config = _effective_agent_pos_config(train_config, state_values) - dof_config = _move_virtual_keys_to_tail(dict(_resolve_dof_config(train_config))) + agent_pos_config = task.agent_pos_config + dof_config = task.dof_config + + # Canonical state values with bare keys. + state_values: dict[str, np.ndarray] = { + "right_ee_cartesian_pos": obs_ndarrays["eef_pos"][None], # (1, 3) + "right_ee_rotation": obs_ndarrays["eef_axisangle"][None], # (1, 3) + "right_gripper": obs_ndarrays["gripper"][None], # (1, 1) + } + from wall_x._vendor.x2robot_utils.geometry import euler_to_matrix_zyx_6d_nb + + rot3 = np.asarray( + obs_ndarrays["eef_axisangle"], dtype=np.float64 + ).reshape(1, 3) + state_values["right_ee_rotation_6D"] = ( + euler_to_matrix_zyx_6d_nb(rot3).reshape(1, 6).astype(np.float32) + ) propri_parts: list[np.ndarray] = [] mask_parts: list[np.ndarray] = [] @@ -201,8 +171,12 @@ def encode_proprio( bare = key.replace("follow_", "").replace("master_", "") if bare in state_values: v = state_values[bare] - if bare == "right_gripper" and v.shape[1] > dim: - v = v[:, :dim] + if bare == "right_gripper" and v.shape[1] < dim: + # The env emits one gripper qpos but some checkpoints train + # with both finger dims, which mirror each other + # (qpos[1] ~= -qpos[0]); mirror-pad instead of OOD zeros. + pad = np.repeat(-v[:, :1], dim - v.shape[1], axis=1) + v = np.concatenate([v, pad], axis=1).astype(np.float32) if v.shape[1] != dim: raise ValueError( f"agent_pos_config[{key!r}]={dim} does not match " @@ -241,12 +215,17 @@ def encode_proprio( # ============================================================ -def decode_chunk(predict_action: np.ndarray, train_config: dict) -> np.ndarray: - """Extract a 7-dof right-arm chunk from model action output.""" +def decode_chunk(predict_action: np.ndarray, task) -> np.ndarray: + """Extract a 7-dof right-arm chunk from model action output. + + ``task`` is the typed ``TaskConfig``. Only the first gripper dim is + kept: 2-dim-gripper checkpoints carry a mirrored second finger dim + (see ``encode_proprio``). + """ if predict_action.ndim == 3: predict_action = predict_action[0] - dof_config = _resolve_dof_config(train_config) + dof_config = task.dof_config slices: dict[str, slice] = {} start = 0 for key, dim in dof_config.items(): @@ -255,105 +234,6 @@ def decode_chunk(predict_action: np.ndarray, train_config: dict) -> np.ndarray: start += dim pos = predict_action[:, slices["right_ee_cartesian_pos"]] - grip = predict_action[:, slices["right_gripper"]] - - if "right_ee_rotation_6D" in slices: - from wall_x._vendor.x2robot_utils.geometry import so3_to_euler_zyx_batch_nb - - rot6d = np.asarray( - predict_action[:, slices["right_ee_rotation_6D"]], dtype=np.float64 - ) - rot = so3_to_euler_zyx_batch_nb(rot6d).astype(np.float32) - elif "right_ee_rotation" in slices: - rot = predict_action[:, slices["right_ee_rotation"]] - else: - raise KeyError( - "dof_config has no right-arm rotation slice " - f"(keys={list(slices.keys())})" - ) - + rot = predict_action[:, slices["right_ee_rotation"]] + grip = predict_action[:, slices["right_gripper"]][:, :1] return np.concatenate([pos, rot, grip], axis=1) - - -def gripper_model_to_libero_osc(grip_2d: np.ndarray) -> float: - """Map model gripper output to robosuite OSC_POSE gripper command in [-1, 1].""" - g = np.asarray(grip_2d, dtype=np.float64).reshape(-1) - if g.size == 0: - raise ValueError("empty gripper action") - cmd = float(g[0]) - if os.environ.get("WALLX_LIBERO_GRIPPER_BINARIZE", "0") == "1": - if abs(cmd) < 1e-6: - cmd = -1.0 - else: - cmd = float(np.sign(cmd)) - if os.environ.get("WALLX_LIBERO_INVERT_GRIPPER", "0") == "1": - cmd *= -1.0 - return cmd - - -def _sanitize_task_description(task_description: str, max_len: int = 50) -> str: - return ( - task_description.lower() - .replace(" ", "_") - .replace("\n", "_") - .replace(".", "_")[:max_len] - ) - - -def save_rollout_video( - rollout_dir: str, - rollout_images: list[np.ndarray], - *, - task_id: int, - episode_idx: int, - success: bool, - task_description: str, - fps: int = 30, -) -> str | None: - """Save an MP4 replay of one LIBERO episode.""" - if not rollout_images: - return None - - import imageio - - os.makedirs(rollout_dir, exist_ok=True) - task_slug = _sanitize_task_description(task_description) - mp4_path = os.path.join( - rollout_dir, - f"task{task_id}_ep{episode_idx}--success={int(success)}--{task_slug}.mp4", - ) - writer = imageio.get_writer(mp4_path, fps=fps, macro_block_size=1) - try: - for img in rollout_images: - writer.append_data(np.asarray(img, dtype=np.uint8)) - finally: - writer.close() - return mp4_path - - -def model_action_to_libero_env(action: np.ndarray) -> np.ndarray: - """Convert model output to the robosuite OSC_POSE 7D action. - - Internal LIBERO evaluation passes the model's 7D chunk directly to - ``env.step``. Keep that as the public default; conversion modes remain - available only for ablations through environment variables. - """ - from scipy.spatial.transform import Rotation as R - - a = np.asarray(action, dtype=np.float64).reshape(-1) - if a.size not in (7, 8): - raise ValueError(f"expected 7D or 8D model action, got shape {a.shape}") - - pos_delta = a[:3] - rot_mode = os.environ.get("WALLX_LIBERO_ROT_MODE", "direct").strip() - if rot_mode == "euler_zyx_to_rotvec": - rot_aa = R.from_euler("zyx", a[3:6]).as_rotvec() - elif rot_mode == "direct": - rot_aa = a[3:6] - else: - raise ValueError( - "WALLX_LIBERO_ROT_MODE must be 'euler_zyx_to_rotvec' or 'direct', " - f"got {rot_mode!r}" - ) - grip = gripper_model_to_libero_osc(a[6:8] if a.size == 8 else a[6:7]) - return np.concatenate([pos_delta, rot_aa, np.array([grip], dtype=np.float64)]) diff --git a/wall_x/_vendor/harrix/envs/libero_sim.py b/wall_x/_vendor/harrix/e2e_infer/envs/libero_sim.py similarity index 58% rename from wall_x/_vendor/harrix/envs/libero_sim.py rename to wall_x/_vendor/harrix/e2e_infer/envs/libero_sim.py index ae510465..9a424980 100644 --- a/wall_x/_vendor/harrix/envs/libero_sim.py +++ b/wall_x/_vendor/harrix/e2e_infer/envs/libero_sim.py @@ -12,43 +12,56 @@ from typing import Any, Optional import numpy as np -from robosuite.wrappers import VisualizationWrapper logger = logging.getLogger(__name__) # ============================================================ -# One-time side effect: auto-create ~/.libero/config.yaml so importing LIBERO -# does not trigger an interactive prompt. +# LIBERO config bootstrap. Called lazily (from the functions that actually use +# LIBERO), NOT at import time — importing this module only registers the env +# class and must not require LIBERO/robosuite or touch ~/.libero. Registration +# must stay import-safe so a driver can import harrix.e2e_infer.envs without +# pulling in LIBERO until a Libero env is actually instantiated. # ============================================================ def _ensure_libero_config() -> None: - import yaml as _yaml + """Create ``$LIBERO_CONFIG_PATH/config.yaml`` (default ``~/.libero``) if absent + so that importing ``libero.libero`` does not drop into its interactive + ``input()`` prompt — which raises EOFError under non-interactive/mpirun runs. + + Locate the package via ``importlib`` instead of ``import libero.libero``: + importing it is precisely what triggers the prompt when the config is missing, + so the old import-to-find-the-path approach could never bootstrap a fresh node. + Idempotent and safe to call on every LIBERO entry point. + """ + import importlib.util libero_config_path = os.environ.get( "LIBERO_CONFIG_PATH", os.path.expanduser("~/.libero") ) config_file = os.path.join(libero_config_path, "config.yaml") + if os.path.exists(config_file): + return - if not os.path.exists(config_file): - os.makedirs(libero_config_path, exist_ok=True) - import libero.libero as _libero_pkg - - benchmark_root = os.path.dirname(os.path.abspath(_libero_pkg.__file__)) - default_paths = { - "benchmark_root": benchmark_root, - "bddl_files": os.path.join(benchmark_root, "./bddl_files"), - "init_states": os.path.join(benchmark_root, "./init_files"), - "datasets": os.path.join(benchmark_root, "../datasets"), - "assets": os.path.join(benchmark_root, "./assets"), - } - with open(config_file, "w") as f: - _yaml.dump(default_paths, f) - logger.info("Auto-created LIBERO config: %s", config_file) + spec = importlib.util.find_spec("libero.libero") + if spec is None or not spec.origin: + return # LIBERO not installed; nothing to pre-configure + import yaml as _yaml -_ensure_libero_config() + benchmark_root = os.path.dirname(os.path.abspath(spec.origin)) + os.makedirs(libero_config_path, exist_ok=True) + default_paths = { + "benchmark_root": benchmark_root, + "bddl_files": os.path.join(benchmark_root, "./bddl_files"), + "init_states": os.path.join(benchmark_root, "./init_files"), + "datasets": os.path.join(benchmark_root, "../datasets"), + "assets": os.path.join(benchmark_root, "./assets"), + } + with open(config_file, "w") as f: + _yaml.dump(default_paths, f) + logger.info("Auto-created LIBERO config: %s", config_file) # ============================================================ @@ -58,6 +71,7 @@ def _ensure_libero_config() -> None: def get_task_suite(task_suite_name: str): """Load a LIBERO task suite.""" + _ensure_libero_config() from libero.libero import benchmark return benchmark.get_benchmark_dict()[task_suite_name]() @@ -78,6 +92,37 @@ def get_libero_dummy_action() -> list[float]: # ============================================================ +def _install_offsamples_hook(offsamples: int) -> None: + """Pin the MSAA sample count on the model before the render context is created. + + The MSAA resolve of EGL off-screen rendering is a source of + run-to-run non-determinism in LIBERO eval: the wrist camera's + first frame alone shows a few pixels of LSB jitter, which the + closed loop amplifies into success-rate noise. Pinning + ``offsamples=0`` makes full runs bit-for-bit reproducible. + + It must take effect *before* robosuite creates the render + context (the context sizes its MSAA buffer from + ``vis.quality.offsamples`` at creation time; changing the field + afterwards has no effect). Since the context is created during + ``OffScreenRenderEnv`` construction, the hook is installed on + ``MjRenderContext.__init__``. + """ + from robosuite.utils import binding_utils + + if getattr(binding_utils.MjRenderContext, "_wallx_offsamples_hook", None) == offsamples: + return + orig_init = binding_utils.MjRenderContext.__init__ + + def patched_init(self, sim, *args, **kwargs): + raw = getattr(sim.model, "_model", sim.model) + raw.vis.quality.offsamples = offsamples + return orig_init(self, sim, *args, **kwargs) + + binding_utils.MjRenderContext.__init__ = patched_init + binding_utils.MjRenderContext._wallx_offsamples_hook = offsamples + + def create_libero_engine( task_id: int, task_suite_name: str, @@ -85,9 +130,15 @@ def create_libero_engine( seed: int = 7, ) -> Any: """Construct one LIBERO robosuite engine.""" + _ensure_libero_config() + from robosuite.wrappers import VisualizationWrapper from libero.libero import get_libero_path from libero.libero.envs import OffScreenRenderEnv + offsamples = os.environ.get("WALLX_RENDER_OFFSAMPLES", "").strip() + if offsamples: + _install_offsamples_hook(int(offsamples)) + task_suite = get_task_suite(task_suite_name) task = task_suite.get_task(task_id) task_bddl_file = os.path.join( @@ -141,7 +192,9 @@ def pick_initial_state( return default_states[episode_idx] if custom_initial_states is None: - raise ValueError(f"custom initial states not loaded for {initial_states_path}") + raise ValueError( + f"custom initial states not loaded for {initial_states_path}" + ) key = task_desc.replace(" ", "_") ep_key = f"demo_{episode_idx}" record = custom_initial_states[key][ep_key] diff --git a/wall_x/_vendor/harrix/envs/registry.py b/wall_x/_vendor/harrix/e2e_infer/envs/registry.py similarity index 70% rename from wall_x/_vendor/harrix/envs/registry.py rename to wall_x/_vendor/harrix/e2e_infer/envs/registry.py index 78d455c3..16fc4e8f 100644 --- a/wall_x/_vendor/harrix/envs/registry.py +++ b/wall_x/_vendor/harrix/e2e_infer/envs/registry.py @@ -4,8 +4,8 @@ from typing import Type -from wall_x._vendor.harrix.eval_config import EvalConfig -from wall_x._vendor.harrix.envs.base import BaseEnv +from wall_x._vendor.harrix.e2e_infer.config import EvalConfig +from wall_x._vendor.harrix.e2e_infer.envs.base import BaseEnv _REGISTRY: dict[str, Type[BaseEnv]] = {} @@ -28,7 +28,7 @@ def _get_class(cfg: EvalConfig) -> Type[BaseEnv]: cls = _REGISTRY.get(t) if cls is None: if t == "libero": - import wall_x._vendor.harrix.envs as _envs + import wall_x._vendor.harrix.e2e_infer.envs as _envs exc = getattr(_envs, "_LIBERO_IMPORT_ERROR", None) if exc is not None: @@ -37,7 +37,9 @@ def _get_class(cfg: EvalConfig) -> Type[BaseEnv]: "Install LIBERO/robosuite and their simulator dependencies " "before using env.type='libero'." ) from exc - raise ValueError(f"unknown env type={t!r}, registered: {sorted(_REGISTRY)}") + raise ValueError( + f"unknown env type={t!r}, registered: {sorted(_REGISTRY)}" + ) return cls @@ -49,5 +51,13 @@ def enumerate_episodes_for(cfg: EvalConfig) -> list[tuple]: return _get_class(cfg).enumerate_episodes(cfg) +def encode_observation_for(cfg: EvalConfig, raw_obs: dict) -> dict: + return _get_class(cfg).encode_observation(raw_obs) + + +def max_infer_rounds_for(cfg: EvalConfig) -> int: + return _get_class(cfg).max_infer_rounds(cfg) + + def registered_envs() -> list[str]: return sorted(_REGISTRY) diff --git a/wall_x/_vendor/harrix/model_executors/__init__.py b/wall_x/_vendor/harrix/model_executors/__init__.py new file mode 100644 index 00000000..1f09366b --- /dev/null +++ b/wall_x/_vendor/harrix/model_executors/__init__.py @@ -0,0 +1,5 @@ +"""Model loading, preprocessing, execution and postprocessing.""" + +from wall_x._vendor.harrix.model_executors.base import ModelExecutor, ServingExecutor + +__all__ = ["ModelExecutor", "ServingExecutor"] diff --git a/wall_x/_vendor/harrix/model_executors/base.py b/wall_x/_vendor/harrix/model_executors/base.py new file mode 100644 index 00000000..38c02653 --- /dev/null +++ b/wall_x/_vendor/harrix/model_executors/base.py @@ -0,0 +1,26 @@ +"""Minimal model-executor contract.""" + +from __future__ import annotations + +from typing import Any, Protocol, Sequence + + +class ModelExecutor(Protocol): + """Model loading and batched execution capability.""" + + def infer_batch(self, requests: Sequence[Any]) -> list[Any]: + ... + + @property + def metadata(self) -> dict[str, Any]: + ... + + +class ServingExecutor(ModelExecutor, Protocol): + """Executor capability required by a stateful serving handler.""" + + def infer(self, request: Any) -> Any: + ... + + def reset(self) -> None: + ... diff --git a/wall_x/_vendor/harrix/adapters/qwen_vlact.py b/wall_x/_vendor/harrix/model_executors/qwen2_5_eval.py similarity index 74% rename from wall_x/_vendor/harrix/adapters/qwen_vlact.py rename to wall_x/_vendor/harrix/model_executors/qwen2_5_eval.py index a7f21183..463894f7 100644 --- a/wall_x/_vendor/harrix/adapters/qwen_vlact.py +++ b/wall_x/_vendor/harrix/model_executors/qwen2_5_eval.py @@ -1,15 +1,15 @@ -"""Shared adapter base for Qwen-VL action models. +"""Qwen2.5 Torch executor used by environment evaluation. -Variant-specific subclasses live in ``harrix.adapters.variants`` and provide -the Wall-X training adapter class via ``_training_adapter``. The shared base -handles checkpoint loading, normalizers, LIBERO observation encoding, prompt -construction, and flow-action inference. +The executor owns checkpoint/model/processor/normalizer initialization and +batched model execution. Environment adapters only select the training-model +adapter and translate executor results to environment actions. """ from __future__ import annotations import copy import logging +from typing import Any import numpy as np import torch @@ -22,24 +22,14 @@ preprocesser_call, ) -from wall_x._vendor.harrix.adapters.base import BaseInferAdapter -from wall_x._vendor.harrix.envs.libero_common import decode_chunk, encode_proprio -from wall_x._vendor.harrix.eval_config import EvalConfig -from wall_x._vendor.harrix.utils.ckpt_load import ( - load_state_dict, - resolve_checkpoint_dir, -) +from wall_x._vendor.harrix.utils.ckpt_load import load_state_dict from wall_x._vendor.harrix.utils.normalizer import build_normalizers from wall_x._vendor.harrix.utils.train_config import ( build_data_config, build_model_config, load_train_config_with_ckpt_overlay, - normalize_train_config_for_inference, - register_data_backend, - resolve_state_bins, - resolve_use_state_string_representation, ) -from wall_x.trainer.trainer_utils import load_wallx_processors +from wall_x.trainer.trainer_utils import load_wallx_processors_from_cfg _SUPPORTED_MODES = {"flow", "ar", "dllm", "vqa", "subtask"} @@ -93,19 +83,11 @@ def _normalize_real_prefix(normalizer, tensor, dataset_names): out[..., width:] = 0 return out -class QwenVLActInferAdapter(BaseInferAdapter): - """Shared constructor and batched flow inference implementation.""" - - # ---- subclass hook ---- - @classmethod - def _training_adapter(cls): - """Return the training-side ModelAdapter subclass.""" - raise NotImplementedError(f"{cls.__name__} must override _training_adapter()") +class Qwen2_5EvalExecutor: + """Checkpoint-backed Qwen2.5 batched evaluator.""" - # ---- ctor ---- - - def __init__(self, cfg: EvalConfig) -> None: + def __init__(self, cfg: Any, training_adapter) -> None: mode = cfg.model.action_mode if mode not in _SUPPORTED_MODES: raise ValueError( @@ -120,58 +102,50 @@ def __init__(self, cfg: EvalConfig) -> None: ) self._action_mode = mode - ta = self._training_adapter() + ta = training_adapter device = "cuda" self._device = device - self._checkpoint_path = resolve_checkpoint_dir(cfg.model.checkpoint_path) - # 1) train_config with checkpoint overlays, then data backend registration. - train_config = load_train_config_with_ckpt_overlay( - cfg.model.train_config_path, self._checkpoint_path + # 1) typed cfg through the single wall_x loader (registers the data + # backend, applies ckpt overlays). + typed_cfg, train_config = load_train_config_with_ckpt_overlay( + cfg.model.train_config_path, cfg.model.checkpoint_path ) - train_config = normalize_train_config_for_inference( - train_config, cfg.model.train_config_path + self._typed_cfg = typed_cfg + + # 2) data_config. Image resizing needs resolution/image_factor/min/max. + self._data_config = build_data_config(typed_cfg, train_config) + + # 3) HF model config + ConfigClass = ta.config_class() + self._model_config = build_model_config( + ConfigClass, cfg.model.checkpoint_path, typed_cfg, train_config ) - register_data_backend(train_config) - self._train_config = train_config - # 2) normalizers first - action_tokenizer setup needs them during processor load. + # 4) normalizers, including norm_key fallback. + # Built before load_wallx_processors: OSS wall-x (refactor_2604) requires + # an explicit normalizer for the action tokenizer; older builds accept it + # via the same kwarg. normalizer_action, normalizer_propri, resolved_norm_key = build_normalizers( - self._checkpoint_path, train_config, cfg.model.norm_key + cfg.model.checkpoint_path, typed_cfg, cfg.model.norm_key ) self._normalizer_action = normalizer_action self._normalizer_propri = normalizer_propri self._norm_key = resolved_norm_key - if normalizer_propri is not None and resolved_norm_key in normalizer_propri.delta: - train_config["_libero_proprio_norm_dim"] = int( - normalizer_propri.delta[resolved_norm_key].shape[0] - ) - - # 3) data_config. Image resizing needs resolution/image_factor/min/max. - self._data_config = build_data_config(cfg.model.train_config_path, train_config) - # 4) HF model config - ConfigClass = ta.config_class() - self._model_config = build_model_config( - ConfigClass, - self._checkpoint_path, - train_config, - cfg.model.train_config_path, - ) - - # 5) processor + tokenizer_mixin (may extend vocab via action_tokenizer) - procs = load_wallx_processors( - train_config, normalizer=normalizer_action, device=device + # 5) processor + tokenizer_mixin (from the typed cfg, so ckpt + # processor/tokenizer overlays are honored). + procs = load_wallx_processors_from_cfg( + typed_cfg, normalizer=normalizer_action, device=device ) self._processor = procs["processor"] self._tokenizer_mixin = procs.get("tokenizer_mixin") - logger.info( - "processor vocab size: %s (action_tokenizer_type=%r)", - len(self._processor.tokenizer), - train_config.get("action_tokenizer_type"), - ) - # 6) model + # 6) model — the OSS package evaluates on the vanilla wall-x + # model class (the internal opt inference engine is not part of + # this release). The fused CUDA kernels are compiled at install + # time and used automatically, so flow-inference accuracy + # matches the internal build. ModelClass = ta.inference_model_class() self._model_class = ModelClass model = ModelClass(self._model_config, self._processor, self._tokenizer_mixin) @@ -179,7 +153,7 @@ def __init__(self, cfg: EvalConfig) -> None: model.to_bfloat16_for_selected_params() # 7) checkpoint weights and model finalization. - state_dict = load_state_dict(self._checkpoint_path, ModelClass) + state_dict = load_state_dict(cfg.model.checkpoint_path, ModelClass) embed_key = "model.embed_tokens.weight" if embed_key in state_dict: ckpt_vocab = state_dict[embed_key].shape[0] @@ -210,7 +184,7 @@ def __init__(self, cfg: EvalConfig) -> None: # 8) cached runtime fields self._cam_names = list(cfg.model.cam_names) self._action_horizon = int(cfg.model.action_horizon) - self._action_dim = sum(train_config["dof_config"].values()) + self._action_dim = sum(typed_cfg.task.dof_config.values()) self._num_inference_timesteps = 10 self._robot_id = "10000" @@ -224,28 +198,40 @@ def chunk_horizon(self) -> int: def action_mode(self) -> str: return self._action_mode - def predict_batch(self, payloads: list[dict]) -> list[np.ndarray]: + @property + def action_dim(self) -> int: + return self._action_dim + + @property + def task_config(self): + return self._typed_cfg.task + + def infer_batch(self, payloads: list[dict]) -> list[np.ndarray]: if self._action_mode == "flow": return self._flow_batch(payloads) raise NotImplementedError( - f"{type(self).__name__}.predict_batch: action_mode={self._action_mode!r} " + f"{type(self).__name__}.infer_batch: action_mode={self._action_mode!r} " "is not dispatched" ) + @property + def metadata(self): + return { + "model_type": "qwen2_5", + "action_mode": self._action_mode, + "action_horizon": self._action_horizon, + } + # ---- batched flow inference ---- def _flow_batch(self, payloads: list[dict]) -> list[np.ndarray]: - # 1) Encode each observation into proprioception, masks, and views. + # 1) Collect model-ready observations. observations: list[dict] = [] instructions: list[str] = [] noises: list[np.ndarray | None] = [] any_noise = False for p in payloads: - observations.append( - encode_proprio( - p["observation"], self._train_config, self._action_horizon - ) - ) + observations.append(p["observation"]) instructions.append(p["instruction"]) n = p.get("noise") if n is not None: @@ -297,23 +283,18 @@ def _flow_batch(self, payloads: list[dict]) -> list[np.ndarray]: **batch_inputs, ) - # 6) Decode one action chunk per payload. + # 6) Return one raw model action chunk per request. predict_action = model_output["predict_action"] # (B, H, D_action) if isinstance(predict_action, torch.Tensor): predict_action = predict_action.detach().cpu().numpy() - return [ - decode_chunk(predict_action[i : i + 1], self._train_config) - for i in range(len(payloads)) - ] + return [predict_action[i : i + 1] for i in range(len(payloads))] # ---- prompt template ---- def _get_flow_prompt(self, instruction: str) -> tuple[str, str]: """Build the flow-action prompt.""" - if self._train_config["data"].get("use_embodied_system_prompt_ratio", 0) > 0: - robot_id = ( - self._robot_id if self._norm_key in ("x2_normal", "ex_normal") else 0 - ) + if getattr(self._typed_cfg.data, "use_embodied_system_prompt_ratio", 0) > 0: + robot_id = self._robot_id if self._norm_key in ("x2_normal", "ex_normal") else 0 cam_name_mapping = {cn: cn for cn in self._cam_names} prologue = get_prologue_with_embodied_information( dataset_name=self._norm_key, @@ -323,17 +304,22 @@ def _get_flow_prompt(self, instruction: str) -> tuple[str, str]: config=self._data_config, ) else: - prologue = f"{_ROLE_START}system\nYou are a helpful assistant.{_ROLE_END}\n" + prologue = ( + f"{_ROLE_START}system\nYou are a helpful assistant.{_ROLE_END}\n" + ) user_request = f"{_ROLE_START}user\nObservation:" for cn in self._cam_names: user_request += ( - f" {_camera_label(cn)}: " f"{_VISION_START}{_IMAGE_PAD}{_VISION_END}" + f" {_camera_label(cn)}: " + f"{_VISION_START}{_IMAGE_PAD}{_VISION_END}" ) user_request += "\nInstruction:" text_prompt = ( f"\nPredict the next action in robot action.\nProprioception: {_PROPRI}\n" ) - user_message = f"{user_request} {instruction}{text_prompt}{_ROLE_END}\n" + user_message = ( + f"{user_request} {instruction}{text_prompt}{_ROLE_END}\n" + ) assistant_message = f"{_ROLE_START}assistant\n" flow_action = _ACTION * self._action_horizon @@ -353,7 +339,7 @@ def _construct_model_input( batch_size = len(observations) dataset_names = [self._norm_key] * batch_size - # Proprioception and masks are prepared as ndarrays in encode_proprio. + # Proprioception and masks are supplied as model-ready ndarrays. agent_pos = torch.cat( [torch.from_numpy(o["proprioception"]) for o in observations], dim=0 ) @@ -393,14 +379,14 @@ def _construct_model_input( pad_prefix_to_same_length=False, norm_state=( agent_pos - if resolve_use_state_string_representation(self._train_config) + if self._typed_cfg.task.use_state_string_representation else None ), agent_pos_mask=agent_pos_mask, state_augmentation_prob=0.0, state_drop_prob=0.0, state_augmentation_ratio=0.0, - state_bins=resolve_state_bins(self._train_config), + state_bins=getattr(self._typed_cfg.data, "state_bins", 256), inference_mode=True, ) diff --git a/wall_x/_vendor/harrix/serving/__init__.py b/wall_x/_vendor/harrix/serving/__init__.py index 87f42b29..b9e1afbd 100644 --- a/wall_x/_vendor/harrix/serving/__init__.py +++ b/wall_x/_vendor/harrix/serving/__init__.py @@ -1,3 +1 @@ -from .websocket_policy_server import WebsocketPolicyServer, BasePolicy - -__all__ = ["WebsocketPolicyServer", "BasePolicy"] +"""Online model serving frontend.""" diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/__init__.py b/wall_x/_vendor/harrix/serving/_wallx_infer/__init__.py deleted file mode 100644 index cb5d7d60..00000000 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Wall-X inference helpers vendored for harrix serving.""" diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/base_dataclass.py b/wall_x/_vendor/harrix/serving/_wallx_infer/base_dataclass.py deleted file mode 100644 index dccdacc6..00000000 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/base_dataclass.py +++ /dev/null @@ -1,476 +0,0 @@ -from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig -from typing import Optional, List -from dataclasses import dataclass, field -import numpy as np -import torch -import wall_x._vendor.x2robot_utils.geometry as data_utils -from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger - -dof_dims = { - "left_ee_cartesian_pos": 3, - "left_ee_cartesian_pos_relative": 3, - "left_ee_rotation": 3, - "left_ee_rotation_relative": 3, - "left_ee_rotation_6D": 6, - "left_ee_rotation_6D_relative": 6, - "left_arm_joint_pos": 7, # 6 joint + 1 gripper - "left_gripper": 1, - "left_gripper_cur": 1, - "left_arm_joint_cur": 7, # 1 -> 7 - "right_ee_cartesian_pos": 3, - "right_ee_cartesian_pos_relative": 3, - "right_ee_rotation": 3, - "right_ee_rotation_relative": 3, - "right_ee_rotation_6D": 6, - "right_ee_rotation_6D_relative": 6, - "right_arm_joint_pos": 7, # 6 joint + 1 gripper - "right_gripper": 1, - "right_gripper_cur": 1, - "right_arm_joint_cur": 7, # 1 -> 7 - "head_actions": 2, - "height": 1, - "car_pose": 3, - "velocity_decomposed": 3, - "velocity_decomposed_odom": 3, - "head_rotation": 2, # match ex001 - "left_joint": 6, - "left_joint_gripper": 1, - "right_joint": 6, - "right_joint_gripper": 1, - "left_rotation_quat": 4, - "right_rotation_quat": 4, - "left_quaternion": 4, - "right_quaternion": 4, - "left_wrench_ext_local_force": 3, - "left_wrench_ext_local_torque": 3, - "right_wrench_ext_local_force": 3, - "right_wrench_ext_local_torque": 3, - "left_wrench_ext_local_force_from_joint": 3, - "left_wrench_ext_local_torque_from_joint": 3, - "right_wrench_ext_local_force_from_joint": 3, - "right_wrench_ext_local_torque_from_joint": 3, - "left_wrench_ext_world_force": 3, - "left_wrench_ext_world_torque": 3, - "left_wrench_ext_world_force_from_joint": 3, - "left_wrench_ext_world_torque_from_joint": 3, - "right_wrench_ext_world_force": 3, - "right_wrench_ext_world_torque": 3, - "right_wrench_ext_world_force_from_joint": 3, - "right_wrench_ext_world_torque_from_joint": 3, - "left_arm_joint_dev": 7, - "right_arm_joint_dev": 7, -} - - -class ComputedDict(dict): - """Dict that registers compute rules and auto-computes None values on get""" - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._compute_rules = {} # key -> compute_function - - def register_compute_rule(self, key, compute_func): - """ - Register a compute rule - - Args: - key: Key to compute - compute_func: Callable taking self, returns computed value - """ - self._compute_rules[key] = compute_func - - def get(self, key, default=None): - """Override get to support auto-compute""" - value = super().get(key, default) - - # If value is None and a rule exists, try to compute - if value is None and key in self._compute_rules: - try: - computed_value = self._compute_rules[key](self) - if computed_value is not None: - # Cache computed value - self[key] = computed_value - return computed_value - except Exception: - pass # On failure return None or default - - return value if value is not None else default - - def __getitem__(self, key): - """Override [] to support auto-compute""" - value = super().get(key, None) - - # If value is None and a rule exists, try to compute - if value is None and key in self._compute_rules: - try: - computed_value = self._compute_rules[key](self) - if computed_value is not None: - # Cache computed value - self[key] = computed_value - return computed_value - except Exception: - pass # On failure raise KeyError or return None - - if key in self: - return super().__getitem__(key) - raise KeyError(key) - - -@dataclass -class RobotStateActionData: - config: InferConfig = None - data: ComputedDict = field( - default_factory=lambda: ComputedDict( - { - # State (formerly pose) - state_ prefix - "state_left_ee_cartesian_pos": None, # (1, 3) - "state_left_ee_rotation": None, # (1, 3) - "state_left_ee_rotation_6D": None, - "state_left_arm_joint_pos": None, - "state_left_gripper": None, # (1, 1) - "state_left_arm_joint_cur": None, - "state_left_gripper_cur": None, - "state_right_ee_cartesian_pos": None, # (1, 3) - "state_right_ee_rotation": None, - "state_right_ee_rotation_6D": None, # (1, 6) - "state_right_arm_joint_pos": None, - "state_right_gripper": None, - "state_right_gripper_cur": None, - "state_right_arm_joint_cur": None, # (1, 1) - "state_head_actions": None, - "state_head_rotation": None, # match ex001 - "state_height": None, - "state_car_pose": None, - "state_velocity_decomposed": None, - "state_velocity_decomposed_odom": None, - # support joint control - "state_left_joint": None, - "state_left_joint_gripper": None, - "state_right_joint": None, - "state_right_joint_gripper": None, - # support quaternion control - "state_left_quaternion": None, - "state_right_quaternion": None, - "state_left_rotation_quat": None, - "state_right_rotation_quat": None, - # support wrench/force-torque observation - "state_left_wrench_ext_local_force": None, - "state_left_wrench_ext_local_torque": None, - "state_right_wrench_ext_local_force": None, - "state_right_wrench_ext_local_torque": None, - "state_left_wrench_ext_local_force_from_joint": None, - "state_left_wrench_ext_local_torque_from_joint": None, - "state_right_wrench_ext_local_force_from_joint": None, - "state_right_wrench_ext_local_torque_from_joint": None, - "state_left_wrench_ext_world_force": None, - "state_left_wrench_ext_world_torque": None, - "state_left_wrench_ext_world_force_from_joint": None, - "state_left_wrench_ext_world_torque_from_joint": None, - "state_right_wrench_ext_world_force": None, - "state_right_wrench_ext_world_torque": None, - "state_right_wrench_ext_world_force_from_joint": None, - "state_right_wrench_ext_world_torque_from_joint": None, - # support joint deviation - "state_left_arm_joint_dev": None, - "state_right_arm_joint_dev": None, - # Action - action_ prefix - "action_left_ee_cartesian_pos": None, - "action_left_ee_cartesian_pos_relative": None, - "action_left_ee_rotation": None, - "action_left_ee_rotation_relative": None, - "action_left_ee_rotation_6D": None, - "action_left_ee_rotation_6D_relative": None, - "action_left_gripper": None, - "action_left_arm_joint_pos": None, - "action_right_ee_cartesian_pos": None, - "action_right_ee_cartesian_pos_relative": None, - "action_right_ee_rotation": None, - "action_right_ee_rotation_relative": None, - "action_right_ee_rotation_6D": None, - "action_right_ee_rotation_6D_relative": None, - "action_right_gripper": None, - "action_right_arm_joint_pos": None, - "action_head_actions": None, - "action_head_rotation": None, # match ex001 - "action_height": None, - "action_car_pose": None, - "action_velocity_decomposed": None, - "action_velocity_decomposed_odom": None, - # support joint control - "action_left_joint": None, - "action_left_joint_gripper": None, - "action_right_joint": None, - "action_right_joint_gripper": None, - # support quaternion control - "action_left_quaternion": None, - "action_right_quaternion": None, - "action_left_rotation_quat": None, - "action_right_rotation_quat": None, - } - ) - ) - dof_mask: np.ndarray = None - logger = InferLogger.get_robot_logger("RobotStateActionData") - - def __post_init__(self): - """Register compute rules""" - # State rules: euler angles -> 6D rotation - self.data.register_compute_rule( - "state_left_ee_rotation_6D", - lambda d: ( - data_utils.euler_to_matrix_zyx_6d_nb(d["state_left_ee_rotation"]) - if d.get("state_left_ee_rotation") is not None - else None - ), - ) - self.data.register_compute_rule( - "state_right_ee_rotation_6D", - lambda d: ( - data_utils.euler_to_matrix_zyx_6d_nb(d["state_right_ee_rotation"]) - if d.get("state_right_ee_rotation") is not None - else None - ), - ) - - # Action rules: absolute position from relative + state - self.data.register_compute_rule( - "action_left_ee_cartesian_pos", - lambda d: ( - d.get("state_left_ee_cartesian_pos") - + d.get("action_left_ee_cartesian_pos_relative") - if d.get("state_left_ee_cartesian_pos") is not None - and d.get("action_left_ee_cartesian_pos_relative") is not None - else None - ), - ) - self.data.register_compute_rule( - "action_right_ee_cartesian_pos", - lambda d: ( - d.get("state_right_ee_cartesian_pos") - + d.get("action_right_ee_cartesian_pos_relative") - if d.get("state_right_ee_cartesian_pos") is not None - and d.get("action_right_ee_cartesian_pos_relative") is not None - else None - ), - ) - - # Action rules: compute absolute rpy - self.data.register_compute_rule( # delta rpy -> abs rpy - "action_left_ee_rotation", - lambda d: ( - data_utils.compose_state_and_delta_to_abs_rpy( - d["action_left_ee_rotation_relative"], - d["state_left_ee_rotation"][0], - ) - if d.get("action_left_ee_rotation_relative") is not None - and d.get("state_left_ee_rotation") is not None - else None - ), - ) - self.data.register_compute_rule( # abs 6D -> abs rpy - "action_left_ee_rotation", - lambda d: ( - data_utils.so3_to_euler_zyx_batch_nb(d["action_left_ee_rotation_6D"]) - if d.get("action_left_ee_rotation_6D") is not None - else None - ), - ) - self.data.register_compute_rule( # delta 6D -> abs 6D - "action_left_ee_rotation_6D", - lambda d: ( - data_utils.compose_state_and_delta_to_abs_6d( - d["action_left_ee_rotation_6D_relative"], - d["state_left_ee_rotation_6D"][0], - ) - if d.get("action_left_ee_rotation_6D_relative") is not None - and d.get("state_left_ee_rotation_6D") is not None - else None - ), - ) - - self.data.register_compute_rule( # delta rpy -> abs rpy - "action_right_ee_rotation", - lambda d: ( - data_utils.compose_state_and_delta_to_abs_rpy( - d["action_right_ee_rotation_relative"], - d["state_right_ee_rotation"][0], - ) - if d.get("action_right_ee_rotation_relative") is not None - and d.get("state_right_ee_rotation") is not None - else None - ), - ) - self.data.register_compute_rule( # abs 6D -> abs rpy - "action_right_ee_rotation", - lambda d: ( - data_utils.so3_to_euler_zyx_batch_nb(d["action_right_ee_rotation_6D"]) - if d.get("action_right_ee_rotation_6D") is not None - else None - ), - ) - self.data.register_compute_rule( # delta 6D -> abs 6D - "action_right_ee_rotation_6D", - lambda d: ( - data_utils.compose_state_and_delta_to_abs_6d( - d["action_right_ee_rotation_6D_relative"], - d["state_right_ee_rotation_6D"][0], - ) - if d.get("action_right_ee_rotation_6D_relative") is not None - and d.get("state_right_ee_rotation_6D") is not None - else None - ), - ) - - def get_agent_pos(self, obs_action_keys=None): - if obs_action_keys is None: - obs_action_keys = self.config.train_config["agent_pos_config"].keys() - - agent_pose_data = [] - for key in obs_action_keys: - # action_padding is a virtual key used to pad agent_pos width. - if key == "action_padding": - dim = self.config.train_config["agent_pos_config"][key] - agent_pose_data.append(np.zeros((1, dim))) - continue - - # Strip follow_ or master_ prefix - if key.startswith("follow_"): - key = key.replace("follow_", "") - elif key.startswith("master_"): - key = key.replace("master_", "") - - # Add state_ prefix to access state data - state_key = f"state_{key}" - - if state_key in self.data: - # get() auto-computes when value is None - value = self.data.get(state_key) - if value is None: - # If still None after compute, use zeros - agent_pose_data.append(np.zeros((1, dof_dims[key]))) - else: - agent_pose_data.append(value) - else: - raise ValueError(f"Key {state_key} not found in data") - - agent_pose_data = np.concatenate(agent_pose_data, axis=1)[None] # (1, 1, D) - - return agent_pose_data - - def get_agent_pos_mask(self, obs_action_keys=None): - if obs_action_keys is None: - obs_action_keys = self.config.train_config["agent_pos_config"].keys() - - agent_pos_mask_data = [] - for key in obs_action_keys: - # action_padding carries no information and should stay masked out. - if key == "action_padding": - dim = self.config.train_config["agent_pos_config"][key] - agent_pos_mask_data.append(np.zeros((1, dim))) - continue - - # Strip follow_ or master_ prefix - if key.startswith("follow_"): - key = key.replace("follow_", "") - elif key.startswith("master_"): - key = key.replace("master_", "") - - # Add state_ prefix to access state data - state_key = f"state_{key}" - - if state_key in self.data: - # get() auto-computes when value is None - value = self.data.get(state_key) - if value is None: - agent_pos_mask_data.append(np.zeros((1, dof_dims[key]))) - else: - agent_pos_mask_data.append(np.ones((1, value.shape[1]))) - else: - raise ValueError(f"Key {state_key} not found in data") - - return np.concatenate(agent_pos_mask_data, axis=1)[None] # (1, 1, D) - - def save_state_data_with_key(self, value, key, gt_dim=None): - # Strip follow_ or master_ prefix - key = key.replace("follow_", "") - key = key.replace("master_", "") - - # if torch, convert to numpy - if isinstance(value, torch.Tensor): - value = value.detach().cpu().numpy() - - if f"state_{key}" not in self.data: # TODO: joint angle control - self.logger.warning(f"{key} is not a valid state key; not recorded") - return - gt_dim = dof_dims[key] if gt_dim is None else gt_dim - - # Shape check: expected (1, D) - if value.shape == (1, gt_dim): - self.data[f"state_{key}"] = value - elif value.shape == (1, 1, gt_dim): - self.data[f"state_{key}"] = value[0] - elif value.shape == (gt_dim,): - self.data[f"state_{key}"] = value[None] - else: - raise ValueError(f"Value shape {value.shape} is not legal") - - def save_action_data_with_key(self, value, key): - key = key.replace("follow_", "") - key = key.replace("master_", "") - - if isinstance(value, torch.Tensor): - value = value.detach().cpu().numpy() - - if value.shape == (dof_dims[key],): - self.data[f"action_{key}"] = value[None] - else: - self.data[f"action_{key}"] = value - - def save_action_data( - self, predict_action, predict_action_keys: Optional[List[str]] = None - ): - if predict_action_keys is None: - predict_action_keys = getattr( - self.config.data_config, "predict_action_keys", None - ) - if predict_action_keys is None: - try: - predict_action_keys = self.config.data_config["predict_action_keys"] - except (KeyError, AttributeError): - predict_action_keys = list( - self.config.train_config["dof_config"].keys() - ) - - if isinstance(predict_action, torch.Tensor): - predict_action = predict_action.detach().cpu().numpy() - - if predict_action.ndim == 3: - predict_action = predict_action[0] - - dof_start = 0 - for action_key in predict_action_keys: - # action_padding is a virtual key that only advances dof_start; no data is written. - if action_key == "action_padding": - dof_dim = self.config.train_config["dof_config"]["action_padding"] - dof_start += dof_dim - continue - - action_key = action_key.replace("follow_", "") - action_key = action_key.replace("master_", "") - dof_dim = dof_dims[action_key] - action_key = f"action_{action_key}" - self.data[action_key] = predict_action[:, dof_start : dof_start + dof_dim] - dof_start += dof_dim - - # Convenience properties for compatibility - @property - def agent_pos(self): - return self.get_agent_pos() - - @property - def agent_pos_mask(self): - return self.get_agent_pos_mask() - - @property - def action(self): - pass # TODO: support action access diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py b/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py deleted file mode 100644 index 5f7f70b7..00000000 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py +++ /dev/null @@ -1,434 +0,0 @@ -import logging -import os -from dataclasses import dataclass, field -from typing import Any - -import yaml -from qwen_vl_utils.vision_process import IMAGE_FACTOR, MAX_PIXELS, MIN_PIXELS - -logger = logging.getLogger(__name__) - - -def _env_bool(name: str, default: bool) -> bool: - v = os.environ.get(name) - if v is None: - return default - return v.strip().lower() in ("1", "true", "yes", "on") - - -@dataclass -class InferenceDataConfig: - """Minimal data config needed by online inference. - - Serving does not build datasets, so it should not require closed-source - data backends just to read image resize settings from a checkpoint yaml. - """ - - resolution: dict[str, int] = field(default_factory=dict) - model_type: str = "qwen2_5" - max_pixels: int = MAX_PIXELS - min_pixels: int = MIN_PIXELS - image_factor: int = IMAGE_FACTOR - use_relative_action: bool = False - predict_action_keys: list[str] = field(default_factory=list) - - def __getitem__(self, key: str) -> Any: - if not hasattr(self, key): - raise KeyError(key) - return getattr(self, key) - - -class InferConfig: - def __init__( - self, - checkpoint_path: str | None = None, - train_config_path: str | None = None, - robot_host: str = "0.0.0.0", - robot_port: int = 41776, - robot_type: str = "desktop", # ["desktop", "turtle", "ex001"] - robot_action_start_ratio: float = 0, # Start ratio for trimming executed actions - robot_action_end_ratio: float = 0.8, # End ratio for trimming executed actions - robot_action_interpolate_multiplier: int = 10, # Action interpolation multiplier - robot_use_joint_angle_control: bool = False, # Joint control (model must predict joints) - turtle_as_desktop: bool = False, # Use turtle as desktop: fixed base/head motion, head camera, base height - action_horizon: int = 32, # Set to the model's action horizon - action_dim: int | None = None, - ar_action_dim: int | None = None, - model_device: str = "cuda:0", - num_inference_timesteps: int = 10, - num_inference_steps: int | None = None, - cfg_scale: float | None = None, - seed: int | None = None, - norm_key: str = "x2_normal", # ["x2_normal", "ex_normal"] - cam_names: list[str] | None = None, - camera_front_key: str = "camera_front", - camera_left_key: str = "camera_left", - camera_right_key: str = "camera_right", - default_instruction: str | None = None, - prompt_template: str | None = None, - qwen25_prompt_template: str | None = None, - prompt_priority_order: str | None = None, - save_video_dir: str = "./videos", - robot_id: str = "10000", - model_type: str = "wallx", # ["wallx", "vga"] - smooth_action: bool = False, - smooth_gripper: bool = True, - ): - # Private attributes for paths - assert checkpoint_path is not None - from wall_x._vendor.harrix.utils.ckpt_load import resolve_checkpoint_dir - - if not os.path.isdir(checkpoint_path) and not os.path.isfile(checkpoint_path): - raise FileNotFoundError( - f"Checkpoint path not found: {checkpoint_path!r}. " - "Serving requires a directory containing model.safetensors and " - "normalizer_action.pth / normalizer_propri.pth (or norm_stats.json)." - ) - self._checkpoint_path = resolve_checkpoint_dir(checkpoint_path) - action_pth = os.path.join(self._checkpoint_path, "normalizer_action.pth") - propri_pth = os.path.join(self._checkpoint_path, "normalizer_propri.pth") - if os.path.exists(action_pth): - self.normalizer_action_path = action_pth - if os.path.exists(propri_pth): - self.normalizer_propri_path = propri_pth - - # Other config attributes - self.robot_host = robot_host - self.robot_port = robot_port - self.robot_type = robot_type # ["desktop", "turtle", "ex001"] - self.robot_action_start_ratio = robot_action_start_ratio - self.robot_action_end_ratio = robot_action_end_ratio - self.robot_action_interpolate_multiplier = robot_action_interpolate_multiplier - self.robot_use_joint_angle_control = ( - robot_use_joint_angle_control # Joint-angle control - ) - self.turtle_as_desktop = turtle_as_desktop - self.robot_id = robot_id - self._action_horizon = ( - action_horizon # Default from train config flow action horizon - ) - self._action_dim = action_dim # Default from train config dof config - self.ar_action_dim = ar_action_dim # Default from train config ar dof config - self.model_device = model_device - self.num_inference_timesteps = ( - num_inference_timesteps # flow matching related config - ) - self.num_inference_steps = num_inference_steps - self.cfg_scale = cfg_scale - self.seed = seed - self.save_video_dir = save_video_dir - self.model_type = model_type # ["wallx", "vga"] - # Initialize config objects - self.train_config: dict = {} - self.model_config = None - self.data_config = None - self.norm_key = norm_key - self.cam_names = cam_names or [ - "face_view", - "left_wrist_view", - "right_wrist_view", - ] - self.camera_front_key = camera_front_key - self.camera_left_key = camera_left_key - self.camera_right_key = camera_right_key - self.default_instruction = default_instruction - self.prompt_template = prompt_template - self.qwen25_prompt_template = qwen25_prompt_template - self.prompt_priority_order = prompt_priority_order - self.smooth_action = _env_bool("WALLX_SMOOTH_ACTION", smooth_action) - self.smooth_gripper = _env_bool("WALLX_SMOOTH_GRIPPER", smooth_gripper) - # Load all configs - self._load_all_configs(train_config_path) - self._apply_cam_names_from_train_config(cam_names) - - def _apply_cam_names_from_train_config( - self, cam_names: list[str] | None - ) -> None: - """Use train YAML camera mapping when CLI did not override cam_names.""" - if cam_names is not None: - return - from wall_x._vendor.harrix.utils.train_config import ( - resolve_cam_names_from_train_config, - ) - - resolved = resolve_cam_names_from_train_config(self.train_config) - if resolved: - self.cam_names = resolved - logger.info( - "[InferConfig] cam_names from train key_mappings: %s", - self.cam_names, - ) - - @property - def checkpoint_path(self) -> str | None: - return self._checkpoint_path - - @checkpoint_path.setter - def checkpoint_path(self, value: str | None): - """Reload all configs when checkpoint_path is updated""" - if self._checkpoint_path != value: - self._checkpoint_path = value - self._load_all_configs() - - @property - def action_horizon(self) -> int: - return self._action_horizon - - @action_horizon.setter - def action_horizon(self, value: int): - self._action_horizon = value - - @property - def action_dim(self) -> int | None: - return self._action_dim - - @action_dim.setter - def action_dim(self, value: int | None): - self._action_dim = value - - @property - def ar_action_dim(self) -> int | None: - return self._ar_action_dim - - @ar_action_dim.setter - def ar_action_dim(self, value: int | None): - self._ar_action_dim = value - - def _load_all_configs(self, train_config_path=None): - """Unified entry to load all configs""" - self._load_train_config(train_config_path) - if self.model_type != "vga": - self._load_model_config_for_wallx() - self._load_data_config() - - # Update action_horizon and action_dim if needed - if self._action_horizon is None: - self._action_horizon = self.train_config.get("data", {}).get( - "action_horizon_flow", 32 - ) - assert self._action_horizon is not None and self._action_horizon > 0 - - if self._action_dim is None: - dof_config = ( - self.train_config.get("dof_config") - or self.train_config.get("task", {}).get("dof_config") - or self.train_config.get("data", {}).get("dof_config", {}) - ) - self._action_dim = sum(dof_config.values()) - - if self._ar_action_dim is None: - ar_dof_config = self.train_config.get("ar_dof_config") or self.train_config.get( - "task", {} - ).get("ar_dof_config", {}) - self._ar_action_dim = sum(ar_dof_config.values()) - - def _load_train_config(self, train_config_path): - if train_config_path is None: - for fname in ("config.yml", "config.yaml"): - candidate = os.path.join(self._checkpoint_path, fname) - if os.path.exists(candidate): - train_config_path = candidate - break - else: - raise FileNotFoundError( - f"No config.yml/config.yaml found in {self._checkpoint_path}" - ) - with open(train_config_path, "r") as f: - self.train_config = yaml.load(f, Loader=yaml.FullLoader) - - ckpt_dir = self._checkpoint_path - preprocessor_file = os.path.join(ckpt_dir, "preprocessor_config.json") - if os.path.exists(preprocessor_file): - logger.info( - "[LoadConfig] Found %s, override processor_path.", - preprocessor_file, - ) - self.train_config["processor_path"] = ckpt_dir - - tokenizer_file = os.path.join(ckpt_dir, "tokenizer.json") - tokenizer_config_file = os.path.join(ckpt_dir, "tokenizer_config.json") - if self.train_config.get( - "action_tokenizer_path", None - ) is not None and not os.path.exists( - self.train_config.get("action_tokenizer_path", None) - ): - if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file): - logger.info( - "[LoadConfig] Found tokenizer files in %s, override action_tokenizer_path.", - ckpt_dir, - ) - self.train_config["action_tokenizer_path"] = ckpt_dir - else: - logger.warning("[LoadConfig] Cannot load action tokenizer! ") - - from wall_x._vendor.harrix.utils.train_config import ( - normalize_train_config_for_inference, - strip_action_tokenizer_fields, - ) - - self.train_config = normalize_train_config_for_inference( - self.train_config, train_config_path - ) - # Flow serving does not load AR action tokenizers; crop embed_tokens instead. - strip_action_tokenizer_fields(self.train_config) - self._train_config_path = train_config_path - - def _load_model_config_for_wallx(self): - model_type = self.train_config["model_type"] - if not os.path.isdir(self._checkpoint_path): - return - # For Qwen models - ckpt_config_path = os.path.join(self._checkpoint_path, "config.json") - resolved_cfg_path = None - - if os.path.exists(ckpt_config_path): - # Prefer checkpoint config - resolved_cfg_path = ckpt_config_path - logger.info( - "[LoadModelConfig] Using checkpoint config.json: %s", - ckpt_config_path, - ) - else: - # Fallback to original config path - fallback_cfg = self.train_config.get("qwen_vl_act_config_path", None) - if fallback_cfg is not None: - resolved_cfg_path = fallback_cfg - logger.info( - "[LoadModelConfig] Using fallback act config: %s", - fallback_cfg, - ) - - if resolved_cfg_path is None or (not os.path.exists(resolved_cfg_path)): - raise ValueError( - f"[LoadModelConfig] Cannot load model config! " - f"Checked:\n" - f" - Checkpoint config.json: {ckpt_config_path}\n" - f" - Fallback path: {self.train_config.get('qwen_vl_act_config_path', None)}" - ) - - # Save back to config for consistency - self.train_config["qwen_vl_act_config_path"] = resolved_cfg_path - - from wall_x.trainer.adapters import resolve_adapter - - adapter_cls = resolve_adapter(model_type) - ConfigClass = adapter_cls.config_class() - - logger.info( - "[LoadModelConfig] Loading model config from: %s", resolved_cfg_path - ) - if resolved_cfg_path.endswith(".json"): - self.model_config = ConfigClass.from_json_file(resolved_cfg_path) - else: - self.model_config = ConfigClass.from_pretrained(resolved_cfg_path) - - self.model_config.update_model_config(self.train_config) - - self.model_config._attn_implementation = "sdpa" - vision_attn = os.environ.get("WALLX_VISION_ATTN_IMPLEMENTATION", "flash_attention_2") - self.model_config.vision_config._attn_implementation = vision_attn - logger.info( - "[LoadModelConfig] vision _attn_implementation=%s (override via WALLX_VISION_ATTN_IMPLEMENTATION)", - vision_attn, - ) - - logger.info("[LoadModelConfig] Model config loaded and updated successfully.") - - def _load_data_config(self): - # Prefer typed TrainConfig path (handles new 8-section schema where - # dof_config lives under task:). Fall back to legacy raw-dict path - # for old flat yamls. See trainer/adapters/base_adapter.py. - typed_dcfg = self._try_typed_data_config() - if typed_dcfg is not None: - self.data_config = typed_dcfg - elif self.model_type != "vga": - self.data_config = self._build_inference_data_config() - else: - # TEMPORARY: direct call to the private _set_data_backend. - # VGA inference still uses the legacy x2robot data config object. - # Wall-X online inference uses the lightweight InferenceDataConfig - # above so serving does not require excluded data backends. - from wall_x.data._registry import _set_data_backend - - _set_data_backend(self.train_config.get("dataset_type", "x2robot_v1")) - from wall_x.model.vga.openloop_visualization import get_data_configs - - dataload_config = get_data_configs(self.train_config.get("data", {})) - dataload_config["predict_action_keys"] = list( - dataload_config.get("dof_config", {}).keys() - ) - self.data_config = dataload_config - - self._ensure_predict_action_keys() - - def _ensure_predict_action_keys(self) -> None: - dof_keys = list((self.train_config.get("dof_config") or {}).keys()) - if not dof_keys: - return - if isinstance(self.data_config, InferenceDataConfig): - if not self.data_config.predict_action_keys: - self.data_config.predict_action_keys = dof_keys - elif isinstance(self.data_config, dict): - self.data_config.setdefault("predict_action_keys", dof_keys) - - def _build_inference_data_config(self) -> InferenceDataConfig: - data = self.train_config.get("data", {}) or {} - - def get(key: str, default: Any) -> Any: - return data.get(key, self.train_config.get(key, default)) - - resolution = get("resolution", {}) or {} - dof_config = self.train_config.get("dof_config") or {} - return InferenceDataConfig( - resolution=dict(resolution), - model_type=get( - "model_type", self.train_config.get("model_type", "qwen2_5") - ), - max_pixels=get("max_pixels", MAX_PIXELS), - min_pixels=get("min_pixels", MIN_PIXELS), - image_factor=get("image_factor", IMAGE_FACTOR), - use_relative_action=get("use_relative_action", False), - predict_action_keys=list(dof_config.keys()), - ) - - def _try_typed_data_config(self): - """Attempt to load TrainConfig and build X2RDataConfig via typed path. - - Returns the X2RDataConfig on success, or None if the yaml is not in - TrainConfig schema (legacy flat yaml). - """ - yml_path = getattr(self, "_train_config_path", None) - if yml_path is None: - for fname in ("config.yml", "config.yaml"): - candidate = os.path.join(self._checkpoint_path, fname) - if os.path.exists(candidate): - yml_path = candidate - break - if yml_path is None: - return None - try: - from wall_x.config.loader import load_config - from wall_x.trainer.adapters.base_adapter import load_trainer_data_config - - typed_cfg = load_config(yml_path) - except Exception as e: - logger.warning( - "[InferConfig] Typed config load failed for %s: %s", - yml_path, - e, - ) - return None - try: - return load_trainer_data_config(typed_cfg) - except Exception as e: - logger.warning( - "[InferConfig] Typed config load failed for %s: %s", - yml_path, - e, - ) - return None - - -if __name__ == "__main__": - config = InferConfig() - logger.info("%s", config.train_config) diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/logger.py b/wall_x/_vendor/harrix/serving/_wallx_infer/logger.py deleted file mode 100644 index ebc60797..00000000 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/logger.py +++ /dev/null @@ -1,277 +0,0 @@ -"""Layered logging helpers for inference components.""" - -import logging -import sys -from typing import Optional -from pathlib import Path -from datetime import datetime - -try: - import colorlog - - HAS_COLORLOG = True -except ImportError: - HAS_COLORLOG = False - logging.getLogger(__name__).warning( - "colorlog not installed. Install with: pip install colorlog" - ) - - -class InferLogger: - """ - Layered inference logging system - """ - - _loggers = {} - _initialized = False - - # Layer identifiers - LEVEL_ENV = "ENV" - LEVEL_ROBOT = "ROBOT" - LEVEL_CONTROLLER = "CONTROLLER" - LEVEL_MODEL = "MODEL" - LEVEL_UTILS = "UTILS" - - # Layer colors (terminal output) - LEVEL_COLORS = { - LEVEL_ENV: "cyan", - LEVEL_ROBOT: "green", - LEVEL_CONTROLLER: "yellow", - LEVEL_MODEL: "purple", # colorlog uses 'purple' not 'magenta' - LEVEL_UTILS: "blue", - } - - @classmethod - def setup( - cls, - log_level: str = "INFO", - log_dir: Optional[str] = None, - console_output: bool = True, - file_output: bool = True, - colorful: bool = True, - ): - """ - Initialize the logging system - - Args: - log_level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL) - log_dir: Log file directory - console_output: Whether to log to console - file_output: Whether to log to file - colorful: Colored console output (requires colorlog) - """ - if cls._initialized: - return - - cls.log_level = getattr(logging, log_level.upper()) - cls.console_output = console_output - cls.file_output = file_output - cls.colorful = colorful and HAS_COLORLOG - - # Create log directory - if file_output and log_dir: - cls.log_dir = Path(log_dir) - cls.log_dir.mkdir(parents=True, exist_ok=True) - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - cls.log_file = cls.log_dir / f"infer_{timestamp}.log" - else: - cls.log_file = None - - cls._initialized = True - - @classmethod - def get_logger(cls, name: str, level: str = None) -> logging.Logger: - """ - Get a logger for the given layer - - Args: - name: Logger name (usually module or class name) - level: Layer tag (ENV, ROBOT, CONTROLLER, MODEL, UTILS) - - Returns: - Configured logger instance - """ - if not cls._initialized: - cls.setup() - - # Auto-detect layer - if level is None: - level = cls._detect_level(name) - - logger_key = f"{level}.{name}" - - if logger_key in cls._loggers: - return cls._loggers[logger_key] - - # Create new logger - logger = logging.getLogger(logger_key) - logger.setLevel(cls.log_level) - logger.propagate = False - - # Clear existing handlers - logger.handlers.clear() - - # Console output - if cls.console_output: - console_handler = logging.StreamHandler(sys.stdout) - console_handler.setLevel(cls.log_level) - - if cls.colorful: - # Colored formatter - color = cls.LEVEL_COLORS.get(level, "white") - console_format = ( - f"%(log_color)s[%(asctime)s]%(reset)s " - f"%(bold_{color})s[{level:^10}]%(reset)s " - f"%(bold_white)s[%(name)s]%(reset)s " - f"%(log_color)s%(levelname)-8s%(reset)s " - f"%(message)s" - ) - - console_formatter = colorlog.ColoredFormatter( - console_format, - datefmt="%H:%M:%S", - log_colors={ - "DEBUG": "cyan", - "INFO": "green", - "WARNING": "yellow", - "ERROR": "red", - "CRITICAL": "bold_red", - }, - ) - else: - # Plain formatter - console_format = ( - f"[%(asctime)s] [{level:^10}] [%(name)s] " - f"%(levelname)-8s %(message)s" - ) - console_formatter = logging.Formatter( - console_format, datefmt="%H:%M:%S" - ) - - console_handler.setFormatter(console_formatter) - logger.addHandler(console_handler) - - # File output - if cls.file_output and cls.log_file: - file_handler = logging.FileHandler(cls.log_file, encoding="utf-8") - file_handler.setLevel(cls.log_level) - - file_format = ( - f"[%(asctime)s] [{level:^10}] [%(name)s] " - f"%(levelname)-8s %(message)s" - ) - file_formatter = logging.Formatter(file_format, datefmt="%Y-%m-%d %H:%M:%S") - file_handler.setFormatter(file_formatter) - logger.addHandler(file_handler) - - cls._loggers[logger_key] = logger - return logger - - @classmethod - def _detect_level(cls, name: str) -> str: - """Auto-detect layer from name""" - name_lower = name.lower() - - if "env" in name_lower: - return cls.LEVEL_ENV - elif "robot" in name_lower and "controller" not in name_lower: - return cls.LEVEL_ROBOT - elif ( - "controller" in name_lower - or "communication" in name_lower - or "socket" in name_lower - ): - return cls.LEVEL_CONTROLLER - elif "model" in name_lower or "wrapper" in name_lower: - return cls.LEVEL_MODEL - else: - return cls.LEVEL_UTILS - - @classmethod - def get_env_logger(cls, name: str = "Environment") -> logging.Logger: - """Get ENV layer logger""" - return cls.get_logger(name, cls.LEVEL_ENV) - - @classmethod - def get_robot_logger(cls, name: str = "Robot") -> logging.Logger: - """Get ROBOT layer logger""" - return cls.get_logger(name, cls.LEVEL_ROBOT) - - @classmethod - def get_controller_logger(cls, name: str = "Controller") -> logging.Logger: - """Get CONTROLLER layer logger""" - return cls.get_logger(name, cls.LEVEL_CONTROLLER) - - @classmethod - def get_model_logger(cls, name: str = "Model") -> logging.Logger: - """Get MODEL layer logger""" - return cls.get_logger(name, cls.LEVEL_MODEL) - - @classmethod - def get_utils_logger(cls, name: str = "Utils") -> logging.Logger: - """Get UTILS layer logger""" - return cls.get_logger(name, cls.LEVEL_UTILS) - - @classmethod - def set_level(cls, level: str): - """Change log level for all loggers""" - new_level = getattr(logging, level.upper()) - cls.log_level = new_level - for logger in cls._loggers.values(): - logger.setLevel(new_level) - for handler in logger.handlers: - handler.setLevel(new_level) - - @classmethod - def close_all(cls): - """Close file handles for all loggers""" - for logger in cls._loggers.values(): - for handler in logger.handlers[:]: - handler.close() - logger.removeHandler(handler) - cls._loggers.clear() - cls._initialized = False - - -# Convenience functions -def get_logger(name: str, level: str = None) -> logging.Logger: - """ - Convenience wrapper to get a logger - - Args: - name: Logger name (usually __name__) - level: Layer tag (optional; auto-detected if omitted) - - Returns: - Configured logger instance - - Example: - from wall_x._vendor.harrix.serving._wallx_infer.logger import get_logger - logger = get_logger(__name__) # auto-detect layer - logger = get_logger(__name__, "ROBOT") # explicit layer - """ - return InferLogger.get_logger(name, level) - - -def setup_logger( - log_level: str = "INFO", - log_dir: Optional[str] = None, - console_output: bool = True, - file_output: bool = True, - colorful: bool = True, -): - """ - Convenience wrapper to configure logging - - Args: - log_level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL) - log_dir: Log file directory - console_output: Whether to log to console - file_output: Whether to log to file - colorful: Colored console output - - Example: - from wall_x._vendor.harrix.serving._wallx_infer.logger import setup_logger - setup_logger(log_level="DEBUG", log_dir="./logs") - """ - InferLogger.setup(log_level, log_dir, console_output, file_output, colorful) diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py b/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py deleted file mode 100644 index 187ea8a2..00000000 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py +++ /dev/null @@ -1,760 +0,0 @@ -import os -import torch -import copy -from safetensors.torch import load_file -import numpy as np -from PIL import Image -from qwen_vl_utils.vision_process import smart_resize -from transformers import BatchFeature - -from wall_x.trainer.trainer_utils import load_wallx_processors - -from wall_x._vendor.x2robot_utils.text_templates import ( - preprocesser_call, - get_prologue_with_embodied_information, -) -from wall_x._vendor.x2robot_utils.grounding import ( - reverse_grounding_points, - extract_grounding_points, -) - -from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig -from wall_x._vendor.harrix.utils.ckpt_load import reshape_compatible_state_dict -from wall_x._vendor.harrix.utils.train_config import ( - resolve_camera_label, - resolve_max_length, - resolve_state_bins, - resolve_use_state_string_representation, -) -from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger -from wall_x.utils.timers import timer, ScopeTimer - - -ENABLE_FAST_PREPROCESS = os.getenv("ENABLE_FAST_PREPROCESS", "False").lower() == "true" - - -def move_to_cuda(obj, device="cuda"): - if isinstance(obj, torch.Tensor): - return obj.to(device) - elif isinstance(obj, (dict, BatchFeature)): - return {k: move_to_cuda(v, device) for k, v in obj.items()} - elif isinstance(obj, list): - return [move_to_cuda(v, device) for v in obj] - elif isinstance(obj, tuple): - return tuple(move_to_cuda(v, device) for v in obj) - else: - return obj - - -class WallxModelWrapper: - def __init__(self, config: InferConfig): - self.config = config - self.logger = InferLogger.get_model_logger("WallxModelWrapper") - self.norm_key = self.config.norm_key - self._register_normalizers() - self.logger.info(f"normalizers {self.norm_key} registered") - self._load_processor() - self._load_model() - self.load_ckpt() - self.logger.info(f"model {self.config.checkpoint_path} loaded") - - self.norm_key = self.config.norm_key - self._register_normalizers() - self.logger.info(f"normalizers {self.norm_key} registered") - - # Initialize robot_type_id (for v3.1 delta tokenizer) - self.robot_type_id = None - if self.tokenizer_mixin is not None: - robot_type = getattr(self.config, "robot_type", None) - if robot_type: - self.tokenizer_mixin.init_inference(robot_type=robot_type) - if hasattr(self.tokenizer_mixin, "robot_type_id"): - self.robot_type_id = self.tokenizer_mixin.robot_type_id - if self.robot_type_id is not None: - self.logger.info( - f"Initialized robot_type_id: {self.robot_type_id} from robot_type: {robot_type}" - ) - - self.role_start_symbol = "<|im_start|>" - self.role_end_symbol = "<|im_end|>" - self.vision_start_symbol = "<|vision_start|>" - self.vision_end_symbol = "<|vision_end|>" - self.image_pad_symbol = "<|image_pad|>" - self.propri_symbol = "<|propri|>" - self.action_symbol = "<|action|>" - - self.cam_names = self.config.cam_names - - def _camera_name_mapping(self): - return self.config.train_config.get("data", {}).get("camera_name_mapping") - - def _load_processor(self): - # Load tokenizer on model_device for inference - device = getattr(self.config, "model_device", "cuda") - processors_dict = load_wallx_processors(self.config.train_config, device=device) - self.processor = processors_dict["processor"] - self.action_mapper = processors_dict["action_mapper"] - self.tokenizer_mixin = processors_dict.get("tokenizer_mixin") - - def _load_model(self): - from wall_x.trainer.adapters import resolve_adapter - - model_type = self.config.train_config["model_type"] - adapter_cls = resolve_adapter(model_type) - ModelClass = adapter_cls.inference_model_class() - self.ModelClass = ModelClass - - self.logger.info(f"initializing model: {model_type} ({ModelClass.__name__})") - - self.model = ModelClass( - self.config.model_config, - self.processor, - self.tokenizer_mixin, - ) - - # log attention implementation - variant-specific layout dispatched - # via the adapter. - adapter_cls.log_attention_implementation(self.logger, self.model) - self.model_type = model_type - - self.logger.info("resizing model token embeddings") - self.model.resize_token_embeddings(len(self.processor.tokenizer)) - self.logger.info("token embedding resize done") - self.logger.info("casting selected params to bfloat16") - self.model.to_bfloat16_for_selected_params() - self.logger.info("bfloat16 cast done") - - def load_ckpt(self, checkpoint_path: str = None): - if checkpoint_path is None: - checkpoint_path = self.config.checkpoint_path - - if os.path.exists(os.path.join(checkpoint_path, "global_step.pth")): - global_step = torch.load(os.path.join(checkpoint_path, "global_step.pth"))[ - "global_step" - ] - self.logger.info(f"checkpoint global_step: {global_step}") - - fsdp_ckpt = os.path.join(checkpoint_path, "pytorch_model_fsdp.bin") - safetensor_ckpt = os.path.join(checkpoint_path, "model.safetensors") - if os.path.exists(fsdp_ckpt): - self.logger.info(f"Loading FSDP checkpoint: {fsdp_ckpt}") - state_dict = torch.load(fsdp_ckpt, map_location="cpu") - # Unwrap outer dict (e.g. {'state_dict': {...}}) - if isinstance(state_dict, dict) and "state_dict" in state_dict: - self.logger.info( - "Using nested state_dict inside pytorch_model_fsdp.bin" - ) - state_dict = state_dict["state_dict"] - elif os.path.exists(safetensor_ckpt): - self.logger.info(f"loading safetensors checkpoint: {safetensor_ckpt}") - state_dict = load_file(safetensor_ckpt, device="cpu") - else: - raise FileNotFoundError( - f"ERROR: No checkpoint found under {checkpoint_path}. " - "Expecting either pytorch_model_fsdp.bin or model.safetensors." - ) - - # Qwen models need fused weight conversion - if not self.ModelClass.is_fused(state_dict): - self.logger.info( - "Converting non-fused weights to fused format...", - ) - state_dict = self.ModelClass.convert_to_fused(state_dict) - else: - self.logger.info( - "The weights is fused, skipping conversion.", - ) - - state_dict = reshape_compatible_state_dict( - state_dict, self.model.state_dict(), log_fn=self.logger.info - ) - msg = self.model.load_state_dict(state_dict, strict=False) - self.model.set_normalizer( - copy.deepcopy(self.normalizer_action), - copy.deepcopy(self.normalizer_propri), - ) - self.logger.info(f"load_state_dict result: {msg}") - self.model.eval() - self.model.to(self.config.model_device) - self.model.to_bfloat16_for_selected_params() - - if hasattr(self.model, "load_optimized_weights"): - self.model.load_optimized_weights(state_dict) - - def _register_normalizers(self): - from wall_x._vendor.harrix.utils.normalizer import build_normalizers - - self.normalizer_action, self.normalizer_propri, resolved = build_normalizers( - self.config.checkpoint_path, - self.config.train_config, - self.config.norm_key, - ) - self.norm_key = resolved - self.config.norm_key = resolved - self._log_norm_debug(self.norm_key) - - def _log_norm_debug(self, norm_key): - if not hasattr(self, "normalizer_action") or not hasattr( - self, "normalizer_propri" - ): - self.logger.warning("[NormDebug] normalizer is not initialized yet") - return - - if ( - norm_key in self.normalizer_action.min - and norm_key in self.normalizer_propri.min - ): - self.logger.debug( - "[NormDebug] norm_key=%s action(min/delta) shape=%s/%s", - norm_key, - tuple(self.normalizer_action.min[norm_key].shape), - tuple(self.normalizer_action.delta[norm_key].shape), - ) - self.logger.debug( - "[NormDebug] action min(all)=%s delta(all)=%s", - self.normalizer_action.min[norm_key].detach().cpu().tolist(), - self.normalizer_action.delta[norm_key].detach().cpu().tolist(), - ) - self.logger.debug( - "[NormDebug] propri(min/delta) shape=%s/%s", - tuple(self.normalizer_propri.min[norm_key].shape), - tuple(self.normalizer_propri.delta[norm_key].shape), - ) - self.logger.debug( - "[NormDebug] propri min(all)=%s delta(all)=%s", - self.normalizer_propri.min[norm_key].detach().cpu().tolist(), - self.normalizer_propri.delta[norm_key].detach().cpu().tolist(), - ) - else: - self.logger.warning( - "[NormDebug] norm_key=%s not found in normalizer", norm_key - ) - - @timer - def construct_model_input( - self, observation, prefix_text, postfix_text, pad_prefix=False - ): - batch_size = len(observation) - dataset_names = [self.norm_key] * batch_size - self.logger.debug("[NormDebug] dataset_names=%s", dataset_names) - - additional_inputs = {} - - # -------- proprioception / masks (batch) -------- - agent_pos_list = [] - agent_pos_mask_list = [] - dof_mask_list = [] - for obs in observation: - if "robot_state_action_data" in obs: - robot_state_action_data = obs["robot_state_action_data"] - - agent_pos = torch.from_numpy(robot_state_action_data.agent_pos) - # Normalize to [1, T, D] for batch cat - if agent_pos.dim() == 2: - agent_pos = agent_pos.unsqueeze(0) - agent_pos_list.append(agent_pos) - - agent_pos_mask = torch.from_numpy( - robot_state_action_data.agent_pos_mask - ) - if agent_pos_mask.dim() == 2: - agent_pos_mask = agent_pos_mask.unsqueeze(0) - agent_pos_mask_list.append(agent_pos_mask) - - dof_mask = torch.from_numpy(robot_state_action_data.dof_mask) - if dof_mask.dim() == 1: - dof_mask = dof_mask.unsqueeze(0) - dof_mask_list.append(dof_mask) - - # cat: [B, T, D] / [B, ...] - if len(agent_pos_list) > 0: - agent_pos = torch.cat(agent_pos_list, dim=0) - agent_pos_mask = torch.cat(agent_pos_mask_list, dim=0) - dof_mask = torch.cat(dof_mask_list, dim=0) - - if self.normalizer_propri is not None: - agent_pos = self.normalizer_propri.normalize_data( - agent_pos, dataset_names - ) - - additional_inputs["proprioception"] = agent_pos.detach() - additional_inputs["agent_pos_mask"] = agent_pos_mask - additional_inputs["dof_mask"] = dof_mask - - # -------- images (flattened, in placeholder scan order) -------- - with ScopeTimer("resize_images"): - # TODO[KC]: optimize this using batch processing - image_inputs = [] - all_image_sizes = [] - if ENABLE_FAST_PREPROCESS: - for obs in observation: - current_image_inputs = self._resize_images_fast(obs) - image_inputs.extend(current_image_inputs) - # tensor shape is (H, W, C), convert to (W, H) to match PIL.size format - all_image_sizes.extend( - [ - (image_i.shape[1], image_i.shape[0]) - for image_i in current_image_inputs - ] - ) - else: - for obs in observation: - current_image_inputs = self._resize_images(obs) - image_inputs.extend(current_image_inputs) - # tensor shape is (H, W, C), convert to (W, H) to match PIL.size format - all_image_sizes.extend( - [ - (image_i.shape[1], image_i.shape[0]) - for image_i in current_image_inputs - ] - ) - - additional_inputs["image_size"] = all_image_sizes - - with ScopeTimer("preprocesser_call"): - inputs = preprocesser_call( - processor=self.model.processor, - prefix_text=prefix_text, - postfix_text=postfix_text, - images=image_inputs, - videos=None, - padding=True, - truncation=True, - return_tensors="pt", - max_length=resolve_max_length(self.config.train_config), - pad_to_128_multiple=False, - pad_prefix_to_same_length=pad_prefix, - norm_state=( - additional_inputs["proprioception"] - if "proprioception" in additional_inputs - and resolve_use_state_string_representation( - self.config.train_config - ) - else None - ), - agent_pos_mask=( - additional_inputs["agent_pos_mask"] - if "agent_pos_mask" in additional_inputs - else None - ), - state_augmentation_prob=0.0, - state_drop_prob=0.0, - state_augmentation_ratio=0.0, - state_bins=resolve_state_bins(self.config.train_config), - inference_mode=True, - ) - - with ScopeTimer("convert_action_token_id and post "): - action_token_id = self.model.processor.tokenizer.convert_tokens_to_ids( - "<|action|>" - ) - flow_action_mask = inputs["input_ids"] == action_token_id - additional_inputs["moe_token_types"] = flow_action_mask - additional_inputs["dataset_names"] = dataset_names - - inputs.update(additional_inputs) - inputs = move_to_cuda(inputs, self.config.model_device) - - return inputs - - def get_text_for_dllm_action(self, instruction): - if ( - self.config.train_config["data"].get("use_embodied_system_prompt_ratio", 0) - > 0 - ): - if self.norm_key != "x2_normal" and self.norm_key != "ex_normal": - self.config.robot_id = 0 - cam_name_mapping = {cam_name: cam_name for cam_name in self.cam_names} - prologue = get_prologue_with_embodied_information( - dataset_name=self.norm_key, - cam_mapping=cam_name_mapping, - robot_id=self.config.robot_id, - uid="", - config=self.config.data_config, - ) - else: - prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n" - user_request = f"{self.role_start_symbol}user\nObservation:" - camera_name_mapping = self._camera_name_mapping() - for cam_name in self.cam_names: - user_request += ( - f" {resolve_camera_label(cam_name, camera_name_mapping)}:" - f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}" - ) - user_request += "\nInstruction:" - text_prompt = f"\nPredict the next action in robot action.\nProprioception: {self.propri_symbol}\n" - user_message = ( - f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n" - ) - placeholder_seq = self.tokenizer_mixin.get_placeholder_for_dllm() - ar_token = "".join(placeholder_seq) + "<|im_end|>\n" - assistant_message = f"{self.role_start_symbol}assistant\n{ar_token}" - flow_action = f"{self.action_symbol * self.config.action_horizon}" - - prefix_text = prologue + user_message + assistant_message - postfix_text = flow_action - - return prefix_text, postfix_text - - @timer - def get_text_for_action(self, instruction): - if ( - self.config.train_config["data"].get("use_embodied_system_prompt_ratio", 0) - > 0 - ): - if self.norm_key not in ["x2_normal", "ex_normal"]: - self.config.robot_id = 0 - cam_name_mapping = {cam_name: cam_name for cam_name in self.cam_names} - prologue = get_prologue_with_embodied_information( - dataset_name=self.norm_key, - cam_mapping=cam_name_mapping, - robot_id=self.config.robot_id, - uid="", - config=self.config.data_config, - ) - else: - prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n" - user_request = f"{self.role_start_symbol}user\nObservation:" - camera_name_mapping = self._camera_name_mapping() - for cam_name in self.cam_names: - user_request += ( - f" {resolve_camera_label(cam_name, camera_name_mapping)}:" - f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}" - ) - user_request += "\nInstruction:" - text_prompt = f"\nPredict the next action in robot action.\nProprioception: {self.propri_symbol}\n" - user_message = ( - f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n" - ) - assistant_message = f"{self.role_start_symbol}assistant\n" - flow_action = f"{self.action_symbol * self.config.action_horizon}" - - prefix_text = prologue + user_message + assistant_message - postfix_text = flow_action - - return prefix_text, postfix_text - - @timer - def get_text_for_subtask_generation(self, instruction): - prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n" - user_request = f"{self.role_start_symbol}user\nObservation:" - camera_name_mapping = self._camera_name_mapping() - for cam_name in self.cam_names: - user_request += ( - f" {resolve_camera_label(cam_name, camera_name_mapping)}:" - f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}" - ) - user_request += "\nInstruction:" - text_prompt = "\nPredict the next action in language.\n" - user_message = ( - f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n" - ) - assistant_message = f"{self.role_start_symbol}assistant\n" - - prefix_text = prologue + user_message + assistant_message - postfix_text = "" - - return prefix_text, postfix_text - - @timer - def _resize_images(self, observation): - image_inputs = [] - for key in self.cam_names: - if key not in observation: - continue - current_obs = observation[key] - if isinstance(current_obs, np.ndarray): - img_pil = Image.fromarray(current_obs) - elif isinstance(current_obs, Image.Image): - img_pil = current_obs - else: - raise ValueError(f"Unsupported image type: {type(current_obs)}") - orig_width, orig_height = img_pil.size - - target_size = self.config.data_config.resolution.get(key, -1) - if target_size != -1: - # Aspect-ratio-preserving resize - if orig_width > orig_height: # landscape - new_width = target_size - new_height = int(target_size * orig_height / orig_width) - else: # portrait - new_height = target_size - new_width = int(target_size * orig_width / orig_height) - img_pil = img_pil.resize((new_width, new_height)) - - # Apply smart resize (Qwen logic) - current_width, current_height = img_pil.size - resized_height, resized_width = smart_resize( - current_height, - current_width, - factor=self.config.data_config.image_factor, - min_pixels=self.config.data_config.min_pixels, - max_pixels=self.config.data_config.max_pixels, - ) - resized_img = img_pil.resize((resized_width, resized_height)) - resized_img = torch.from_numpy(np.array(resized_img)).to( - self.config.model_device - ) - image_inputs.append(resized_img) - - return image_inputs - - def _resize_images_fast(self, observation): - import cv2 - - image_inputs = [] - for key in self.cam_names: - if key not in observation: - continue - current_obs = observation[key] - orig_height, orig_width, _ = current_obs.shape - - target_size = self.config.data_config.resolution.get(key, -1) - current_width, current_height = orig_width, orig_height - if target_size != -1: - # Aspect-ratio-preserving resize - if orig_width > orig_height: # landscape - new_width = target_size - new_height = int(target_size * orig_height / orig_width) - else: # portrait - new_height = target_size - new_width = int(target_size * orig_width / orig_height) - current_width = new_width - current_height = new_height - - # Apply smart resize (Qwen logic) - resized_height, resized_width = smart_resize( - current_height, - current_width, - factor=self.config.data_config.image_factor, # FIXME - min_pixels=self.config.data_config.min_pixels, # FIXME - max_pixels=self.config.data_config.max_pixels, # FIXME - ) - - resized_img = cv2.resize( - current_obs, - (resized_width, resized_height), - interpolation=cv2.INTER_CUBIC, - ) - resized_img = torch.from_numpy(resized_img).to(self.config.model_device) - image_inputs.append(resized_img) - - return image_inputs - - def infer_flow_action(self, observation, instruction): - self.logger.info("generating flow action") - self.logger.info(f"flow action instruction: {instruction}") - - prefix_text, postfix_text = self.get_text_for_action(instruction) - model_input = self.construct_model_input( - [observation], [prefix_text], [postfix_text] - ) - - padding = ( - torch.zeros_like( - self.normalizer_action.delta[model_input["dataset_names"][0]] - ) - .unsqueeze(0) - .to("cpu") - ) - padding_action = self.normalizer_action.normalize_data( - padding, model_input["dataset_names"] - ).to(model_input["input_ids"].device) - - self.logger.info( - "generate_flow_action start (horizon=%s, flow_steps=%s, device=%s)", - self.config.action_horizon, - self.config.num_inference_timesteps, - self.config.model_device, - ) - if torch.cuda.is_available(): - try: - free_b, total_b = torch.cuda.mem_get_info( - torch.device(self.config.model_device) - ) - self.logger.info( - "CUDA mem before flow: free=%.2f GiB / total=%.2f GiB", - free_b / (1024**3), - total_b / (1024**3), - ) - except Exception as e: - self.logger.warning("CUDA mem_get_info failed: %s", e) - - with ScopeTimer("generate_flow_action"): - model_output = self.model.generate_flow_action( - action_horizon=self.config.action_horizon, - action_dim=self.config.action_dim, - num_inference_timesteps=self.config.num_inference_timesteps, - padding_action=padding_action, - **model_input, - ) - - self.logger.info("flow action generation done") - - model_output["robot_state_action_data"] = observation["robot_state_action_data"] - model_output["robot_state_action_data"].save_action_data( - model_output["predict_action"] - ) - self.logger.info("saved flow action to robot_state_action_data") - return model_output - - def infer_flow_action_batch(self, observations, instructions): - """ - Batch flow action inference: - - observations: List[Dict], same format as single inference - - instructions: List[str], aligned with observations - Returns List[model_output] of length batch size - """ - assert len(observations) == len( - instructions - ), "observations and instructions must have the same length" - batch_size = len(observations) - - prefix_list = [] - postfix_list = [] - for ins in instructions: - prefix_text, postfix_text = self.get_text_for_action(ins) - prefix_list.append(prefix_text) - postfix_list.append(postfix_text) - - # Build batch input in one preprocesser_call; avoid per-sample cat - batch_inputs = self.construct_model_input( - observations, prefix_list, postfix_list - ) - - padding_list = [] - for ds_name in batch_inputs["dataset_names"]: - padding = ( - torch.zeros_like(self.normalizer_action.delta[ds_name]) - .unsqueeze(0) - .to("cpu") - ) - padding_list.append(padding) - padding = torch.cat(padding_list, dim=0) - padding_action = self.normalizer_action.normalize_data( - padding, batch_inputs["dataset_names"] - ).to(batch_inputs["input_ids"].device) - - with ScopeTimer("generate_flow_action_batch"): - model_output = self.model.generate_flow_action( - action_horizon=self.config.action_horizon, - action_dim=self.config.action_dim, - num_inference_timesteps=self.config.num_inference_timesteps, - padding_action=padding_action, - **batch_inputs, - ) - - predict_action = model_output["predict_action"] # [B, H, D] - - outputs = [] - for i in range(batch_size): - single_action = predict_action[i : i + 1] - single_output = { - "predict_action": single_action, - "robot_state_action_data": observations[i]["robot_state_action_data"], - } - single_output["robot_state_action_data"].save_action_data( - single_output["predict_action"] - ) - outputs.append(single_output) - - return outputs - - def infer_ar_action(self, observation, instruction): - self.logger.info("generating ar action") - self.logger.info(f"ar action instruction: {instruction}") - - prefix_text, _ = self.get_text_for_action(instruction) - model_input = self.construct_model_input([observation], [prefix_text], [""]) - - model_output = self.model.generate_ar_action( - action_horizon=self.config.action_horizon, - action_dim=self.config.ar_action_dim, - num_inference_timesteps=self.config.num_inference_timesteps, - robot_type_id=self.robot_type_id, - **model_input, - ) - self.logger.info("ar action generation done") - - model_output["robot_state_action_data"] = observation["robot_state_action_data"] - model_output["robot_state_action_data"].save_action_data( - model_output["predict_action"] - ) - self.logger.info("saved ar action to robot_state_action_data") - return model_output - - def infer_subtask(self, observation, instruction): - self.logger.info("generating subtask") - self.logger.info(f"subtask instruction: {instruction}") - - prefix_text, postfix_text = self.get_text_for_subtask_generation(instruction) - model_input = self.construct_model_input([observation], [prefix_text], [""]) - - model_output = self.model.generate_text(**model_input) - subtask = model_output["predict_output_text"][0].split("<|im_end|>")[0].strip() - self.logger.info(f"subtask generation done, subtask: {subtask}") - return subtask - - def infer_vqa(self, observation, instruction): - self.logger.info("generating vqa answer") - - if isinstance(observation["multi_modal"], list): - orig_size = observation["multi_modal"][0].size - else: - orig_size = observation["multi_modal"].size - - prefix_text = instruction - model_input = self.construct_model_input([observation], [prefix_text], [""]) - - model_output = self.model.generate_text(**model_input) - answer = model_output["predict_output_text"][0].split("<|im_end|>")[0].strip() - self.logger.info(f"vqa answer done, answer: {answer}") - answer = reverse_grounding_points( - answer, - orig_size[1], - orig_size[0], - model_input["image_size"][0][1], - model_input["image_size"][0][0], - self.config.data_config.model_type, - ) - points = extract_grounding_points(answer) - - return { - "answer": answer, - "points": points, - } - - def infer_dllm_action( - self, observation, instruction, use_ar_action=False, dataset_name="x2_normal" - ): - assert self.model_type in ["qwen2_5"], "DLLM only supports qwen2_5" - prefix_text, postfix_text = self.get_text_for_dllm_action(instruction) - model_input = self.construct_model_input( - [observation], [prefix_text], [postfix_text], pad_prefix=True - ) - model_input = self.model.update_infer_dllm_position_mask(model_input) - total_ar_step = self.tokenizer_mixin.inference_ar_steps_for_dllm - cnt = 0 - while cnt < 3: # fast mode: at most 3 retries - model_output = self.model.generate_dllm_action( - action_horizon=self.config.action_horizon, - action_dim=self.config.action_dim, - ar_action_dim=self.config.ar_action_dim, - num_inference_timesteps=self.config.num_inference_timesteps, - use_ar_action=use_ar_action, - total_ar_step=total_ar_step, - robot_type_id=self.robot_type_id, # for v3.1 delta decode - **model_input, - ) - if model_output["predict_action"] is not None: - break - cnt += 1 - self.logger.warning(f"dllm action generation failed, retry {cnt}") - self.logger.info("dllm action generation done") - - model_output["robot_state_action_data"] = observation["robot_state_action_data"] - model_output["robot_state_action_data"].save_action_data( - model_output["predict_action"] - ) - self.logger.info("saved dllm action to robot_state_action_data") - return model_output diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py b/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py deleted file mode 100644 index 641bbc43..00000000 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py +++ /dev/null @@ -1,1181 +0,0 @@ -import numpy as np -from abc import ABC -import re - -from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig -from wall_x._vendor.harrix.serving._wallx_infer.utils import VehiclePoseHandler, UnifiedTrajectoryProcessor -from wall_x._vendor.harrix.serving._wallx_infer.base_dataclass import ( - RobotStateActionData, - dof_dims, -) -from wall_x._vendor.harrix.serving._wallx_infer.socket_controller import RobotController - -from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger - -robot_action_key_mapping = { - "follow_left_ee_cartesian_pos": "follow1_pos[:3]", - "follow_left_ee_rotation": "follow1_pos[3:6]", - "follow_left_gripper": "follow1_pos[6:7]", - "follow_right_ee_cartesian_pos": "follow2_pos[:3]", - "follow_right_ee_rotation": "follow2_pos[3:6]", - "follow_right_gripper": "follow2_pos[6:7]", - "head_actions": "head_pos", - "head_rotation": "head_pos", # match ex001 - "height": "lift", - "velocity_decomposed": "velocity_decomposed", - "velocity_decomposed_odom": "velocity_decomposed_odom", # match ex001 - "follow_left_arm_joint_cur": "follow1_joints_cur", # 1 -> 7 - "follow_right_arm_joint_cur": "follow2_joints_cur", # 1 -> 7 - "follow_left_gripper_cur": "follow1_joints_cur[-1:]", - "follow_right_gripper_cur": "follow2_joints_cur[-1:]", - "follow_left_arm_joint_pos": "follow1_joints", - "follow_right_arm_joint_pos": "follow2_joints", - "follow_left_wrench_ext_local_force": "follow1_end_effort_force", - "follow_left_wrench_ext_local_torque": "follow1_end_effort_torque", - "follow_right_wrench_ext_local_force": "follow2_end_effort_force", - "follow_right_wrench_ext_local_torque": "follow2_end_effort_torque", - "follow_left_wrench_ext_local_force_from_joint": "follow1_end_effort_force_from_joint", - "follow_left_wrench_ext_local_torque_from_joint": "follow1_end_effort_torque_from_joint", - "follow_right_wrench_ext_local_force_from_joint": "follow2_end_effort_force_from_joint", - "follow_right_wrench_ext_local_torque_from_joint": "follow2_end_effort_torque_from_joint", - "follow_left_wrench_ext_world_force": "follow1_wrench_ext_world_force", - "follow_left_wrench_ext_world_torque": "follow1_wrench_ext_world_torque", - "follow_left_wrench_ext_world_force_from_joint": "follow1_wrench_ext_world_force_from_joint", - "follow_left_wrench_ext_world_torque_from_joint": "follow1_wrench_ext_world_torque_from_joint", - "follow_right_wrench_ext_world_force": "follow2_wrench_ext_world_force", - "follow_right_wrench_ext_world_torque": "follow2_wrench_ext_world_torque", - "follow_right_wrench_ext_world_force_from_joint": "follow2_wrench_ext_world_force_from_joint", - "follow_right_wrench_ext_world_torque_from_joint": "follow2_wrench_ext_world_torque_from_joint", - "follow_left_arm_joint_dev": "follow1_joints_dev", - "follow_right_arm_joint_dev": "follow2_joints_dev", -} - - -def _parse_follow_pos(follow_pos) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Parse websocket ``follow{1,2}_pos``: pos3 + euler3 + one or more gripper dims.""" - arr = np.asarray(follow_pos, dtype=np.float64).reshape(-1) - if arr.shape[0] < 7: - raise ValueError(f"follow_pos expects 7 dims, got shape {arr.shape}") - return arr[:3], arr[3:6], arr[6:] - - -def ingest_websocket_follow_pos( - state: dict, - config: InferConfig, - robot_state_action_data: RobotStateActionData, -) -> None: - """Convert client ``follow1_pos``/``follow2_pos`` (14D) into train-config proprio.""" - from wall_x._vendor.harrix.utils.train_config import resolve_agent_pos_config - from wall_x._vendor.x2robot_utils import geometry as geom - - agent_cfg = resolve_agent_pos_config(config.train_config) - arm_ws_keys = ( - ("left", "follow1_pos", "follow_left_"), - ("right", "follow2_pos", "follow_right_"), - ) - for _side, ws_key, cfg_prefix in arm_ws_keys: - if ws_key not in state or state[ws_key] is None: - continue - pos, euler, grip = _parse_follow_pos(state[ws_key]) - rot6d = geom.euler_to_matrix_zyx_6d_nb(euler.reshape(1, 3)) - for key in agent_cfg: - if key == "action_padding" or not key.startswith(cfg_prefix): - continue - if "cartesian_pos" in key: - robot_state_action_data.save_state_data_with_key(pos, key) - elif "rotation_6d" in key.lower(): - robot_state_action_data.save_state_data_with_key( - rot6d.reshape(1, -1), key - ) - elif "rotation" in key: - robot_state_action_data.save_state_data_with_key(euler, key) - elif "gripper" in key: - robot_state_action_data.save_state_data_with_key( - grip[: int(agent_cfg[key])], key, gt_dim=int(agent_cfg[key]) - ) - - for key, dim in agent_cfg.items(): - if key == "action_padding" or key.startswith(("follow_left_", "follow_right_")): - continue - norm_key = key.replace("follow_", "").replace("master_", "") - if robot_state_action_data.data.get(f"state_{norm_key}") is None: - robot_state_action_data.save_state_data_with_key( - np.zeros(int(dim), dtype=np.float64), key - ) - - -def export_websocket_follow_pos( - robot_state_action_data: RobotStateActionData, -) -> dict[str, list]: - """Convert model ``robot_state_action_data`` back to ``follow1_pos``/``follow2_pos``.""" - left_arm_action = _stack_state_action_series(robot_state_action_data, side="left") - right_arm_action = _stack_state_action_series(robot_state_action_data, side="right") - return { - "follow1_pos": left_arm_action.tolist(), - "follow2_pos": right_arm_action.tolist(), - } - - -def _zero_state_series(key: str) -> np.ndarray: - """Fallback proprio when websocket state omits an arm (e.g. only follow2_pos).""" - dim = dof_dims.get(key, 1) - return np.zeros((1, dim), dtype=np.float64) - - -def _identity_state_rotation_6d() -> np.ndarray: - from wall_x._vendor.x2robot_utils import geometry as geom - - return geom.euler_to_matrix_zyx_6d_nb(np.zeros((1, 3), dtype=np.float64)) - - -def _as_2d_series(value, name: str) -> np.ndarray: - """Normalize state/action arrays to (T, D) for follow*_pos serialization.""" - arr = np.asarray(value, dtype=np.float64) - if arr.ndim == 0: - arr = arr.reshape(1, 1) - elif arr.ndim == 1: - arr = arr.reshape(1, -1) if arr.size <= 7 else arr.reshape(-1, 1) - elif arr.ndim > 2: - arr = arr.reshape(arr.shape[0], -1) - if arr.ndim != 2: - raise ValueError(f"{name}: expected 2D series after normalize, got {arr.shape}") - return arr - - -def _stack_state_action_series( - robot_state_action_data: RobotStateActionData, - *, - side: str, - pos_key: str = "ee_cartesian_pos", - rot_key: str = "ee_rotation", - grip_key: str = "gripper", -) -> np.ndarray: - """Build (1+horizon, 7) follow arm pose: pos3 + rot3 + grip1.""" - from wall_x._vendor.x2robot_utils import geometry as geom - - data = robot_state_action_data.data - prefix = "left" if side == "left" else "right" - - state_pos = data.get(f"state_{prefix}_{pos_key}") - action_pos = data.get(f"action_{prefix}_{pos_key}") - rel_pos = data.get(f"action_{prefix}_{pos_key}_relative") - if action_pos is None and rel_pos is not None: - base_pos = ( - _as_2d_series(state_pos, f"state_{prefix}_{pos_key}") - if state_pos is not None - else _zero_state_series(f"{prefix}_{pos_key}") - ) - action_pos = base_pos + np.asarray(rel_pos, dtype=np.float64) - if state_pos is None: - state_pos = _zero_state_series(f"{prefix}_{pos_key}") - - state_rot = data.get(f"state_{prefix}_{rot_key}") - action_rot = data.get(f"action_{prefix}_{rot_key}") - - rot6d = data.get(f"action_{prefix}_{rot_key}_6D") - if action_rot is None and rot6d is not None: - action_rot = geom.so3_to_euler_zyx_batch_nb(np.asarray(rot6d, dtype=np.float64)) - if action_rot is None: - rel6d = data.get(f"action_{prefix}_{rot_key}_6D_relative") - state6d = data.get(f"state_{prefix}_{rot_key}_6D") - if state6d is None and state_rot is not None: - state6d = geom.euler_to_matrix_zyx_6d_nb( - _as_2d_series(state_rot, f"state_{prefix}_{rot_key}") - ) - if state6d is None: - state6d = _identity_state_rotation_6d() - if rel6d is not None: - rot6d = geom.compose_state_and_delta_to_abs_6d( - np.asarray(rel6d, dtype=np.float64), - np.asarray(state6d, dtype=np.float64).reshape(-1), - ) - action_rot = geom.so3_to_euler_zyx_batch_nb(rot6d) - - if state_rot is None: - state6d = data.get(f"state_{prefix}_{rot_key}_6D") - if state6d is not None: - state_rot = geom.so3_to_euler_zyx_batch_nb( - np.asarray(state6d, dtype=np.float64) - ) - else: - state_rot = _zero_state_series(f"{prefix}_{rot_key}") - - state_grip = data.get(f"state_{prefix}_{grip_key}") - if state_grip is None: - state_grip = _zero_state_series(f"{prefix}_{grip_key}") - action_grip = data.get(f"action_{prefix}_{grip_key}") - - if action_pos is None or action_rot is None or action_grip is None: - missing = [ - name - for name, value in ( - (f"action_{prefix}_{pos_key}", action_pos), - (f"action_{prefix}_{rot_key}", action_rot), - (f"action_{prefix}_{grip_key}", action_grip), - ) - if value is None - ] - raise ValueError( - f"Cannot serialize {prefix} arm action; missing action fields: {missing}" - ) - - pos = np.concatenate( - [ - _as_2d_series(state_pos, f"state_{prefix}_{pos_key}"), - _as_2d_series(action_pos, f"action_{prefix}_{pos_key}"), - ], - axis=0, - ) - rot = np.concatenate( - [ - _as_2d_series(state_rot, f"state_{prefix}_{rot_key}"), - _as_2d_series(action_rot, f"action_{prefix}_{rot_key}"), - ], - axis=0, - ) - grip = np.concatenate( - [ - _as_2d_series(state_grip, f"state_{prefix}_{grip_key}"), - _as_2d_series(action_grip, f"action_{prefix}_{grip_key}"), - ], - axis=0, - ) - return np.concatenate([pos, rot, grip], axis=1) - - -def _views_to_camera_observation( - config: InferConfig, views: dict -) -> dict[str, np.ndarray]: - """Map websocket camera keys to model observation keys. - - Only cameras listed in ``config.cam_names`` are required (e.g. LIBERO uses - face_view + right_wrist_view without left_wrist_view). - """ - camera_mappings = { - "face_view": [config.camera_front_key, "face_view", "front_view"], - "left_wrist_view": [ - config.camera_left_key, - "left_wrist_view", - "left_view", - ], - "right_wrist_view": [ - config.camera_right_key, - "right_wrist_view", - "right_view", - ], - } - observation: dict[str, np.ndarray] = {} - for obs_key in config.cam_names: - possible_keys = camera_mappings.get(obs_key, [obs_key]) - for view_key in possible_keys: - if view_key in views and views[view_key] is not None: - view_data = views[view_key] - if ( - isinstance(view_data, np.ndarray) - and view_data.ndim == 4 - and view_data.shape[0] == 1 - ): - observation[obs_key] = view_data[0] - else: - observation[obs_key] = view_data - break - else: - raise KeyError(f"Missing view for {obs_key!r}, tried keys: {possible_keys}") - return observation - - -class Robot(ABC): - def __init__(self, config: InferConfig, robot_id=10) -> None: - self.config = config - - self.robot_controller = RobotController( - robot_id=robot_id, host=config.robot_host, port=config.robot_port - ) - - self.robot_controller.connect() - - self.is_received = False # blocking flag - self.current_robot_state_action = None - - self.logger = InferLogger.get_robot_logger("Robot") - - def _get_views_and_state(self): - state = self.robot_controller.recv_action() - views = self.robot_controller.recv_image( - [ - self.config.camera_left_key, - self.config.camera_front_key, - self.config.camera_right_key, - ] - ) - self.is_received = True - return state, views - - def _get_dof_mask(self): - dof_config = self.config.train_config["dof_config"] - total_dof = sum(dof_config.values()) - dof_mask = np.ones((1, self.config.action_horizon, total_dof)) - return dof_mask - - @staticmethod - def _try_compose_arm_joint_pos_with_gripper(state: dict, key: str): - """Build 7D arm joint position from runtime state when possible. - - Why this is needed: - - Joint-control training expects `*_arm_joint_pos` to be 7D - (6 arm joints + 1 gripper). - - Some online robot states expose arm joints as 6D (`follow*_joints`) - and gripper separately in `follow*_pos[6]`. - - If we pass raw 6D joints directly, `save_state_data_with_key` will fail - shape validation against dof_dims=7. - - Behavior: - - Only handles `follow_left_arm_joint_pos` / `follow_right_arm_joint_pos`. - - Returns a composed 7D vector when both required sources exist and shapes - match expected runtime format. - - Returns None for all other keys or incompatible state shapes, so caller - falls back to original mapping logic. - """ - - if key == "follow_left_arm_joint_pos": - joints_key = "follow1_joints" - ee_key = "follow1_pos" - elif key == "follow_right_arm_joint_pos": - joints_key = "follow2_joints" - ee_key = "follow2_pos" - else: - return None - - if joints_key not in state or ee_key not in state: - return None - - joints = np.asarray(state[joints_key]).reshape(-1) - ee = np.asarray(state[ee_key]).reshape(-1) - if joints.shape[0] == 6 and ee.shape[0] >= 7: - return np.concatenate([joints, ee[6:7]], axis=0) - - return None - - def get_observation(self): - state, views = self._get_views_and_state() - - robot_state_action_data = RobotStateActionData(config=self.config) - for key in robot_action_key_mapping.keys(): - state_key_str = robot_action_key_mapping[key] - value = None - # Prefer a robust 7D composition path for arm_joint_pos: - # online state often provides `follow*_joints` as 6D plus gripper - # in `follow*_pos[6]`. Compose first to match model/dof schema. - composed_value = self._try_compose_arm_joint_pos_with_gripper(state, key) - if composed_value is not None: - value = composed_value - robot_state_action_data.save_state_data_with_key( - np.asarray(value)[None], key - ) - continue - - match = re.match(r"(\w+)\[(.*)\]", state_key_str) - if match: - base_key = match.group(1) - if base_key in state: - slicing_str = match.group(2) - slice_parts = slicing_str.split(":") - slice_args = [(int(p) if p.strip() else None) for p in slice_parts] - s = slice(*slice_args) - value = state[base_key][s] - else: - if state_key_str in state: - value = state[state_key_str] - if value is not None: - robot_state_action_data.save_state_data_with_key( - np.asarray(value)[None], key - ) - - dof_mask = self._get_dof_mask() - robot_state_action_data.dof_mask = dof_mask - self.current_robot_state_action = robot_state_action_data - return { - "robot_state_action_data": robot_state_action_data, - **_views_to_camera_observation(self.config, views), - } - - def go_home(self): - if self.current_robot_state_action is None: - self.get_observation() - - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 3)), "left_ee_cartesian_pos" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 3)), "left_ee_rotation" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 1)), "left_gripper" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 3)), "right_ee_cartesian_pos" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 3)), "right_ee_rotation" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 1)), "right_gripper" - ) - - action_dict = {"robot_state_action_data": self.current_robot_state_action} - self.apply_action( - action_dict, - robot_action_interpolate_multiplier=150, - robot_action_start_ratio=0, - robot_action_end_ratio=1, - ) - self.logger.info("robot returned to initial pose") - - def _get_left_arm_action( - self, robot_state_action_data: RobotStateActionData - ) -> np.ndarray: - if not self.config.robot_use_joint_angle_control: - return _stack_state_action_series(robot_state_action_data, side="left") - left_arm_joint_pos = np.concatenate( - [ - _as_2d_series( - robot_state_action_data.data["state_left_arm_joint_pos"], - "state_left_arm_joint_pos", - ), - _as_2d_series( - robot_state_action_data.data["action_left_arm_joint_pos"], - "action_left_arm_joint_pos", - ), - ], - axis=0, - ) - return left_arm_joint_pos - - def _get_right_arm_action( - self, robot_state_action_data: RobotStateActionData - ) -> np.ndarray: - if not self.config.robot_use_joint_angle_control: - return _stack_state_action_series(robot_state_action_data, side="right") - right_arm_joint_pos = np.concatenate( - [ - _as_2d_series( - robot_state_action_data.data["state_right_arm_joint_pos"], - "state_right_arm_joint_pos", - ), - _as_2d_series( - robot_state_action_data.data["action_right_arm_joint_pos"], - "action_right_arm_joint_pos", - ), - ], - axis=0, - ) - return right_arm_joint_pos - - def get_serialized_actions( - self, - input: dict, - robot_action_interpolate_multiplier=None, - robot_action_start_ratio=None, - robot_action_end_ratio=None, - ) -> dict: - """Build trimmed+interpolated serialized actions from model_output (with robot_state_action_data); does not send. - - Used by apply_action and serving policy; returns a dict ready for robot or client. - """ - assert "robot_state_action_data" in input - - left_arm_action = self._get_left_arm_action(input["robot_state_action_data"]) - right_arm_action = self._get_right_arm_action(input["robot_state_action_data"]) - - # Trim action - action_length = len(left_arm_action) - if robot_action_start_ratio is None: - robot_action_start_ratio = self.config.robot_action_start_ratio - if robot_action_end_ratio is None: - robot_action_end_ratio = self.config.robot_action_end_ratio - start_frame = int(robot_action_start_ratio * action_length) - end_frame = int(robot_action_end_ratio * action_length) - - left_arm_action = left_arm_action[start_frame:end_frame] - right_arm_action = right_arm_action[start_frame:end_frame] - - # Interpolate - if robot_action_interpolate_multiplier is None: - robot_action_interpolate_multiplier = ( - self.config.robot_action_interpolate_multiplier - ) - target_length = robot_action_interpolate_multiplier * len(left_arm_action) - left_arm_action, right_arm_action = ( - UnifiedTrajectoryProcessor.interpolate_trajectory_batch( - [left_arm_action, right_arm_action], target_length - ) - ) - - if not self.config.robot_use_joint_angle_control: - return { - "follow1_pos": left_arm_action.tolist(), - "follow2_pos": right_arm_action.tolist(), - } - return { - "follow1_joints": left_arm_action.tolist(), - "follow2_joints": right_arm_action.tolist(), - } - - def apply_action( - self, - input: dict, - robot_action_interpolate_multiplier=None, - robot_action_start_ratio=None, - robot_action_end_ratio=None, - ) -> None: - serialized_actions = self.get_serialized_actions( - input, - robot_action_interpolate_multiplier=robot_action_interpolate_multiplier, - robot_action_start_ratio=robot_action_start_ratio, - robot_action_end_ratio=robot_action_end_ratio, - ) - self._send_actions(serialized_actions) - - def _send_actions(self, serialized_actions: dict) -> None: - """Send actions to robot controller""" - self.robot_controller.robot_comm.send_dict(serialized_actions) - self.is_received = False - self.current_robot_state_action = None - - -class DesktopRobot(Robot): - def __init__(self, config: InferConfig, robot_id=10) -> None: - self.is_received = False # blocking flag - self.current_robot_state_action = None - - self.logger = InferLogger.get_robot_logger("Robot") - - def _get_dof_mask(self): - from wall_x._vendor.harrix.utils.train_config import resolve_dof_config - - dof_config = resolve_dof_config(self.config.train_config) - total_dof = sum(dof_config.values()) - dof_mask = np.ones((1, self.config.action_horizon, total_dof)) - - # For desktop, mask out head_actions, height, and velocity_decomposed - mask_keys = ["head_actions", "height", "velocity_decomposed"] - start_idx = 0 - for key, dof_size in dof_config.items(): - if key in mask_keys: - # Set mask for these dims to 0 - dof_mask[:, :, start_idx : start_idx + dof_size] = 0 - start_idx += dof_size - - return dof_mask - - -class DesktopRobotVGA(Robot): - def __init__(self, config: InferConfig, robot_id=10) -> None: - super().__init__(config, robot_id) - self.is_received = False # blocking flag - self.current_robot_state_action = None - - self.logger = InferLogger.get_robot_logger("Robot") - - def _get_dof_mask(self): - dof_config = self.config.train_config["data"]["dof_config"] - total_dof = sum(dof_config.values()) - dof_mask = np.ones((1, self.config.action_horizon, total_dof)) - - # For desktop, mask out head_actions, height, and velocity_decomposed - mask_keys = ["head_actions", "height", "velocity_decomposed"] - start_idx = 0 - for key, dof_size in dof_config.items(): - if key in mask_keys: - # Set mask for these dims to 0 - dof_mask[:, :, start_idx : start_idx + dof_size] = 0 - start_idx += dof_size - - return dof_mask - - -class DesktopRobotPreprocessor(DesktopRobot): - def __init__(self, config: InferConfig, robot_id=10) -> None: - self.config = config - - self.is_received = False # blocking flag - self.current_robot_state_action = None - self.logger = InferLogger.get_robot_logger("Robot") - - def _send_actions(self, serialized_actions: dict) -> None: - # bypass - return - - def get_observation(self, state, views): - robot_state_action_data = RobotStateActionData(config=self.config) - if "follow1_pos" in state or "follow2_pos" in state: - ingest_websocket_follow_pos(state, self.config, robot_state_action_data) - for key in robot_action_key_mapping.keys(): - if key.startswith(("follow_left_", "follow_right_")): - continue - state_key_str = robot_action_key_mapping[key] - value = None - match = re.match(r"(\w+)\[(.*)\]", state_key_str) - if match: - base_key = match.group(1) - if base_key in state: - slicing_str = match.group(2) - slice_parts = slicing_str.split(":") - slice_args = [(int(p) if p.strip() else None) for p in slice_parts] - s = slice(*slice_args) - value = state[base_key][s] - else: - if state_key_str in state: - value = state[state_key_str] - if value is not None: - robot_state_action_data.save_state_data_with_key( - np.asarray(value)[None], key - ) - - dof_mask = self._get_dof_mask() - robot_state_action_data.dof_mask = dof_mask - self.current_robot_state_action = robot_state_action_data - return { - "robot_state_action_data": robot_state_action_data, - **_views_to_camera_observation(self.config, views), - } - - -class TurtleRobot(Robot): - def __init__(self, config: InferConfig, robot_id=10) -> None: - self.vehicle_pose_handler = VehiclePoseHandler() - self.last_speed = [0, 0, 0] - - def _get_state_from_controller(self): - state = super()._get_state_from_controller() - - self.vehicle_pose_handler.update_pose(state["car_pose"]) - state["velocity_decomposed"] = self.last_speed - - if self.config.turtle_as_desktop: - state["head_pos"] = [0, -1] - state["lift"] = [0.4] - - return state - - def _calculate_car_pose(self, base_velocity_pred): - # Vectorized integrate velocity to pose - dt = 1 / 20 - current_pose = ( - self.vehicle_pose_handler.current_pose.copy() - if self.vehicle_pose_handler.current_pose is not None - else np.array([0.0, 0.0, 0.0]) - ) - - # Batch integrate positions - poses_frames = [] - for i in range(len(base_velocity_pred)): - current_pose = self.vehicle_pose_handler.velocity_to_pose( - base_velocity_pred[i, 0], - base_velocity_pred[i, 1], - base_velocity_pred[i, 2], - dt, - current_pose, - ) - poses_frames.append(current_pose.copy()) - - return poses_frames - - def _get_car_velocity( - self, robot_state_action_data: RobotStateActionData - ) -> np.ndarray: - car_velocity = robot_state_action_data.data["action_velocity_decomposed"] - self.last_speed = car_velocity[-1, :].copy().tolist() - return car_velocity - - def _get_head_action( - self, robot_state_action_data: RobotStateActionData - ) -> np.ndarray: - head_actions = np.concatenate( - [ - robot_state_action_data.data["state_head_actions"], - robot_state_action_data.data["action_head_actions"], - ], - axis=0, - ) - return head_actions - - def _get_height_action( - self, robot_state_action_data: RobotStateActionData - ) -> np.ndarray: - height = np.concatenate( - [ - robot_state_action_data.data["state_height"], - robot_state_action_data.data["action_height"], - ], - axis=0, - ) - return height - - def go_home(self): - if self.current_robot_state_action is None: - self.get_observation() - - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 3)), "left_ee_cartesian_pos" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 3)), "left_ee_rotation" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 1)), "left_gripper" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 3)), "right_ee_cartesian_pos" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 3)), "right_ee_rotation" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 1)), "right_gripper" - ) - - self.current_robot_state_action.save_action_data_with_key( - np.array([[0, -1]]), "head_actions" - ) - self.current_robot_state_action.save_action_data_with_key( - np.array([[0.4]]), "height" - ) - self.current_robot_state_action.save_action_data_with_key( - np.array([self.last_speed]), "velocity_decomposed" - ) - - action_dict = {"robot_state_action_data": self.current_robot_state_action} - self.apply_action( - action_dict, - robot_action_interpolate_multiplier=150, - robot_action_start_ratio=0, - robot_action_end_ratio=1, - ) - self.logger.info("robot returned to initial pose") - - def get_serialized_actions( - self, - input: dict, - robot_action_interpolate_multiplier=None, - robot_action_start_ratio=None, - robot_action_end_ratio=None, - ) -> dict: - """Build serialized Turtle actions from model_output (dual arms + head/lift/car_pose); does not send.""" - assert "robot_state_action_data" in input - - left_arm_action = self._get_left_arm_action(input["robot_state_action_data"]) - right_arm_action = self._get_right_arm_action(input["robot_state_action_data"]) - action_length = len(left_arm_action) - if robot_action_start_ratio is None: - robot_action_start_ratio = self.config.robot_action_start_ratio - if robot_action_end_ratio is None: - robot_action_end_ratio = self.config.robot_action_end_ratio - start_frame = int(robot_action_start_ratio * action_length) - end_frame = int(robot_action_end_ratio * action_length) - - left_arm_action = left_arm_action[start_frame:end_frame] - right_arm_action = right_arm_action[start_frame:end_frame] - if robot_action_interpolate_multiplier is None: - robot_action_interpolate_multiplier = ( - self.config.robot_action_interpolate_multiplier - ) - target_length = robot_action_interpolate_multiplier * action_length - ( - left_arm_action, - right_arm_action, - ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( - [left_arm_action, right_arm_action], - target_length, - ) - serialized_actions = { - "follow1_pos": left_arm_action.tolist(), - "follow2_pos": right_arm_action.tolist(), - } - - if not self.config.turtle_as_desktop: - head_action = self._get_head_action(input["robot_state_action_data"]) - height_action = self._get_height_action(input["robot_state_action_data"]) - car_velocity_action = self._get_car_velocity( - input["robot_state_action_data"] - ) - head_action = head_action[start_frame:end_frame] - height_action = height_action[start_frame:end_frame] - car_velocity_action = car_velocity_action[start_frame:end_frame] - car_pose_action = np.array(self._calculate_car_pose(car_velocity_action)) - ( - head_action, - height_action, - car_pose_action, - ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( - [head_action, height_action, car_pose_action], - target_length, - ) - serialized_actions["head_pos"] = head_action.tolist() - serialized_actions["lift"] = height_action.tolist() - serialized_actions["car_pose"] = car_pose_action.tolist() - else: - serialized_actions["head_pos"] = [ - [0, -1] for _ in range(len(left_arm_action)) - ] - serialized_actions["lift"] = [0.4 for _ in range(len(left_arm_action))] - serialized_actions["car_pose"] = [ - [0.0, 0.0, 0.0] for _ in range(len(left_arm_action)) - ] - return serialized_actions - - def apply_action( - self, - input: dict, - robot_action_interpolate_multiplier=None, - robot_action_start_ratio=None, - robot_action_end_ratio=None, - ) -> None: - serialized_actions = self.get_serialized_actions( - input, - robot_action_interpolate_multiplier=robot_action_interpolate_multiplier, - robot_action_start_ratio=robot_action_start_ratio, - robot_action_end_ratio=robot_action_end_ratio, - ) - self._send_actions(serialized_actions) - - -class TurtleRobotPreprocessor(TurtleRobot): - """Turtle preprocessor: no real robot; builds observations and get_serialized_actions for serving.""" - - def __init__(self, config: InferConfig, robot_id=10) -> None: - self.config = config - self.is_received = False - self.current_robot_state_action = None - self.logger = InferLogger.get_robot_logger("Robot") - self.vehicle_pose_handler = VehiclePoseHandler() - self.last_speed = [0, 0, 0] - - def _send_actions(self, serialized_actions: dict) -> None: - return - - def get_observation(self, state, views): - robot_state_action_data = RobotStateActionData(config=self.config) - for key in robot_action_key_mapping.keys(): - state_key_str = robot_action_key_mapping[key] - value = None - match = re.match(r"(\w+)\[(.*)\]", state_key_str) - if match: - base_key = match.group(1) - if base_key in state: - slicing_str = match.group(2) - slice_parts = slicing_str.split(":") - slice_args = [(int(p) if p.strip() else None) for p in slice_parts] - s = slice(*slice_args) - value = state[base_key][s] - else: - if state_key_str in state: - value = state[state_key_str] - if value is not None: - robot_state_action_data.save_state_data_with_key( - np.asarray(value)[None], key - ) - dof_mask = self._get_dof_mask() - robot_state_action_data.dof_mask = dof_mask - self.current_robot_state_action = robot_state_action_data - return { - "robot_state_action_data": robot_state_action_data, - **_views_to_camera_observation(self.config, views), - } - - -class EX001Robot(Robot): - """EX001 robot: no pose integration; predict from velocity_decomposed_odom state and send to client.""" - - def __init__(self, config: InferConfig, robot_id=10) -> None: - super().__init__(config, robot_id) - self.last_speed = [0, 0, 0] - - def _get_velocity_decomposed( - self, robot_state_action_data: RobotStateActionData - ) -> np.ndarray: - """Get decomposed velocity actions: concat state and action velocity_decomposed_odom""" - velocity_decomposed = np.concatenate( - [ - robot_state_action_data.data["state_velocity_decomposed_odom"], - robot_state_action_data.data["action_velocity_decomposed_odom"], - ], - axis=0, - ) - self.last_speed = velocity_decomposed[-1, :].copy().tolist() - return velocity_decomposed - - def _get_head_action( - self, robot_state_action_data: RobotStateActionData - ) -> np.ndarray: - """Get head actions: concat state and action head sequence; supports head_rotation / head_actions""" - if robot_state_action_data.data.get("state_head_rotation") is not None: - state_key, action_key = "state_head_rotation", "action_head_rotation" - else: - state_key, action_key = "state_head_actions", "action_head_actions" - head_actions = np.concatenate( - [ - robot_state_action_data.data[state_key], - robot_state_action_data.data[action_key], - ], - axis=0, - ) - return head_actions - - def _get_height_action( - self, robot_state_action_data: RobotStateActionData - ) -> np.ndarray: - """Get lift height actions: concat state and action height sequence""" - height = np.concatenate( - [ - robot_state_action_data.data["state_height"], - robot_state_action_data.data["action_height"], - ], - axis=0, - ) - return height - - def go_home(self): - """Return to initial pose""" - if self.current_robot_state_action is None: - self.get_observation() - - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 3)), "left_ee_cartesian_pos" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 3)), "left_ee_rotation" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 1)), "left_gripper" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 3)), "right_ee_cartesian_pos" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 3)), "right_ee_rotation" - ) - self.current_robot_state_action.save_action_data_with_key( - np.zeros((1, 1)), "right_gripper" - ) - - self.current_robot_state_action.save_action_data_with_key( - np.array([[0, -1]]), "head_actions" - ) - self.current_robot_state_action.save_action_data_with_key( - np.array([[0.4]]), "height" - ) - self.current_robot_state_action.save_action_data_with_key( - np.array([self.last_speed]), "velocity_decomposed_odom" - ) - - action_dict = {"robot_state_action_data": self.current_robot_state_action} - self.apply_action( - action_dict, - robot_action_interpolate_multiplier=150, - robot_action_start_ratio=0, - robot_action_end_ratio=1, - ) - self.logger.info("robot returned to initial pose") - - def get_serialized_actions( - self, - input: dict, - robot_action_interpolate_multiplier=None, - robot_action_start_ratio=None, - robot_action_end_ratio=None, - ) -> dict: - """Build serialized EX001 actions from model_output (dual arms + head/lift/velocity_decomposed_odom); does not send. - - Uses predicted velocity_decomposed_odom directly; no pose integration. - """ - assert "robot_state_action_data" in input - - left_arm_action = self._get_left_arm_action(input["robot_state_action_data"]) - right_arm_action = self._get_right_arm_action(input["robot_state_action_data"]) - action_length = len(left_arm_action) - - # action_length is 1 state frame + action_horizon action frames - # Actual action length is action_length - 1 - actual_action_length = action_length - 1 - - if robot_action_start_ratio is None: - robot_action_start_ratio = self.config.robot_action_start_ratio - if robot_action_end_ratio is None: - robot_action_end_ratio = self.config.robot_action_end_ratio - - # Trim indices in action section (skip frame 0 state) - start_frame = 1 + int(robot_action_start_ratio * actual_action_length) - end_frame = 1 + int(robot_action_end_ratio * actual_action_length) - - left_arm_action = left_arm_action[start_frame:end_frame] - right_arm_action = right_arm_action[start_frame:end_frame] - - if robot_action_interpolate_multiplier is None: - robot_action_interpolate_multiplier = ( - self.config.robot_action_interpolate_multiplier - ) - target_length = robot_action_interpolate_multiplier * len(left_arm_action) - - ( - left_arm_action, - right_arm_action, - ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( - [left_arm_action, right_arm_action], - target_length, - ) - - serialized_actions = { - "follow1_pos": left_arm_action.tolist(), - "follow2_pos": right_arm_action.tolist(), - } - - # Get head, lift, and velocity actions - head_action = self._get_head_action(input["robot_state_action_data"]) - height_action = self._get_height_action(input["robot_state_action_data"]) - velocity_decomposed_action = self._get_velocity_decomposed( - input["robot_state_action_data"] - ) - - # Trim actions (same start_frame and end_frame) - head_action = head_action[start_frame:end_frame] - height_action = height_action[start_frame:end_frame] - velocity_decomposed_action = velocity_decomposed_action[start_frame:end_frame] - - # Interpolate - ( - head_action, - height_action, - velocity_decomposed_action, - ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( - [head_action, height_action, velocity_decomposed_action], - target_length, - ) - - serialized_actions["head_pos"] = head_action.tolist() - serialized_actions["lift"] = height_action.tolist() - # Send velocity_decomposed_odom directly; no pose integration - serialized_actions["velocity_decomposed_odom"] = ( - velocity_decomposed_action.tolist() - ) - return serialized_actions - - def apply_action( - self, - input: dict, - robot_action_interpolate_multiplier=None, - robot_action_start_ratio=None, - robot_action_end_ratio=None, - ) -> None: - serialized_actions = self.get_serialized_actions( - input, - robot_action_interpolate_multiplier=robot_action_interpolate_multiplier, - robot_action_start_ratio=robot_action_start_ratio, - robot_action_end_ratio=robot_action_end_ratio, - ) - self._send_actions(serialized_actions) - - -class EX001RobotPreprocessor(EX001Robot): - """EX001 preprocessor: no real robot; builds observations and get_serialized_actions for serving.""" - - def __init__(self, config: InferConfig, robot_id=10) -> None: - self.config = config - self.is_received = False - self.current_robot_state_action = None - self.logger = InferLogger.get_robot_logger("Robot") - self.last_speed = [0, 0, 0] - - def _send_actions(self, serialized_actions: dict) -> None: - """No action send; preprocessing only""" - return - - def get_observation(self, state, views): - """Build observation from external state and views""" - robot_state_action_data = RobotStateActionData(config=self.config) - for key in robot_action_key_mapping.keys(): - state_key_str = robot_action_key_mapping[key] - value = None - - # Keep the same 7D adaptation behavior as Robot.get_observation(), - # so offline/serving preprocessor path is consistent with runtime path. - composed_value = self._try_compose_arm_joint_pos_with_gripper(state, key) - if composed_value is not None: - value = composed_value - robot_state_action_data.save_state_data_with_key( - np.asarray(value)[None], key - ) - continue - - match = re.match(r"(\w+)\[(.*)\]", state_key_str) - if match: - base_key = match.group(1) - if base_key in state: - slicing_str = match.group(2) - slice_parts = slicing_str.split(":") - slice_args = [(int(p) if p.strip() else None) for p in slice_parts] - s = slice(*slice_args) - value = state[base_key][s] - else: - if state_key_str in state: - value = state[state_key_str] - if value is not None: - robot_state_action_data.save_state_data_with_key( - np.asarray(value)[None], key - ) - dof_mask = self._get_dof_mask() - robot_state_action_data.dof_mask = dof_mask - self.current_robot_state_action = robot_state_action_data - - # Build observation; handle key name mismatches - observation = { - "robot_state_action_data": robot_state_action_data, - } - - # Try multiple possible key names - camera_mappings = { - "face_view": [ - self.config.camera_front_key, - "face_view", - "front_view", - ], - "left_wrist_view": [ - self.config.camera_left_key, - "left_wrist_view", - "left_view", - ], - "right_wrist_view": [ - self.config.camera_right_key, - "right_wrist_view", - "right_view", - ], - } - - for obs_key in self.config.cam_names: - possible_view_keys = camera_mappings.get(obs_key, [obs_key]) - found = False - for view_key in possible_view_keys: - if view_key in views and views[view_key] is not None: - view_data = views[view_key] - # If view_data is (1, H, W, 3), take [0] - if ( - isinstance(view_data, np.ndarray) - and len(view_data.shape) == 4 - and view_data.shape[0] == 1 - ): - observation[obs_key] = view_data[0] - else: - observation[obs_key] = view_data - found = True - self.logger.info( - f"EX001RobotPreprocessor: using '{view_key}' as '{obs_key}'" - ) - break - if not found: - self.logger.error( - f"EX001RobotPreprocessor: no view for '{obs_key}', tried keys: {possible_view_keys}" - ) - # Black image placeholder - observation[obs_key] = np.zeros((480, 640, 3), dtype=np.uint8) - - return observation diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py b/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py deleted file mode 100644 index aa8f72a7..00000000 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py +++ /dev/null @@ -1,423 +0,0 @@ -import socket -import struct -import json -import cv2 -import numpy as np -import time -import select -import errno -import os -import pickle - -import threading - -from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger - -_DEFAULT_ROBOT_CONFIG = { - 0: {"host": "127.0.0.1", "action_port": 57749, "keyboard_port": 58849}, - 1: {"host": "127.0.0.1", "action_port": 57750, "keyboard_port": 58850}, - 2: {"host": "127.0.0.1", "action_port": 57751, "keyboard_port": 58851}, - 3: {"host": "127.0.0.1", "action_port": 57761, "keyboard_port": 58861}, -} - - -class RobotRegistry: - def __init__(self): - self.registry = _DEFAULT_ROBOT_CONFIG - self.logger = InferLogger.get_controller_logger("RobotRegistry") - self.logger.debug(f"init robot registry, registered {len(self.registry)} robots") - - def register_robot(self, robot_id, host, action_port, keyboard_port): - if robot_id in self.registry: - self.logger.warning(f"robot ID {robot_id} already registered") - return False - else: - self.registry[robot_id] = { - "host": host, - "action_port": action_port, - "keyboard_port": keyboard_port, - } - self.logger.info( - f"register robot ID={robot_id}: host={host}, action_port={action_port}, keyboard_port={keyboard_port}" - ) - return True - - def get_robot_info(self, robot_id): - if robot_id in self.registry: - info = self.registry[robot_id] - self.logger.debug(f"get robot ID={robot_id} info: {info}") - return info - else: - self.logger.error(f"robot ID={robot_id} not found") - return None - - def exist(self, robot_id): - return robot_id in self.registry - - -class RobotCommunication: - def __init__(self, robot_id, host=None, action_port=None, keyboard_port=None): - self.logger = InferLogger.get_controller_logger("RobotCommunication") - self.logger.debug(f"init robot comms: robot_id={robot_id}") - - self.robot_register = RobotRegistry() - if self.robot_register.exist(robot_id): - self.robot_info = self.robot_register.get_robot_info(robot_id) - self.logger.debug(f"using registered robot config: {robot_id}") - else: - self.robot_register.register_robot( - robot_id, host, action_port, keyboard_port - ) - self.robot_info = self.robot_register.get_robot_info(robot_id) - self.logger.debug(f"register new robot config: {robot_id}") - - self.action_sock = None - self.action_conn = None - - self.keyboard_sock = None - self.keyboard_conn = None - - self.client_socks = [] - - def connect(self): - self.logger.info("establishing socket connection...") - - self.action_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - self.action_sock.setblocking(True) - self.action_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - - host = self.robot_info["host"] - port = self.robot_info["action_port"] - self.action_sock.bind((host, port)) - self.action_sock.listen(1) - self.logger.info(f"listening on action port: {host}:{self.action_sock.getsockname()[1]}") - - action_thread = threading.Thread(target=self.handle_action_client) - - self.logger.debug("starting connection handler thread...") - action_thread.start() - - action_thread.join() - self.logger.info("socket connection established") - - def handle_action_client(self): - self.logger.debug("waiting for client connection...") - self.action_conn, addr = self.action_sock.accept() - self.logger.info(f"accepted connection from {addr}") - - def recv_image(self, index): - self.logger.debug(f"receiving image index={index}...") - image_size = struct.unpack(" dict: - if not self.views_path: - raise RuntimeError( - "DummyRobotController requires views_path or " - "WALL_X_DUMMY_VIEWS_PATH." - ) - self.logger.debug(f"load views from pickle: {cam_names}") - with open(self.views_path, "rb") as f: - all_views = pickle.load(f) - self.logger.debug(f"views loaded: {len(all_views)} views") - return all_views - - def recv_action(self): - if not self.state_path: - raise RuntimeError( - "DummyRobotController requires state_path or " - "WALL_X_DUMMY_STATE_PATH." - ) - self.logger.debug("load actions from pickle") - with open(self.state_path, "rb") as f: - all_actions = pickle.load(f) - self.logger.debug(f"actions loaded: {len(all_actions)} fields") - return all_actions - - -class RobotController: - def __init__( - self, - robot_id: int, - host: str = None, - port: str = None, - max_time_step: int = 10000, - ): - self.logger = InferLogger.get_controller_logger("RobotController") - self.logger.info( - f"init robot controller: robot_id={robot_id}, max_time_step={max_time_step}" - ) - - self.robot_comm = RobotCommunication(robot_id, host, port) - self.max_time_step = max_time_step - self.global_step = 0 - - def connect(self): - self.logger.info("connecting robot controller...") - self.robot_comm.connect() - self.logger.info("robot controller connected") - - def close(self): - self.logger.info("closing robot controller...") - self.robot_comm.close() - self.logger.info("robot controller closed") - - def prediction(self, views, actions) -> dict: - raise NotImplementedError - - def recv_image( - self, cam_names: list = ["camera_left", "camera_front", "camera_right"] - ) -> dict: - self.logger.debug(f"receive images: {cam_names}") - views = {} - for i, name in enumerate(cam_names): - image = np.array(self.robot_comm.recv_image(i)) - views[name] = image[None, :] - self.logger.debug(f"received {name}: shape={image.shape}") - self.logger.debug(f"all images received: {len(views)} views") - return views - - def recv_action(self): - self.logger.debug("receiving action data...") - action_data = self.robot_comm.recv_action_data() - return action_data - - def recv_keyboard_input(self): - self.robot_comm.accept_connections() - json_data = self.robot_comm.recv_keyboard_input() - if json_data is not None: - self.logger.debug( - f"keyboard: motionStatus={json_data.get('motionStatus')}, armMode={json_data.get('armMode')}" - ) - return json_data["motionStatus"], json_data["armMode"] - - return None, None - - def reset(self): - self.logger.info(f"reset controller: global_step {self.global_step} -> 0") - self.global_step = 0 - - def record_start(self): - self.logger.info("start recording...") - record_signal = {"cmd": "RECORD_START"} - self.robot_comm.send_dict(record_signal) - - def record_continue(self): - self.logger.debug("continue recording...") - record_signal = {"cmd": "RECORD_CONTINUE"} - self.robot_comm.send_dict(record_signal) - - def record_start_process(self): - if self.global_step == 0: - self.logger.info("waiting for START signal...") - action = None - while action != "START": - action, _ = self.recv_keyboard_input() - self.logger.info("received START signal") - self.record_start() - - self.record_continue() - - def set_zero(self): - self.logger.info("setting zero pose...") - record_signal = {"cmd": "INIT_ZERO", "gripper": [0.0, 0.0]} - self.robot_comm.send_dict(record_signal) - self.reset() - self.logger.info("zero pose set") - - def record_stop(self): - self.logger.info("stop recording...") - record_signal = {"cmd": "RECORD_STOP"} - self.robot_comm.send_dict(record_signal) - time.sleep(0.01) - - def recover_from_failure(self): - self.logger.warning("recovering from failure...") - record_signal = {"cmd": "TO_MASTER_SLAVE"} - self.robot_comm.send_dict(record_signal) - time.sleep(0.01) - - self.logger.info("waiting for START signal...") - action = None - while action != "START": - action, _ = self.recv_keyboard_input() - - self.logger.info("recovery done, resume recording") - self.record_start() - - def reset_to_zero(self): - self.logger.info("reset to zero pose...") - self.record_stop() - self.set_zero() - - def to_slave(self): - self.logger.info("switching to slave mode...") - record_signal = {"cmd": "TO_SLAVE"} - self.robot_comm.send_dict(record_signal) - time.sleep(0.01) - - def run( - self, - record_mode=False, - cam_names: list = ["camera_left", "camera_front", "camera_right"], - ): - - while self.global_step <= self.max_time_step: - if record_mode: - self.record_start_process() - - self.global_step += 1 - - action = self.recv_action() - view = self.recv_image(cam_names) - - pred = self.prediction(view, action) - self.robot_comm.send_dict(pred) - - if record_mode: - action, arm_mode = self.recv_keyboard_input() - - if action is not None and arm_mode is not None: - self.logger.info("action: %s, arm_mode: %s", action, arm_mode) - if action == "STOP" and arm_mode == "ARM_TEST_MODE_MS": - self.record_stop() - self.recover_from_failure() - action, arm_mode = None, None - while action != "STOP" or arm_mode != "ARM_TEST_MODE_S": - action, arm_mode = self.recv_keyboard_input() - self.record_stop() - self.to_slave() - self.reset() - elif arm_mode == "ARM_TEST_MODE_S" and action == "INIT": - self.reset_to_zero() diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/utils.py b/wall_x/_vendor/harrix/serving/_wallx_infer/utils.py deleted file mode 100644 index 41637683..00000000 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/utils.py +++ /dev/null @@ -1,280 +0,0 @@ -import numpy as np -from scipy.signal import savgol_filter -from scipy.spatial.transform import Rotation as R # TODO: convert to numba -from collections import deque -import threading - -from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger - - -class KeyboardThread(threading.Thread): - """ - Simple keyboard listener thread with stop and reset - """ - - def __init__(self): - self.should_reset = False - self.should_stop = False - self.new_instruction_index = None # Stores new instruction index - self.logger = InferLogger.get_utils_logger("KeyboardThread") - - super(KeyboardThread, self).__init__(name="keyboard-thread", daemon=True) - self.show_help() - self.start() - - def run(self): - """Listen for keyboard input""" - while True: - try: - user_input = input().strip().lower() - - if user_input in ["s", "stop"]: - self.should_stop = not self.should_stop - self.logger.info("[keyboard] stop signal sent") - - elif user_input in ["r", "reset"]: - self.logger.info("[keyboard] resetting...") - self.should_reset = True - self.logger.info("[keyboard] reset signal sent") - - elif user_input.isdigit(): - # Numeric input: switch instruction index - index = int(user_input) - self.new_instruction_index = index - self.logger.info(f"[keyboard] switched to instruction index: {index}") - - else: - self.logger.info(f"[keyboard] input received: {user_input}. no action taken.") - - except EOFError: - break - except Exception as e: - self.logger.error(f"[keyboard] error: {e}") - - def show_help(self): - self.logger.info( - "[keyboard] controls: 's' stop, 'r' reset, digit switches instruction index" - ) - - -# Arm trajectory parameters -ARM_MAX_VELOCITY = 0.02 -ARM_EXECUTION_HZ = 20 -ARM_MIN_EXECUTION_TIME = 5.0 -ARM_MAX_EXECUTION_TIME = 15.0 - - -class UnifiedTrajectoryProcessor: - """Unified trajectory processor""" - - @staticmethod - def interpolate_trajectory_batch(trajectories, target_length, smooth=True): - """ - Interpolate multiple trajectories to a common length - Args: - trajectories: list of np.array, each shape (N, D) - target_length: int, target length - smooth: bool, whether to smooth - Returns: - list of np.array, interpolated trajectories - """ - if not trajectories: - return [] - - results = [] - for traj in trajectories: - if len(traj) == 0: - results.append(np.zeros((target_length, traj.shape[1]))) - continue - - if len(traj) == target_length: - results.append(traj) - continue - - # Vectorized interpolation - original_indices = np.linspace(0, len(traj) - 1, len(traj)) - target_indices = np.linspace(0, len(traj) - 1, target_length) - - # Handle different data types - if traj.shape[1] == 7: # Arm data [x,y,z,rx,ry,rz,gripper] - interpolated = UnifiedTrajectoryProcessor._interpolate_arm_trajectory( - traj, original_indices, target_indices, target_length - ) - else: # Other data (height, current, etc.) - interpolated = np.zeros((target_length, traj.shape[1])) - for i in range(traj.shape[1]): - interpolated[:, i] = np.interp( - target_indices, original_indices, traj[:, i] - ) - - # Smoothing - if smooth and len(interpolated) >= 5: - interpolated = UnifiedTrajectoryProcessor._smooth_trajectory( - interpolated - ) - - results.append(interpolated) - - return results - - @staticmethod - def _interpolate_arm_trajectory( - traj, original_indices, target_indices, target_length - ): - """Optimized arm trajectory interpolation""" - interpolated = np.zeros((target_length, 7)) - - # Vectorized interp for position and gripper - for i in [0, 1, 2, 6]: # x, y, z, gripper - interpolated[:, i] = np.interp(target_indices, original_indices, traj[:, i]) - - # Quaternion interpolation (vectorized) - quaternions = R.from_euler("xyz", traj[:, 3:6]).as_quat() - interpolated_quats = np.zeros((target_length, 4)) - for i in range(4): - interpolated_quats[:, i] = np.interp( - target_indices, original_indices, quaternions[:, i] - ) - - # Batch normalize - norms = np.linalg.norm(interpolated_quats, axis=1, keepdims=True) - interpolated_quats = interpolated_quats / norms - - # Batch convert back to euler - interpolated[:, 3:6] = R.from_quat(interpolated_quats).as_euler("xyz") - - return interpolated - - @staticmethod - def _interpolate_position_trajectory( - traj, original_indices, target_indices, target_length - ): - """Optimized position trajectory interpolation""" - interpolated = np.zeros((target_length, 3)) - for i in range(3): - interpolated[:, i] = np.interp(target_indices, original_indices, traj[:, i]) - return interpolated - - @staticmethod - def _smooth_trajectory(trajectory): - """Vectorized smoothing""" - if len(trajectory) < 5: - return trajectory - - try: - # Smooth all dimensions in batch - smoothed = np.zeros_like(trajectory) - for dim in range(trajectory.shape[1]): - smoothed[:, dim] = savgol_filter( - trajectory[:, dim], - min( - 5, - ( - len(trajectory) - if len(trajectory) % 2 == 1 - else len(trajectory) - 1 - ), - ), - 3, - mode="nearest", - ) - return smoothed - except Exception: - return trajectory - - @staticmethod - def calculate_optimal_trajectory_length(left_traj, right_traj): - """Compute optimal trajectory length""" - - # Vectorized distance computation - def calc_distance(traj): - if len(traj) < 2: - return 0.0 - pos_diff = traj[1:, :3] - traj[:-1, :3] - return np.sum(np.linalg.norm(pos_diff, axis=1)) - - distances = [calc_distance(left_traj), calc_distance(right_traj)] - max_distance = max(distances) - - if max_distance > 1e-6: - execution_time = np.clip( - max_distance / ARM_MAX_VELOCITY, - ARM_MIN_EXECUTION_TIME, - ARM_MAX_EXECUTION_TIME, - ) - else: - execution_time = ARM_MIN_EXECUTION_TIME - - return max(int(execution_time * ARM_EXECUTION_HZ), len(left_traj)) - - -class VehiclePoseHandler: - """Vehicle pose and velocity computation""" - - def __init__(self): - self.current_pose = None - self.previous_pose = None - self.pose_history = deque(maxlen=10) - self.logger = InferLogger.get_utils_logger("VehiclePoseHandler") - - def update_pose(self, new_pose): - """Update vehicle pose""" - if new_pose is not None: - self.previous_pose = self.current_pose - self.current_pose = np.array(new_pose) - self.pose_history.append(self.current_pose.copy()) - self.logger.info("current_pose %s", self.current_pose) - return self.current_pose - - def velocity_to_pose(self, vx_body, vy_body, vyaw, dt, start_pose=None): - """Convert body-frame velocity to global pose""" - if start_pose is None: - if self.current_pose is not None: - start_pose = self.current_pose.copy() - else: - start_pose = np.array([0.0, 0.0, 0.0]) - - x, y, theta = start_pose - - # Body velocity to global displacement - cos_theta = np.cos(theta) - sin_theta = np.sin(theta) - - # Transform: body frame -> global frame - dx_global = (vx_body * cos_theta - vy_body * sin_theta) * dt - dy_global = (vx_body * sin_theta + vy_body * cos_theta) * dt - dtheta = vyaw * dt - - # Compute new pose - x_new = x + dx_global - y_new = y + dy_global - theta_new = theta + dtheta - - # Wrap angle to [-pi, pi] - theta_new = (theta_new + np.pi) % (2 * np.pi) - np.pi - - return np.array([x_new, y_new, theta_new]) - - def compute_body_velocities_from_poses( - self, current_pose, previous_pose, dt=1 / 20 - ): - """Compute body-frame velocity from pose delta""" - if current_pose is None or previous_pose is None: - return np.array([0.0, 0.0, 0.0]) - - # Global-frame displacement - dx_global = current_pose[0] - previous_pose[0] - dy_global = current_pose[1] - previous_pose[1] - dtheta = current_pose[2] - previous_pose[2] - - # Use previous frame angle for transform - theta = previous_pose[2] - cos_theta = np.cos(theta) - sin_theta = np.sin(theta) - - # Global displacement to body velocity - vx_body = (dx_global * cos_theta + dy_global * sin_theta) / dt - vy_body = (-dx_global * sin_theta + dy_global * cos_theta) / dt - vyaw = dtheta / dt - - return np.array([vx_body, vy_body, vyaw]) diff --git a/wall_x/_vendor/harrix/serving/handlers/__init__.py b/wall_x/_vendor/harrix/serving/handlers/__init__.py new file mode 100644 index 00000000..37946f6a --- /dev/null +++ b/wall_x/_vendor/harrix/serving/handlers/__init__.py @@ -0,0 +1 @@ +"""Serving request handlers and their lazy registry.""" diff --git a/wall_x/_vendor/harrix/serving/handlers/base.py b/wall_x/_vendor/harrix/serving/handlers/base.py new file mode 100644 index 00000000..f97d1c5c --- /dev/null +++ b/wall_x/_vendor/harrix/serving/handlers/base.py @@ -0,0 +1,78 @@ +"""Serving request-handler contract.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, Dict, List + +from wall_x._vendor.harrix.model_executors.base import ServingExecutor + + +class ServingHandler: + """Translate serving requests into inference results. + + ``infer``/``infer_batch`` remain the compatibility ABI during migration. + """ + + def handle(self, request: Dict, session: Any = None) -> Dict: + """Canonical request entry point; ``session`` is reserved for migration.""" + + return self.infer(request) + + def handle_batch(self, requests: List[Dict]) -> List[Dict]: + return self.infer_batch(requests) + + def infer(self, obs: Dict) -> Dict: + raise NotImplementedError + + def infer_batch(self, obs_list: List[Dict]) -> List[Dict]: + return [self.infer(obs) for obs in obs_list] + + def reset(self) -> None: + pass + + @property + def metadata(self) -> Dict[str, Any]: + return {} + + +class ExecutorServingHandler(ServingHandler): + """Expose an executor through the serving lifecycle contract.""" + + def __init__(self, executor: ServingExecutor) -> None: + missing = [ + name + for name in ("infer", "infer_batch", "reset") + if not callable(getattr(executor, name, None)) + ] + if not hasattr(executor, "metadata"): + missing.append("metadata") + if missing: + raise TypeError( + "ExecutorServingHandler requires a ServingExecutor; " + f"missing capabilities: {', '.join(missing)}" + ) + if not isinstance(executor.metadata, Mapping): + raise TypeError( + "ExecutorServingHandler requires ServingExecutor.metadata " + f"to be a mapping property, got {type(executor.metadata).__name__}" + ) + self.executor = executor + + def infer(self, obs: Dict) -> Dict: + return self.executor.infer(obs) + + def infer_batch(self, obs_list: List[Dict]) -> List[Dict]: + return self.executor.infer_batch(obs_list) + + def reset(self) -> None: + self.executor.reset() + + @property + def metadata(self) -> Dict[str, Any]: + return self.executor.metadata + + +BasePolicy = ServingHandler + +__all__ = ["BasePolicy", "ExecutorServingHandler", "ServingHandler"] diff --git a/wall_x/_vendor/harrix/serving/handlers/wall_x_vanilla_policy.py b/wall_x/_vendor/harrix/serving/handlers/wall_x_vanilla_policy.py new file mode 100644 index 00000000..952da884 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/handlers/wall_x_vanilla_policy.py @@ -0,0 +1,146 @@ +"""Vanilla (open-source) Wall-X serving policy. + +This is the opt-engine-free counterpart to :class:`WallXPolicy`. It runs flow +inference through the same vanilla model path the LIBERO evaluation adapter uses +(``QwenVLActInferAdapter``), so it has no dependency on the closed +``harrix.model_executors`` inference engine or on internal robot preprocessors. + +Contract (single observation): + + obs = { + "state": { # proprioception, one entry per model dof group + "eef_pos": [x, y, z], # end-effector position (3) + "eef_axisangle": [rx, ry, rz], # end-effector rotation, axis-angle (3) + "gripper": [g0, ...], # gripper qpos (1 or 2 dims) + }, + "views": { # camera images, base64-encoded or raw ndarray + "face_view": , + "right_wrist_view": , + }, + "instruction": "pick up the red block", + } + +``infer`` returns the raw model action chunk:: + + {"action": np.ndarray of shape (action_horizon, action_dim)} + +Downstream robot-specific serialization/interpolation is intentionally left to +the client, since real-robot preprocessors are not part of the open-source cut. +""" + +import base64 +import logging +from typing import Any, Dict + +import numpy as np + +from wall_x._vendor.harrix.serving.handlers.base import BasePolicy + +logger = logging.getLogger(__name__) + +INFER_MODE_FLOW = "flow" + + +class WallXVanillaPolicy(BasePolicy): + """Opt-free Wall-X policy: vanilla flow inference, raw action chunks.""" + + def __init__( + self, + eval_config, + image_passing_mode: str = "base64", + ): + """Build the vanilla inference adapter from a harrix ``EvalConfig``. + + Args: + eval_config: ``harrix.e2e_infer.config.EvalConfig`` with + ``model.checkpoint_path`` + (and optionally ``model.train_config_path`` / ``model.action_horizon``) + populated. + image_passing_mode: ``"base64"`` (client sends JPEG/PNG bytes, base64) or + ``"numpy"`` (client sends raw H×W×C uint8 arrays). + """ + import wall_x._vendor.harrix.e2e_infer.adapters # noqa: F401 registers model adapters + from wall_x._vendor.harrix.e2e_infer.adapters.registry import build_adapter + + if image_passing_mode not in ("base64", "numpy"): + raise ValueError( + f"image_passing_mode must be 'base64' or 'numpy', got " + f"{image_passing_mode!r}" + ) + self.image_passing_mode = image_passing_mode + self._adapter = build_adapter(eval_config) + self._action_horizon = int(self._adapter.chunk_horizon) + self._action_dim = int(self._adapter.action_dim) + logger.info( + "WallXVanillaPolicy ready: action_horizon=%d action_dim=%d image_mode=%s", + self._action_horizon, + self._action_dim, + self.image_passing_mode, + ) + + # ── BasePolicy API ─────────────────────────────────────────────────── + + def reset(self) -> None: + """Stateless policy: nothing to reset.""" + + @property + def metadata(self) -> Dict[str, Any]: + return { + "engine": "vanilla", + "infer_mode": INFER_MODE_FLOW, + "action_horizon": self._action_horizon, + "action_dim": self._action_dim, + "image_passing_mode": self.image_passing_mode, + } + + def infer(self, obs: Dict) -> Dict: + """Run one flow-inference pass; return the raw action chunk. + + Delegates to the adapter's validated ``predict_batch`` (the same path the + LIBERO evaluation uses), which handles proprio encoding, prompt + construction, and flow sampling internally. + """ + payload = { + "observation": self._assemble_observation(obs), + "instruction": obs.get("instruction", ""), + "noise": obs.get("noise"), + } + actions = self._adapter.predict_batch([payload]) + return {"action": np.asarray(actions[0])} + + # ── observation assembly ───────────────────────────────────────────── + + def _decode_view(self, value) -> np.ndarray: + """Decode one camera view into an (H, W, 3) RGB uint8 array.""" + if self.image_passing_mode == "base64": + import cv2 + + img_bytes = base64.b64decode(value) + img_array = np.frombuffer(img_bytes, dtype=np.uint8) + decoded = cv2.imdecode(img_array, cv2.IMREAD_COLOR) + if decoded is None: + raise ValueError("failed to decode base64 image (cv2.imdecode)") + return cv2.cvtColor(decoded, cv2.COLOR_BGR2RGB) + arr = np.asarray(value, dtype=np.uint8) + if arr.ndim != 3 or arr.shape[-1] != 3: + raise ValueError( + f"numpy image must be (H, W, 3), got shape {arr.shape}" + ) + return arr + + def _assemble_observation(self, obs: Dict) -> Dict: + """Build the compact ndarray payload the adapter's encoder expects. + + The adapter's ``predict_batch`` runs ``encode_proprio`` on this dict, so + the keys mirror the LIBERO evaluation payload + (see ``harrix.e2e_infer.envs.libero_common``). + """ + state = obs["state"] + views = obs["views"] + return { + "eef_pos": np.asarray(state["eef_pos"], dtype=np.float32), + "eef_axisangle": np.asarray(state["eef_axisangle"], dtype=np.float32), + "gripper": np.asarray(state["gripper"], dtype=np.float32), + "face_view": self._decode_view(views["face_view"]), + "wrist_view": self._decode_view(views["right_wrist_view"]), + } diff --git a/wall_x/_vendor/harrix/serving/launch_serving.py b/wall_x/_vendor/harrix/serving/launch_serving.py deleted file mode 100644 index 9b9c0f4f..00000000 --- a/wall_x/_vendor/harrix/serving/launch_serving.py +++ /dev/null @@ -1,280 +0,0 @@ -#!/usr/bin/env python3 -""" -Server script for Wall-X model. - -This script serves a Wall-X model using a websocket server, allowing -clients to connect and get action predictions from observations. - -""" - -import dataclasses -import enum -import inspect -import logging -import os -import socket -import sys -import yaml -import traceback -import tyro - -from wall_x._vendor.harrix.serving.websocket_policy_server import WebsocketPolicyServer - - -def _server_model_config_to_infer_kwargs(model_config) -> dict: - from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig - - infer_params = inspect.signature(InferConfig.__init__).parameters - return { - k: v - for k, v in vars(model_config).items() - if v is not None and k in infer_params - } - - -def get_wallx_policy(model_config, image_passing_mode, serialize_actions=True): - from wall_x._vendor.harrix.serving.policy.wall_x_policy import WallXPolicy - from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig - - config = InferConfig(**_server_model_config_to_infer_kwargs(model_config)) - return WallXPolicy(config=config, image_passing_mode=image_passing_mode, serialize_actions=serialize_actions) - - - - -logger = logging.getLogger(__name__) - - -class EnvMode(enum.Enum): - """Supported environments/datasets.""" - - X2ROBOT = "x2robot" - LIBERO = "libero" - - -@dataclasses.dataclass -class ServerModelConfig: - """Configuration for loading a Wall-X model.""" - - checkpoint_path: str | None = None - train_config_path: str | None = None - # robot_host: str = '0.0.0.0' - # robot_port: int = 33723 - robot_type: str = "desktop" # ["desktop", "turtle"] - robot_action_start_ratio: float = ( - 0.0 # proportion of action execution to start from - ) - robot_action_end_ratio: float = 1.0 # proportion of action execution to end at - robot_action_interpolate_multiplier: int = 10 # action interpolation multiplier - robot_use_joint_angle_control: bool = ( - False # use joint angle control (model must predict joints) - ) - turtle_as_desktop: bool = ( - False # use turtle platform as desktop with fixed base/head/camera/height - ) - action_horizon: int = 32 # specify the correct horizon for the model - action_dim: int | None = None - model_device: str = "cuda" - num_inference_timesteps: int = 10 - num_inference_steps: int | None = None - cfg_scale: float | None = None - seed: int | None = None - save_video_dir: str = "./videos" - - # Please specify explicitly if the checkpoint was not trained on the x2robot dataset. - norm_key: str | None = None - # Model cameras; None = infer from train config ``data.key_mappings.camera``. - cam_names: list[str] | None = None - # Robot camera keys in incoming websocket observations. - camera_front_key: str = "camera_front" - camera_left_key: str = "camera_left" - camera_right_key: str = "camera_right" - # Serving prompt controls. If the client request has no instruction, - # default_instruction is used. prompt_template follows train-config semantics. - default_instruction: str | None = None - prompt_template: str | None = None - qwen25_prompt_template: str | None = None - prompt_priority_order: str | None = None - -@dataclasses.dataclass -class Args: - """Arguments for the serve_wall_x script.""" - - # Environment mode (used for default configurations) - env: EnvMode = EnvMode.X2ROBOT - - # Model configuration. If not provided, uses default config for the environment - model_config: ServerModelConfig | None = None - - # Default text prompt to use if not provided in observation - default_prompt: str | None = None - - # Port to serve the policy on - port: int = 43007 - - # Host to bind the server to - host: str = "0.0.0.0" - - # Enable debug logging - debug: bool = False - - # Image passing mode - image_passing_mode: str = "base64" # ["numpy", "base64"] - - # Model type - model_type: str = "wallx" # OSS supports qwen2.5 Wall-X only - - # Serialize actions via robot_preprocessor (True for robot control, False for raw output) - serialize_actions: bool = True - - # -- Dynamic batching ----------------------------------------- - # Set max_batch_size to enable dynamic batching. None = single mode. - max_batch_size: int | None = None - max_wait_time_ms: float = 0 - max_queue_size: int = 100 - timeout_ms: float = 30000 - - # -- Engine flags --------------------------------------------- - enable_experimental_engine: bool = False - enable_cuda_graph: bool = False - - -# Default model configurations for each environment -DEFAULT_CONFIGS: dict[EnvMode, ServerModelConfig] = { - EnvMode.X2ROBOT: ServerModelConfig( - checkpoint_path=None, - train_config_path=None, - robot_action_start_ratio=0.0, - robot_action_end_ratio=1.0, - robot_action_interpolate_multiplier=10, - robot_use_joint_angle_control=False, - turtle_as_desktop=False, - action_horizon=32, - action_dim=None, - model_device="cuda", - num_inference_timesteps=10, - ), - EnvMode.LIBERO: ServerModelConfig( - checkpoint_path=None, - train_config_path=None, - robot_type="desktop", - robot_action_start_ratio=0.0, - robot_action_end_ratio=1.0, - robot_action_interpolate_multiplier=1, - action_horizon=10, - action_dim=None, - model_device="cuda", - num_inference_timesteps=10, - cam_names=None, - ), -} - - -def get_model_config(args: Args) -> ServerModelConfig: - """Get model configuration from args or defaults.""" - if args.model_config is not None: - return args.model_config - - if config := DEFAULT_CONFIGS.get(args.env): - logger.info(f"Using default configuration for {args.env.value}") - return config - - raise ValueError( - f"No default configuration for {args.env.value}. " - f"Please provide --model-config with model_path and action_tokenizer_path." - ) - - -def create_policy(args: Args): - """Create a policy from the given arguments.""" - model_config = get_model_config(args) - if args.model_type != "wallx": - raise ValueError( - f"Unsupported model type: {args.model_type!r}. " - "The public package only supports model_type='wallx'." - ) - policy = get_wallx_policy( - model_config, args.image_passing_mode, args.serialize_actions - ) - return policy - - -def main(args: Args) -> None: - """Main function to start the Wall-X model server.""" - log_level = logging.DEBUG if args.debug else logging.INFO - logging.basicConfig( - level=log_level, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", - ) - - # Set engine environment variables - if args.enable_experimental_engine: - os.environ["ENABLE_EXPERIMENTAL_INFERENCE_ENGINE"] = "true" - logger.info("ENABLE_EXPERIMENTAL_INFERENCE_ENGINE=true") - - if args.enable_cuda_graph: - os.environ["ENABLE_CUDA_GRAPH"] = "true" - logger.info("ENABLE_CUDA_GRAPH=true") - - logger.info("Starting model server") - logger.info(f"Model type: {args.model_type}") - logger.info(f"Environment: {args.env.value}") - logger.info(f"Port: {args.port}") - logger.info(f"Host: {args.host}") - logger.info(f"Serialize actions: {args.serialize_actions}") - - # Create policy - try: - policy = create_policy(args) - except Exception as e: - logger.error(f"Failed to create policy: {e}") - logger.error(traceback.format_exc()) - sys.exit(1) - - # Get policy metadata - policy_metadata = policy.metadata - policy_metadata["env"] = args.env.value - - # Get network info - hostname = socket.gethostname() - try: - local_ip = socket.gethostbyname(hostname) - except Exception: - local_ip = "unknown" - - logger.info(f"Server hostname: {hostname}") - logger.info(f"Server IP: {local_ip}") - logger.info(f"Server will be available at: ws://{args.host}:{args.port}") - logger.info(f"Health check endpoint: http://{args.host}:{args.port}/healthz") - - batching_str = ( - f"batch_size={args.max_batch_size}, wait={args.max_wait_time_ms}ms" - if args.max_batch_size - else "disabled" - ) - logger.info(f"Batching: {batching_str}") - - # Create and start server - server = WebsocketPolicyServer( - policy=policy, - host=args.host, - port=args.port, - metadata=policy_metadata, - max_batch_size=args.max_batch_size, - max_wait_time_ms=args.max_wait_time_ms, - max_queue_size=args.max_queue_size, - timeout_ms=args.timeout_ms, - ) - - logger.info("Starting server...") - try: - server.serve_forever() - except KeyboardInterrupt: - logger.info("Server stopped by user") - except Exception as e: - logger.error(f"Server error: {e}") - sys.exit(1) - - -if __name__ == "__main__": - main(tyro.cli(Args)) diff --git a/wall_x/_vendor/harrix/serving/launch_vanilla_serving.py b/wall_x/_vendor/harrix/serving/launch_vanilla_serving.py new file mode 100644 index 00000000..12d48ea8 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/launch_vanilla_serving.py @@ -0,0 +1,126 @@ +"""Launch a vanilla (open-source) Wall-X websocket policy server. + +Opt-free counterpart to ``launch_serving.py``: builds :class:`WallXVanillaPolicy` +(no ``harrix.model_executors`` engine, no CUDA-graph / flashinfer) and serves it over the +same websocket transport. Intended for the open-source LIBERO/qwen2.5 release. + +Usage:: + + python -m harrix.serving.launch_vanilla_serving \ + --checkpoint-path /path/to/checkpoint \ + --port 32195 + +The client sends observations of the form:: + + {"state": {"eef_pos": [...], "eef_axisangle": [...], "gripper": [...]}, + "views": {"face_view": , "right_wrist_view": }, + "instruction": "..."} + +and receives ``{"action": }``. +""" + +import argparse +import logging +import os +import sys + +logger = logging.getLogger(__name__) + + +def parse_args(argv=None) -> argparse.Namespace: + p = argparse.ArgumentParser(description="Vanilla Wall-X websocket policy server") + p.add_argument("--checkpoint-path", required=True, help="Checkpoint dir or file.") + p.add_argument( + "--train-config-path", + default=None, + help="Training config used by the checkpoint (default: /config.yml).", + ) + p.add_argument("--norm-key", default="libero_all", help="Normalizer key.") + p.add_argument( + "--action-horizon", + type=int, + default=None, + help="Model action horizon (default: read from the train config).", + ) + p.add_argument( + "--cam-names", + nargs="+", + default=["face_view", "right_wrist_view"], + help="Camera view names, in model input order.", + ) + p.add_argument("--host", default="0.0.0.0", help="Bind host.") + p.add_argument("--port", type=int, default=32195, help="WebSocket port.") + p.add_argument( + "--image-passing-mode", + choices=["base64", "numpy"], + default="base64", + help="How the client passes images.", + ) + p.add_argument( + "--cuda-id", default="0", help="Sets CUDA_VISIBLE_DEVICES for the server." + ) + p.add_argument( + "--max-batch-size", + type=int, + default=0, + help="Enable dynamic batching (0 = disabled).", + ) + p.add_argument("--max-wait-time-ms", type=int, default=10) + p.add_argument("--max-queue-size", type=int, default=64) + p.add_argument("--timeout-ms", type=int, default=60000) + p.add_argument("--debug", action="store_true") + return p.parse_args(argv) + + +def main(argv=None) -> int: + args = parse_args(argv) + logging.basicConfig( + level=logging.DEBUG if args.debug else logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) + os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_id) + + from wall_x._vendor.harrix.e2e_infer.config import EvalConfig, autofill_from_checkpoint + from wall_x._vendor.harrix.serving.handlers.wall_x_vanilla_policy import WallXVanillaPolicy + from wall_x._vendor.harrix.serving.replica.server import WebsocketPolicyServer + + cfg = EvalConfig() + cfg.model.checkpoint_path = args.checkpoint_path + cfg.model.train_config_path = args.train_config_path + cfg.model.norm_key = args.norm_key + cfg.model.cam_names = list(args.cam_names) + cfg.model.action_mode = "flow" + if args.action_horizon is not None: + cfg.model.action_horizon = args.action_horizon + # Resolve train_config_path / action_horizon from the checkpoint when unset. + autofill_from_checkpoint(cfg) + + logger.info("Building vanilla Wall-X policy from %s", args.checkpoint_path) + policy = WallXVanillaPolicy(cfg, image_passing_mode=args.image_passing_mode) + + logger.info("Serving at ws://%s:%d", args.host, args.port) + logger.info( + "Batching: %s", + f"batch_size={args.max_batch_size}, wait={args.max_wait_time_ms}ms" + if args.max_batch_size + else "disabled", + ) + server = WebsocketPolicyServer( + policy=policy, + host=args.host, + port=args.port, + metadata=policy.metadata, + max_batch_size=args.max_batch_size or None, + max_wait_time_ms=args.max_wait_time_ms, + max_queue_size=args.max_queue_size, + timeout_ms=args.timeout_ms, + ) + try: + server.serve_forever() + except KeyboardInterrupt: + logger.info("Server stopped by user") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/wall_x/_vendor/harrix/serving/monitoring/__init__.py b/wall_x/_vendor/harrix/serving/monitoring/__init__.py new file mode 100644 index 00000000..1c18ab8f --- /dev/null +++ b/wall_x/_vendor/harrix/serving/monitoring/__init__.py @@ -0,0 +1 @@ +"""Serving metrics, health probes and deployment assets.""" diff --git a/wall_x/_vendor/harrix/serving/monitoring/metrics.py b/wall_x/_vendor/harrix/serving/monitoring/metrics.py new file mode 100644 index 00000000..d0564761 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/monitoring/metrics.py @@ -0,0 +1,469 @@ +"""Prometheus metrics exposition for the model service. + +Dependency-free: renders the existing router / worker stats snapshots into the +Prometheus text exposition format (version 0.0.4). The router and each worker +serve this at ``/metrics``; Prometheus distinguishes workers by the scrape +``instance`` label, and the router additionally labels per-worker series by +``worker_id`` / ``device``. + +Counters (``*_total``) come straight from the monotonic counters already in +the snapshots; gauges are instantaneous values. Request latency is exposed as +a Prometheus Histogram (``*_bucket`` / ``*_sum`` / ``*_count``) so the unified +platform can derive mean and any quantile, aggregated correctly across workers: + + mean : rate(..._sum[5m]) / rate(..._count[5m]) + p50 : histogram_quantile(0.50, sum by (le) (rate(..._bucket[5m]))) + p95 : histogram_quantile(0.95, sum by (le) (rate(..._bucket[5m]))) + p99 : histogram_quantile(0.99, sum by (le) (rate(..._bucket[5m]))) + +We render here (not prometheus_client objects) so there is no cross-process +registry to manage in multi-instance; the only instrumentation is a cheap +``LatencyHistogram.observe`` at the two points latency is already recorded. +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple + +CONTENT_TYPE = "text/plain; version=0.0.4; charset=utf-8" + +Sample = Tuple[Mapping[str, Any], Any] + +# Latency histogram bucket upper bounds in SECONDS (Prometheus convention). +# Dense around 50-500ms where typical inference latency lives, with a +# tail to 5s for queue-induced spikes. +DEFAULT_LATENCY_BUCKETS_S: Tuple[float, ...] = ( + 0.025, + 0.05, + 0.075, + 0.1, + 0.15, + 0.2, + 0.25, + 0.3, + 0.4, + 0.5, + 0.75, + 1.0, + 1.5, + 2.0, + 3.0, + 5.0, +) + + +@dataclass +class LatencyHistogram: + """Minimal cumulative-bucket latency histogram. + + Lock-free by design — not by tolerating races, but because there are + none: observe is only ever called from its process's single asyncio + event-loop thread (router: _finish_worker_request, under the router + lock; worker: _record_request_latency in the connection handler; the + model forward in asyncio.to_thread never touches the histogram), and + observe contains no await — so its _sum / _count / + _counts[idx] increments cannot be interleaved with another coroutine + and are effectively atomic. It also never raises (bad / NaN inputs are + dropped). NOTE: this is NOT thread-safe; if observe ever gets called + from a real worker thread (e.g. inside asyncio.to_thread), guard the + three increments with a threading.Lock. + """ + + bounds: Tuple[float, ...] = DEFAULT_LATENCY_BUCKETS_S + _counts: List[int] = field(default_factory=list) + _sum: float = 0.0 + _count: int = 0 + + def __post_init__(self) -> None: + # one extra slot for the implicit +Inf overflow bucket + if not self._counts: + self._counts = [0] * (len(self.bounds) + 1) + + def observe(self, seconds: Optional[float]) -> None: + if seconds is None: + return + try: + s = float(seconds) + except (TypeError, ValueError): + return + if s != s: # NaN + return + self._sum += s + self._count += 1 + idx = len(self.bounds) # default to +Inf overflow bucket + for i, b in enumerate(self.bounds): + if s <= b: + idx = i + break + self._counts[idx] += 1 + + def snapshot(self) -> Dict[str, Any]: + """Return cumulative (le, count) buckets + sum + count for rendering.""" + cum: List[Tuple[float, int]] = [] + running = 0 + for i, b in enumerate(self.bounds): + running += self._counts[i] + cum.append((b, running)) + running += self._counts[-1] + cum.append((float("inf"), running)) + return {"buckets": cum, "sum": self._sum, "count": self._count} + + +def _escape_label(value: Any) -> str: + s = str(value) + return s.replace("\\", "\\\\").replace("\n", "\\n").replace('"', '\\"') + + +def _format_labels(labels: Optional[Mapping[str, Any]]) -> str: + if not labels: + return "" + inner = ",".join(f'{k}="{_escape_label(v)}"' for k, v in labels.items()) + return "{" + inner + "}" + + +def _num(value: Any) -> Any: + """5.0 -> 5, 0.025 -> 0.025 (compact, valid Prometheus number).""" + f = float(value) + return int(f) if f.is_integer() else f + + +def _coerce(value: Any) -> Optional[float]: + if value is None: + return None + if isinstance(value, bool): + return 1.0 if value else 0.0 + try: + f = float(value) + except (TypeError, ValueError): + return None + if f != f or f in (float("inf"), float("-inf")): + return None + return f + + +class _Writer: + def __init__(self) -> None: + self._chunks: List[str] = [] + + def family( + self, name: str, mtype: str, help_text: str, samples: Iterable[Sample] + ) -> None: + lines = [] + for labels, value in samples: + v = _coerce(value) + if v is None: + continue + lines.append(f"{name}{_format_labels(labels)} {_num(v)}") + if not lines: + return + self._chunks.append(f"# HELP {name} {help_text}") + self._chunks.append(f"# TYPE {name} {mtype}") + self._chunks.extend(lines) + + def scalar( + self, + name: str, + mtype: str, + help_text: str, + value: Any, + labels: Optional[Mapping[str, Any]] = None, + ) -> None: + self.family(name, mtype, help_text, [(labels or {}, value)]) + + def histogram( + self, + name: str, + help_text: str, + series: Sequence[Tuple[Mapping[str, Any], Optional[Dict[str, Any]]]], + ) -> None: + """series: list of (labels, snapshot-dict-or-None). Emits one HELP/TYPE + then ``_bucket``/``_sum``/``_count`` lines for every series.""" + valid = [(lbl, s) for lbl, s in series if s and s.get("count") is not None] + if not valid: + return + self._chunks.append(f"# HELP {name} {help_text}") + self._chunks.append(f"# TYPE {name} histogram") + for labels, snap in valid: + base = dict(labels or {}) + for le, cum in snap["buckets"]: + le_str = "+Inf" if le == float("inf") else _num(le) + self._chunks.append( + f"{name}_bucket{_format_labels({**base, 'le': le_str})} {int(cum)}" + ) + self._chunks.append(f"{name}_sum{_format_labels(base)} {_num(snap['sum'])}") + self._chunks.append( + f"{name}_count{_format_labels(base)} {int(snap['count'])}" + ) + + def render(self) -> str: + return "\n".join(self._chunks) + "\n" + + +# --------------------------------------------------------------------------- # +# Router: fleet + per-worker routing view (single public scrape target). +# --------------------------------------------------------------------------- # +def render_router_metrics(snapshot: Dict[str, Any]) -> str: + w = _Writer() + router = snapshot.get("router", {}) + w.scalar( + "model_service_router_workers_total", + "gauge", + "Total worker instances behind the router.", + router.get("total_workers"), + ) + w.scalar( + "model_service_router_workers_healthy", + "gauge", + "Worker instances passing health checks.", + router.get("healthy_workers"), + ) + w.scalar( + "model_service_router_workers_available", + "gauge", + "Worker instances available for routing (healthy and not draining).", + router.get("available_workers"), + ) + w.scalar( + "model_service_router_active_connections", + "gauge", + "Active client (robot) connections currently held by the router " + "(number of robots being served).", + router.get("active_connections"), + ) + + ru = snapshot.get("rolling_update") or {} + w.scalar( + "model_service_rolling_update_running", + "gauge", + "1 while a rolling update is in progress, else 0.", + 1 if ru.get("running") else 0, + ) + wu = snapshot.get("weight_update") or {} + w.scalar( + "model_service_weight_update_running", + "gauge", + "1 while a rolling weight/actor update (rollover) is in progress, " "else 0.", + 1 if wu.get("running") else 0, + ) + + workers = snapshot.get("workers", []) or [] + + def labels(wk): + return {"worker_id": str(wk.get("worker_id")), "device": str(wk.get("device"))} + + def per(name: str, mtype: str, help_text: str, key: str, transform=lambda x: x): + w.family( + name, + mtype, + help_text, + [(labels(wk), transform(wk.get(key))) for wk in workers], + ) + + per( + "model_service_worker_active_connections", + "gauge", + "Active client connections proxied to this worker.", + "active_connections", + ) + per( + "model_service_worker_in_flight", + "gauge", + "Requests currently being processed by this worker.", + "in_flight", + ) + per( + "model_service_worker_requests_total", + "counter", + "Requests proxied to this worker.", + "total_requests", + ) + per( + "model_service_worker_connections_total", + "counter", + "Connections proxied to this worker.", + "total_connections", + ) + per( + "model_service_worker_errors_total", + "counter", + "Proxy/inference errors for this worker.", + "error_count", + ) + per( + "model_service_worker_healthy", + "gauge", + "1 if the worker is healthy, else 0.", + "healthy", + lambda x: 1 if x else 0, + ) + per( + "model_service_worker_draining", + "gauge", + "1 if the worker is draining (rolling update), else 0.", + "draining", + lambda x: 1 if x else 0, + ) + per( + "model_service_worker_available", + "gauge", + "1 if the worker is available for routing, else 0.", + "available_for_routing", + lambda x: 1 if x else 0, + ) + per( + "model_service_worker_reloading", + "gauge", + "1 while the worker is reloading weights in place (rollover), else 0.", + "reloading", + lambda x: 1 if x else 0, + ) + per( + "model_service_worker_restart_count", + "counter", + "Process restarts for this worker.", + "restart_count", + ) + per( + "model_service_worker_rolling_restart_count", + "counter", + "Rolling-update restarts for this worker.", + "rolling_restart_count", + ) + + w.family( + "model_service_worker_module_version", + "gauge", + "Per-worker loaded version of each model module " + "(base / encoder / actor); steps up worker-by-worker during rollover.", + [ + ({**labels(wk), "module": str(mod)}, ver) + for wk in workers + for mod, ver in (wk.get("module_versions") or {}).items() + ], + ) + # Per-worker request round-trip latency (router's view) as a histogram. + w.histogram( + "model_service_worker_request_duration_seconds", + "Request round-trip latency (seconds) per worker, as seen by the router.", + [(labels(wk), wk.get("latency_histogram")) for wk in workers], + ) + return w.render() + + +# --------------------------------------------------------------------------- # +# Worker: its own detailed view incl. scheduler / batching (per-port scrape). +# --------------------------------------------------------------------------- # +def render_worker_metrics(snapshot: Dict[str, Any]) -> str: + w = _Writer() + wk = snapshot.get("worker", {}) + rq = snapshot.get("requests", {}) + sch = snapshot.get("scheduler") + + w.scalar("model_service_up", "gauge", "1 if the worker server is up.", 1) + w.scalar( + "model_service_uptime_seconds", + "gauge", + "Seconds since the worker started.", + wk.get("uptime_s"), + ) + w.scalar( + "model_service_batching_enabled", + "gauge", + "1 if dynamic batching is enabled on this worker.", + 1 if wk.get("batching_enabled") else 0, + ) + + w.scalar( + "model_service_requests_total", + "counter", + "Inference requests handled by this worker.", + rq.get("total_requests"), + ) + w.scalar( + "model_service_connections_total", + "counter", + "Connections accepted by this worker.", + rq.get("total_connections"), + ) + w.scalar( + "model_service_errors_total", + "counter", + "Request errors on this worker.", + rq.get("error_count"), + ) + w.scalar( + "model_service_in_flight", + "gauge", + "Requests currently in flight on this worker.", + rq.get("in_flight"), + ) + w.scalar( + "model_service_active_connections", + "gauge", + "Active client connections on this worker.", + rq.get("active_connections"), + ) + + # Per-request handling latency as a histogram (mean/p50/p95/p99 derived in + # Grafana; see module docstring for the queries). + w.histogram( + "model_service_request_duration_seconds", + "Per-request handling latency (seconds) on this worker.", + [({}, rq.get("latency_histogram"))], + ) + + if sch: + w.scalar( + "model_service_scheduler_queue_size", + "gauge", + "Requests waiting in the batching queue.", + sch.get("queue_size"), + ) + w.scalar( + "model_service_scheduler_max_batch_size", + "gauge", + "Configured max batch size.", + sch.get("max_batch_size"), + ) + w.scalar( + "model_service_scheduler_batches_total", + "counter", + "Batches executed.", + sch.get("total_batches"), + ) + w.scalar( + "model_service_scheduler_batch_items_total", + "counter", + "Total requests folded into batches.", + sch.get("total_batch_items"), + ) + w.scalar( + "model_service_scheduler_last_batch_size", + "gauge", + "Size of the most recent batch.", + sch.get("last_batch_size"), + ) + w.scalar( + "model_service_scheduler_requests_completed_total", + "counter", + "Requests completed via the scheduler.", + sch.get("total_requests_completed"), + ) + w.scalar( + "model_service_scheduler_requests_failed_total", + "counter", + "Requests failed in the scheduler.", + sch.get("total_requests_failed"), + ) + w.scalar( + "model_service_scheduler_timeouts_total", + "counter", + "Requests that hit the scheduler timeout.", + sch.get("total_timeouts"), + ) + w.scalar( + "model_service_scheduler_queue_full_total", + "counter", + "Requests rejected because the queue was full.", + sch.get("queue_full_count"), + ) + return w.render() diff --git a/wall_x/_vendor/harrix/serving/policy/__init__.py b/wall_x/_vendor/harrix/serving/policy/__init__.py deleted file mode 100644 index e59a287b..00000000 --- a/wall_x/_vendor/harrix/serving/policy/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Wall-X policy for harrix websocket serving.""" - -from .wall_x_policy import WallXPolicy - -__all__ = ["WallXPolicy"] diff --git a/wall_x/_vendor/harrix/serving/policy/_smoothing.py b/wall_x/_vendor/harrix/serving/policy/_smoothing.py deleted file mode 100644 index 9b6c127d..00000000 --- a/wall_x/_vendor/harrix/serving/policy/_smoothing.py +++ /dev/null @@ -1,104 +0,0 @@ -"""Test-time Laplacian smoothing for serving model outputs. - -Mirrors the open-loop implementation in -`wallx_bus2604/wall-x/run_scripts/infer_openloop.py` -(_laplacian_smooth @ 203-209, smooth_action / smooth_gripper fields @ 186-187, -application @ 458-463). Invoked by WallXPolicy before downstream 6D->euler -conversion so that smoothed rotations propagate through serialization. -""" - -from __future__ import annotations - -from typing import List, Optional, Sequence - -import numpy as np -import torch - -from wall_x._vendor.harrix.serving._wallx_infer.base_dataclass import dof_dims - - -def _laplacian_smooth(a: np.ndarray, lam: float = 1.0, iters: int = 30) -> np.ndarray: - """Iterative Laplacian smoothing along axis 0; endpoints pinned.""" - a = a.copy() - orig = a.copy() - for _ in range(iters): - a[1:-1] = (orig[1:-1] + lam * (a[:-2] + a[2:])) / (1 + 2 * lam) - return a - - -def _gripper_column_indices( - predict_action_keys: Sequence[str], - action_padding_dof: Optional[int] = None, -) -> List[int]: - """Column indices in the flat (T, D) layout that correspond to gripper dofs.""" - cols: List[int] = [] - dof_start = 0 - for key in predict_action_keys: - if key == "action_padding": - dof_start += action_padding_dof or 0 - continue - short = key.replace("follow_", "").replace("master_", "") - width = dof_dims[short] - if "gripper" in short: - cols.extend(range(dof_start, dof_start + width)) - dof_start += width - return cols - - -_LAZY_ACTION_KEYS = ( - "action_left_ee_cartesian_pos", - "action_right_ee_cartesian_pos", - "action_left_ee_rotation", - "action_right_ee_rotation", - "action_left_ee_rotation_6D", - "action_right_ee_rotation_6D", -) - - -def apply_smoothing( - model_output: dict, - smooth_action: bool, - smooth_gripper: bool, - predict_action_keys: Sequence[str], - action_padding_dof: Optional[int] = None, - action_dim: Optional[int] = None, -) -> None: - """Smooth `model_output['predict_action']` in place and refresh per-arm keys.""" - if not smooth_action: - return - pa = model_output.get("predict_action") - if pa is None: - return - - was_tensor = isinstance(pa, torch.Tensor) - arr = pa.detach().cpu().numpy() if was_tensor else np.asarray(pa) - - orig_ndim = arr.ndim - if orig_ndim == 3: - if arr.shape[0] != 1: - return - arr = arr[0] - if arr.ndim != 2 or arr.shape[0] < 3: - return - if action_dim is not None and arr.shape[-1] != action_dim: - return - - orig = arr.copy() - smoothed = _laplacian_smooth(arr) - - if not smooth_gripper: - for c in _gripper_column_indices(predict_action_keys, action_padding_dof): - if c < smoothed.shape[-1]: - smoothed[:, c] = orig[:, c] - - out = smoothed[None] if orig_ndim == 3 else smoothed - if was_tensor: - out = torch.from_numpy(out).to(device=pa.device, dtype=pa.dtype) - model_output["predict_action"] = out - - rsd = model_output.get("robot_state_action_data") - if rsd is not None: - for k in _LAZY_ACTION_KEYS: - if k in rsd.data: - rsd.data[k] = None - rsd.save_action_data(out) diff --git a/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py b/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py deleted file mode 100644 index 155d6024..00000000 --- a/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py +++ /dev/null @@ -1,328 +0,0 @@ -import base64 -import logging -from typing import Dict, Any, List - -import cv2 -import numpy as np -import torch -from wall_x._vendor.harrix.serving.websocket_policy_server import BasePolicy -from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig -from wall_x._vendor.harrix.serving._wallx_infer.model_wrapper import WallxModelWrapper -from wall_x._vendor.harrix.serving._wallx_infer.robot import ( - DesktopRobotPreprocessor, - TurtleRobotPreprocessor, - EX001RobotPreprocessor, -) -from wall_x._vendor.harrix.serving.policy._smoothing import apply_smoothing -from wall_x.utils.timers import ScopeTimer - -logger = logging.getLogger(__name__) - -# Inference modes aligned with model_wrapper for vqa / batch extensions -INFER_MODE_FLOW = "flow" -INFER_MODE_AR = "ar" -INFER_MODE_FLOW_WITH_SUBTASK = "flow_with_subtask" -INFER_MODE_DLLM_FLOW = "dllm" -INFER_MODE_DLLM_DD = "discrete_diffusion" -ACTION_INFER_MODES = ( - INFER_MODE_FLOW, - INFER_MODE_AR, - INFER_MODE_FLOW_WITH_SUBTASK, - INFER_MODE_DLLM_FLOW, - INFER_MODE_DLLM_DD, -) - - -class WallXPolicy(BasePolicy): - """Policy wrapper for Wall-X model that implements the BasePolicy interface.""" - - def __init__( - self, - config: InferConfig, - image_passing_mode: str = "base64", - default_infer_mode: str = INFER_MODE_FLOW, - serialize_actions: bool = True, - ): - """Initialize the Wall-X policy. - - Args: - config: Inference configuration dataclass. - image_passing_mode: How images are passed from client ('base64' or 'numpy'). - default_infer_mode: Default inference mode ('flow', 'ar', 'flow_with_subtask', etc.). - serialize_actions: If True, actions are serialized via robot_preprocessor; - if False, raw model output is returned directly. - """ - self.config = config - self.model_wrapper = WallxModelWrapper(config) - self.robot_preprocessor = self._register_robot_preprocessor() - self.image_passing_mode = image_passing_mode - self.default_infer_mode = default_infer_mode - self.serialize_actions = serialize_actions - logger.info( - "Image passing mode: %s, robot_type: %s, default_infer_mode: %s, " - "serialize_actions: %s, smooth_action: %s, smooth_gripper: %s", - self.image_passing_mode, - config.robot_type, - self.default_infer_mode, - self.serialize_actions, - config.smooth_action, - config.smooth_gripper, - ) - - def _register_robot_preprocessor(self): - """Select preprocessor by config.robot_type; mirrors env._register_robot.""" - if self.config.robot_type == "desktop": - return DesktopRobotPreprocessor(self.config) - if self.config.robot_type == "turtle": - return TurtleRobotPreprocessor(self.config) - if self.config.robot_type == "ex001": - return EX001RobotPreprocessor(self.config) - raise ValueError(f"Invalid robot_type: {self.config.robot_type!r}") - - @staticmethod - def _decode_base64_image(image_b64: str) -> np.ndarray: - """Decode client base64 JPEG/PNG payloads to RGB images.""" - img_bytes = base64.b64decode(image_b64) - img_array = np.frombuffer(img_bytes, dtype=np.uint8) - decoded_img = cv2.imdecode(img_array, cv2.IMREAD_COLOR) - if decoded_img is None: - raise ValueError("Failed to decode base64 image payload") - return cv2.cvtColor(decoded_img, cv2.COLOR_BGR2RGB) - - def reset(self) -> None: - """Reset the policy state.""" - self.action_buffer = [] - self.buffer_index = 0 - logger.debug("Policy reset") - - def _get_dof_config(self) -> dict: - train_config = self.config.train_config or {} - return ( - train_config.get("dof_config") - or train_config.get("task", {}).get("dof_config", {}) - or {} - ) - - def _is_single_arm_right_only(self) -> bool: - """True for single-arm LIBERO-style configs (right arm only, no left arm).""" - train_config = self.config.train_config or {} - agent_pos = ( - train_config.get("agent_pos_config") - or train_config.get("task", {}).get("agent_pos_config") - or {} - ) - skip = {"action_padding"} - has_left = any(k.startswith("follow_left_") for k in agent_pos if k not in skip) - has_right = any( - k.startswith("follow_right_") for k in agent_pos if k not in skip - ) - return has_right and not has_left - - def _pack_action_chunk_response( - self, model_output: Dict[str, Any] - ) -> Dict[str, Any]: - """Return a msgpack-safe action chunk for websocket clients.""" - predict_action = model_output["predict_action"] - if isinstance(predict_action, torch.Tensor): - predict_action = predict_action.detach().cpu().numpy() - response: Dict[str, Any] = { - "predict_action": np.asarray(predict_action, dtype=np.float32), - } - if "subtask" in model_output: - response["subtask"] = model_output["subtask"] - return response - - def _get_predict_action_keys(self) -> List[str]: - """Resolve flat action key order for smoothing / serialization.""" - data_cfg = self.config.data_config - keys = None - if isinstance(data_cfg, dict): - keys = data_cfg.get("predict_action_keys") - else: - keys = getattr(data_cfg, "predict_action_keys", None) - if keys: - return list(keys) - return list(self._get_dof_config().keys()) - - def _apply_smoothing(self, model_output: Dict[str, Any]) -> None: - """Apply Laplacian smoothing on predict_action before 6D->euler conversion.""" - cfg = self.config - if not getattr(cfg, "smooth_action", False): - return - dof_config = self._get_dof_config() - apply_smoothing( - model_output, - smooth_action=True, - smooth_gripper=cfg.smooth_gripper, - predict_action_keys=self._get_predict_action_keys(), - action_padding_dof=dof_config.get("action_padding"), - action_dim=cfg.action_dim, - ) - - def _run_action_infer( - self, observation: Dict, instruction: str, mode: str - ) -> Dict[str, Any]: - """Run model_wrapper action inference by mode; returns model_output with robot_state_action_data. - - Supports flow | ar | flow_with_subtask. VQA can be added here later (different return format). - """ - if mode == INFER_MODE_FLOW: - return self.model_wrapper.infer_flow_action(observation, instruction) - if mode == INFER_MODE_AR: - return self.model_wrapper.infer_ar_action(observation, instruction) - if mode == INFER_MODE_FLOW_WITH_SUBTASK: - with ScopeTimer("infer_subtask"): - subtask = self.model_wrapper.infer_subtask(observation, instruction) - with ScopeTimer("infer_flow_action"): - model_output = self.model_wrapper.infer_flow_action( - observation, subtask - ) - model_output["subtask"] = subtask - return model_output - if mode == INFER_MODE_DLLM_FLOW: - return self.model_wrapper.infer_dllm_action( - observation, instruction, use_ar_action=False - ) - if mode == INFER_MODE_DLLM_DD: - return self.model_wrapper.infer_dllm_action( - observation, instruction, use_ar_action=True - ) - raise ValueError( - f"Unsupported infer_mode={mode!r}, expected one of {ACTION_INFER_MODES}" - ) - - def infer(self, obs: Dict) -> Dict: - """Infer action from observation. - - Args: - obs: Dictionary containing: - - 'state': Robot state - - 'views': Camera views (keyed by camera name) - - 'instruction': Task instruction - - Optional: 'infer_mode' - one of 'flow' | 'ar' | 'flow_with_subtask' - - Optional: 'robot_action_start_ratio' / 'robot_action_end_ratio' / - 'robot_action_interpolate_multiplier' - override config for action trim/interpolate - - Returns: - When serialize_actions=True (default): Serialized action dict - (e.g. follow1_pos/follow2_pos or follow1_joints/follow2_joints). - When serialize_actions=False: Raw model_output dict. - If infer_mode is flow_with_subtask, also includes 'subtask'. - """ - state = obs["state"] - views = obs["views"] - instruction = obs.get("instruction") or self.config.default_instruction or "" - if self.image_passing_mode == "base64": - for k, v in views.items(): - views[k] = np.expand_dims(self._decode_base64_image(v), axis=0) - - with ScopeTimer("get_observation"): - observation = self.robot_preprocessor.get_observation(state, views) - - mode = obs.get("infer_mode", self.default_infer_mode) - with ScopeTimer(f"infer_{mode}"): - model_output = self._run_action_infer(observation, instruction, mode) - - self._apply_smoothing(model_output) - if ( - "robot_state_action_data" in model_output - and "predict_action" in model_output - ): - model_output["robot_state_action_data"].save_action_data( - model_output["predict_action"] - ) - - # Single-arm LIBERO has no left-arm EE to serialize; dual-arm follow1/follow2 - # layout would fail in get_serialized_actions. Return raw action chunks instead. - if not self.serialize_actions or self._is_single_arm_right_only(): - return self._pack_action_chunk_response(model_output) - - return self.robot_preprocessor.get_serialized_actions( - model_output, robot_action_interpolate_multiplier=1 - ) # interpolation on robot websocket client - - # -- Batch inference ---------------------------------------------- - - def _preprocess_obs(self, obs: Dict): - """Preprocess a single observation dict into (observation, instruction). - - Handles both base64 and raw image modes. - """ - state = obs["state"] - views = obs["views"] - instruction = obs.get("instruction") or self.config.default_instruction or "" - - if self.image_passing_mode == "base64": - for k, v in views.items(): - views[k] = np.expand_dims(self._decode_base64_image(v), axis=0) - else: - # raw numpy: ensure (1, H, W, C) - for k, img in views.items(): - if isinstance(img, np.ndarray) and img.ndim == 3: - views[k] = np.expand_dims(img, axis=0) - - observation = self.robot_preprocessor.get_observation(state, views) - return observation, instruction - - def infer_batch( - self, obs_list: List[Dict[str, Any]], skip_serialize: bool = False - ) -> List[Dict[str, Any]]: - """Perform batch inference. - - Args: - obs_list: List of observations, each containing: - - "views": dict of camera images - - "state": robot state dict or array - - "instruction": text instruction - - Returns: - List of action dicts. - """ - batch_size = len(obs_list) - logger.info(f"WallXPolicy.infer_batch: processing {batch_size} observations") - - try: - observations = [] - instructions = [] - for obs in obs_list: - observation, instruction = self._preprocess_obs(obs) - observations.append(observation) - instructions.append(instruction) - - mode = obs_list[0].get("infer_mode", self.default_infer_mode) - - with torch.no_grad(): - if mode == INFER_MODE_FLOW: - model_outputs = self.model_wrapper.infer_flow_action_batch( - observations, instructions - ) - else: - # AR / flow_with_subtask: fall back to per-sample inference - model_outputs = [] - for obs_dict, instruction in zip(observations, instructions): - output = self._run_action_infer(obs_dict, instruction, mode) - model_outputs.append(output) - - if skip_serialize: - return [{}] * len(model_outputs) - - results = [] - for model_output in model_outputs: - self._apply_smoothing(model_output) - if self.serialize_actions and not self._is_single_arm_right_only(): - action = self.robot_preprocessor.get_serialized_actions( - model_output, robot_action_interpolate_multiplier=1 - ) - results.append(action) - else: - results.append(self._pack_action_chunk_response(model_output)) - - return results - - except Exception as e: - logger.error(f"Batch inference failed: {e}", exc_info=True) - return [{"action": {}, "error": str(e)} for _ in obs_list] - - @property - def metadata(self) -> Dict[str, Any]: - return {"batch_enabled": True, "model": "wall-x"} diff --git a/wall_x/_vendor/harrix/serving/replica/__init__.py b/wall_x/_vendor/harrix/serving/replica/__init__.py new file mode 100644 index 00000000..e7da5f45 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/replica/__init__.py @@ -0,0 +1 @@ +"""Single serving-replica process, transport, queue and batching.""" diff --git a/wall_x/_vendor/harrix/serving/replica/scheduler.py b/wall_x/_vendor/harrix/serving/replica/scheduler.py new file mode 100644 index 00000000..487123fe --- /dev/null +++ b/wall_x/_vendor/harrix/serving/replica/scheduler.py @@ -0,0 +1,284 @@ +"""Request queue and dynamic batching for one serving replica.""" + +from __future__ import annotations + +import asyncio +import logging +import time +import uuid +from collections import deque +from dataclasses import dataclass +from typing import Any, Callable, Dict, List + +logger = logging.getLogger(__name__) + + +@dataclass +class Request: + request_id: str + obs: Any + future: asyncio.Future + timestamp: float + packed: bool = False + policy_mode_override: str | None = None + + +class RequestScheduler: + def __init__( + self, + policy, + max_batch_size: int = 8, + max_wait_time_ms: float = 100, + max_queue_size: int = 128, + timeout_ms: float = 5000, + policy_lock: Any = None, + request_transform: Callable[[Any, str | None], Any] | None = None, + ): + self.policy = policy + self.max_batch_size = max_batch_size + self.max_wait_time = max_wait_time_ms / 1000.0 + self.max_queue_size = max_queue_size + self.timeout = timeout_ms / 1000.0 + self.policy_lock = policy_lock + self.request_transform = request_transform + + self.queue = deque() + self.queue_lock = asyncio.Lock() + self.queue_not_empty = asyncio.Condition(self.queue_lock) + + self.running = False + self.batch_task = None + self.started_at = None + self.total_requests_enqueued = 0 + self.total_requests_completed = 0 + self.total_requests_failed = 0 + self.total_timeouts = 0 + self.queue_full_count = 0 + self.total_batches = 0 + self.total_batch_items = 0 + self.total_batch_latency_ms = 0.0 + self.last_batch_size = 0 + self.last_batch_latency_ms = None + self.last_batch_started_at = None + self.last_batch_finished_at = None + + async def start(self): + self.running = True + self.started_at = time.time() + self.batch_task = asyncio.create_task(self._batch_loop()) + logger.info( + f"RequestScheduler started: max_batch={self.max_batch_size}, max_wait={self.max_wait_time * 1000}ms" + ) + + async def stop(self): + # Signal batch_loop to stop under the lock so _collect_batch + # wakes up and sees running=False atomically. + async with self.queue_lock: + self.running = False + self.queue_not_empty.notify_all() + + # Wait for batch_loop to fully exit before draining. + if self.batch_task: + await self.batch_task + + # Now drain any remaining requests (no race — batch_loop is done). + async with self.queue_lock: + remaining = list(self.queue) + self.queue.clear() + if remaining: + logger.info( + f"Graceful shutdown: processing {len(remaining)} remaining request(s)" + ) + for i in range(0, len(remaining), self.max_batch_size): + batch = remaining[i : i + self.max_batch_size] + await self._process_batch(batch) + logger.info("RequestScheduler stopped") + + async def add_request( + self, + obs: Dict[str, Any] | bytes, + packed: bool = False, + policy_mode_override: str | None = None, + ) -> Dict[str, Any]: + request_id = str(uuid.uuid4()) + future = asyncio.Future() + request = Request( + request_id=request_id, + obs=obs, + future=future, + timestamp=time.monotonic(), + packed=packed, + policy_mode_override=policy_mode_override, + ) + + async with self.queue_lock: + if len(self.queue) >= self.max_queue_size: + self.queue_full_count += 1 + raise RuntimeError( + f"Queue full: {len(self.queue)}/{self.max_queue_size}" + ) + self.queue.append(request) + self.total_requests_enqueued += 1 + self.queue_not_empty.notify() + + try: + result = await asyncio.wait_for(future, timeout=self.timeout) + return result + except asyncio.TimeoutError: + self.total_timeouts += 1 + logger.error(f"Request {request_id} timeout after {self.timeout}s") + raise + + def stats_snapshot(self) -> Dict[str, Any]: + queue_size = len(self.queue) + average_batch_latency_ms = None + if self.total_batches > 0: + average_batch_latency_ms = self.total_batch_latency_ms / self.total_batches + return { + "running": self.running, + "started_at": self.started_at, + "queue_size": queue_size, + "max_queue_size": self.max_queue_size, + "max_batch_size": self.max_batch_size, + "max_wait_time_ms": self.max_wait_time * 1000, + "timeout_ms": self.timeout * 1000, + "total_requests_enqueued": self.total_requests_enqueued, + "total_requests_completed": self.total_requests_completed, + "total_requests_failed": self.total_requests_failed, + "total_timeouts": self.total_timeouts, + "queue_full_count": self.queue_full_count, + "total_batches": self.total_batches, + "total_batch_items": self.total_batch_items, + "last_batch_size": self.last_batch_size, + "last_batch_latency_ms": self.last_batch_latency_ms, + "average_batch_latency_ms": average_batch_latency_ms, + "last_batch_started_at": self.last_batch_started_at, + "last_batch_finished_at": self.last_batch_finished_at, + } + + async def _batch_loop(self): + while self.running: + try: + batch = await self._collect_batch() + if batch: + await self._process_batch(batch) + except Exception as e: + logger.error(f"Unexpected error in batch loop: {e}", exc_info=True) + + async def _collect_batch(self) -> List[Request]: + async with self.queue_not_empty: + while self.running and len(self.queue) == 0: + await self.queue_not_empty.wait() + + if not self.running: + return [] + + batch = [] + deadline = time.monotonic() + self.max_wait_time + + while len(batch) < self.max_batch_size: + if len(self.queue) > 0: + batch.append(self.queue.popleft()) + + if len(batch) >= self.max_batch_size: + break + + if len(self.queue) == 0: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + await asyncio.wait_for( + self.queue_not_empty.wait(), timeout=remaining + ) + except asyncio.TimeoutError: + break + + return batch + + async def _process_batch(self, batch: List[Request]): + if not batch: + return + + start_time = time.monotonic() + self.last_batch_started_at = time.time() + logger.debug("Processing batch of %d requests", len(batch)) + + try: + results = await asyncio.to_thread(self._infer_batch_payloads, batch) + if len(results) != len(batch): + raise RuntimeError( + "Policy infer_batch returned " + f"{len(results)} result(s) for {len(batch)} request(s)" + ) + + # Split queue-wait (time spent in the batch queue) from model- + # forward time (this batch's processing). The old single infer_ms + # conflated both, hiding queue buildup during rollover. + batch_ms = (time.monotonic() - start_time) * 1000.0 + completed_count = 0 + for req, result in zip(batch, results): + if isinstance(result, dict): + result["__sched_timing__"] = { + "queue_ms": round((start_time - req.timestamp) * 1000.0, 2), + "batch_ms": round(batch_ms, 2), + "batch_size": len(batch), + } + if not req.future.done(): + req.future.set_result(result) + completed_count += 1 + + infer_time = time.monotonic() - start_time + infer_time_ms = infer_time * 1000 + self.total_batches += 1 + self.total_batch_items += len(batch) + self.total_requests_completed += completed_count + self.total_batch_latency_ms += infer_time_ms + self.last_batch_size = len(batch) + self.last_batch_latency_ms = infer_time_ms + self.last_batch_finished_at = time.time() + logger.debug( + "Batch processed in %.1fms, throughput: %.1f req/s", + infer_time_ms, + len(batch) / infer_time, + ) + + except Exception as e: + infer_time_ms = (time.monotonic() - start_time) * 1000 + failed_count = 0 + self.last_batch_size = len(batch) + self.last_batch_latency_ms = infer_time_ms + self.last_batch_finished_at = time.time() + logger.error(f"Batch processing failed: {e}", exc_info=True) + for req in batch: + if not req.future.done(): + req.future.set_exception(e) + failed_count += 1 + self.total_requests_failed += failed_count + + def _infer_batch_payloads(self, batch: List[Request]) -> List[Dict[str, Any]]: + obs_list = [] + has_packed = any(req.packed for req in batch) + if has_packed: + try: + import msgpack + import msgpack_numpy as m + + m.patch() + except ImportError as exc: + raise RuntimeError( + "msgpack and msgpack-numpy are required for packed requests" + ) from exc + + for req in batch: + if req.packed: + obs = msgpack.unpackb(req.obs) + else: + obs = req.obs + if self.request_transform is not None: + obs = self.request_transform(obs, req.policy_mode_override) + obs_list.append(obs) + if self.policy_lock is None: + return self.policy.infer_batch(obs_list) + with self.policy_lock: + return self.policy.infer_batch(obs_list) diff --git a/wall_x/_vendor/harrix/serving/replica/server.py b/wall_x/_vendor/harrix/serving/replica/server.py new file mode 100644 index 00000000..05b6d893 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/replica/server.py @@ -0,0 +1,385 @@ +"""WebSocket and health endpoints for one serving replica.""" + +from __future__ import annotations + +import asyncio +import http +import json +import logging +import os +import time +import threading +import uuid +from contextlib import AsyncExitStack +from functools import partial +from typing import Any, Callable, Dict, List + +try: + import msgpack + import msgpack_numpy as m + + m.patch() +except ImportError: + logging.warning( + "msgpack-numpy not installed. Install with: pip install msgpack-numpy" + ) + msgpack = None + +import websockets.asyncio.server as _server +import websockets.frames + +from wall_x._vendor.harrix.serving.monitoring.metrics import LatencyHistogram +from wall_x._vendor.harrix.serving.handlers.base import BasePolicy + +logger = logging.getLogger(__name__) + + +class WebsocketPolicyServer: + """Serves a policy using the websocket protocol. + + Implements a websocket server that: + 1. Sends policy metadata on connection + 2. Receives observations + 3. Returns predicted actions (single or batched) + 4. Tracks timing information + + When batching parameters are provided, dynamic batching is enabled: + requests from multiple clients are queued and processed in batches. + """ + + def __init__( + self, + policy: BasePolicy, + host: str = "0.0.0.0", + port: int = 8000, + metadata: Dict | None = None, + # Dynamic batching parameters (None = disabled) + max_batch_size: int | None = None, + max_wait_time_ms: float | None = None, + max_queue_size: int = 100, + timeout_ms: float = 30000, + policy_mode_override: str | None = None, + additional_policy_mode_ports: Dict[int, str] | None = None, + request_transform: Callable[[Any, str | None], Any] | None = None, + listener_metadata: Dict[str, Dict[str, Any]] | None = None, + lifecycle_hooks: List[Any] | None = None, + ) -> None: + self._policy = policy + self._host = host + self._port = port + self._metadata = metadata or {} + self._started_at = time.time() + self._active_connections = 0 + self._active_connection_ids = set() + self._in_flight = 0 + self._total_connections = 0 + self._total_requests = 0 + self._error_count = 0 + self._last_error = None + self._last_request_started_at = None + self._last_request_finished_at = None + self._last_request_latency_ms = None + self._ewma_latency_ms = None + self._latency_hist = LatencyHistogram() + self._listener_modes = {port: policy_mode_override} + self._listener_modes.update(additional_policy_mode_ports or {}) + self._request_transform = request_transform + self._listener_metadata = listener_metadata or {} + self._policy_lock = threading.RLock() + self._lifecycle_hooks = list(lifecycle_hooks or []) + + # Dynamic batching + self._scheduler = None + if max_batch_size is not None: + from .scheduler import RequestScheduler + + self._scheduler = RequestScheduler( + policy=policy, + max_batch_size=max_batch_size, + max_wait_time_ms=( + max_wait_time_ms if max_wait_time_ms is not None else 0 + ), + max_queue_size=max_queue_size, + timeout_ms=timeout_ms, + policy_lock=self._policy_lock, + request_transform=request_transform, + ) + + logging.getLogger("websockets.server").setLevel(logging.WARNING) + + @property + def batching_enabled(self) -> bool: + return self._scheduler is not None + + @property + def policy_lock(self): + return self._policy_lock + + @property + def in_flight(self) -> int: + return self._in_flight + + def add_lifecycle_hook(self, hook) -> None: + self._lifecycle_hooks.append(hook) + + def serve_forever(self) -> None: + asyncio.run(self.run()) + + async def run(self): + async with AsyncExitStack() as stack: + if self._scheduler is not None: + await self._scheduler.start() + stack.push_async_callback(self._scheduler.stop) + for hook in self._lifecycle_hooks: + hook.start() + stack.callback(hook.stop) + + servers = [] + for port, override in self._listener_modes.items(): + server = await stack.enter_async_context( + _server.serve( + partial(self._handler, policy_mode_override=override), + self._host, + port, + compression=None, + max_size=None, + ping_interval=None, + ping_timeout=None, + process_request=self._process_request, + ) + ) + servers.append(server) + mode_str = "batched" if self.batching_enabled else "single" + for port, override in self._listener_modes.items(): + logger.info( + "Server started on %s:%s (mode=%s policy_mode_override=%s)", + self._host, + port, + mode_str, + override or "request", + ) + await asyncio.gather(*(server.serve_forever() for server in servers)) + + async def _handler( + self, + websocket: _server.ServerConnection, + *, + policy_mode_override: str | None = None, + ): + connection_id = uuid.uuid4().hex[:12] + logger.info( + "Worker connection opened connection_id=%s client=%s", + connection_id, + websocket.remote_address, + ) + self._active_connections += 1 + self._active_connection_ids.add(connection_id) + self._total_connections += 1 + + try: + if msgpack is None: + await websocket.close( + code=websockets.frames.CloseCode.INTERNAL_ERROR, + reason="msgpack-numpy not installed on server", + ) + return + + # Send metadata, refreshing the live per-module weight versions so + # the router tracks this worker's current version after in-place + # swaps (self._metadata is captured once at startup). + meta = dict(self._metadata) + meta["policy_mode_override"] = policy_mode_override + if policy_mode_override is not None: + meta.update( + self._listener_metadata.get(policy_mode_override, {}) + ) + ws_fn = getattr(self._policy, "weights_status", None) + if callable(ws_fn): + versions = ws_fn().get("versions") + if versions: + meta["module_versions"] = versions + await websocket.send(msgpack.packb(meta)) + + prev_total_time = None + request_index = 0 + while True: + try: + packed_obs = await websocket.recv() + except websockets.ConnectionClosed: + logger.info( + "Worker connection closed connection_id=%s client=%s", + connection_id, + websocket.remote_address, + ) + break + + request_index += 1 + self._in_flight += 1 + self._total_requests += 1 + recv_ts = time.time() + self._last_request_started_at = recv_ts + start_time = time.monotonic() + try: + infer_time = time.monotonic() + + if self._scheduler is not None: + # Dynamic batching path. Queue the packed payload + # immediately; unpacking runs inside the scheduler's + # batch thread so the worker event loop can enqueue + # concurrent requests without serial msgpack work. + action = await self._scheduler.add_request( + packed_obs, + packed=True, + policy_mode_override=policy_mode_override, + ) + else: + obs = msgpack.unpackb(packed_obs) + if self._request_transform is not None: + obs = self._request_transform( + obs, + policy_mode_override, + ) + # Single inference path: run synchronously on the + # event-loop thread. torch.compile/CUDA-graph state + # set up by warmup is per-thread TLS — dispatching to + # `asyncio.to_thread` lands on a fresh worker thread + # whose TLS is empty, which makes the compiled + # sample_actions trip an internal assertion (see + # cudagraph_trees.get_obj). For a single-robot WS + # session the ~100ms block per request is fine. + with self._policy_lock: + action = self._policy.infer(obs) + + infer_time = time.monotonic() - infer_time + + sched_timing = ( + action.pop("__sched_timing__", None) + if isinstance(action, dict) + else None + ) + # t_recv/t_send: epoch seconds on the worker clock. Cross- + # machine differences vs client timestamps include clock + # offset — only durations (t_send - t_recv) and intra- + # episode trends are trustworthy (robot_monitor.md §3). + action["server_timing"] = { + "infer_ms": infer_time * 1000, + "t_recv": recv_ts, + "t_send": time.time(), + } + if sched_timing: + # queue_ms = batch-queue wait; batch_ms = model forward + action["server_timing"].update(sched_timing) + if prev_total_time is not None: + action["server_timing"]["prev_total_ms"] = ( + prev_total_time * 1000 + ) + + await websocket.send(msgpack.packb(action)) + total_time = time.monotonic() - start_time + self._record_request_latency(total_time * 1000) + prev_total_time = total_time + + except websockets.ConnectionClosed: + logger.info( + "Worker connection closed connection_id=%s client=%s", + connection_id, + websocket.remote_address, + ) + break + except Exception as e: + self._record_error(e) + logger.error( + "Error handling request connection_id=%s request_index=%s: %s", + connection_id, + request_index, + e, + exc_info=True, + ) + await websocket.close( + code=websockets.frames.CloseCode.INTERNAL_ERROR, + reason="Internal server error.", + ) + raise + finally: + self._in_flight = max(0, self._in_flight - 1) + finally: + self._active_connections = max(0, self._active_connections - 1) + self._active_connection_ids.discard(connection_id) + + def _record_request_latency(self, latency_ms: float) -> None: + self._last_request_latency_ms = latency_ms + self._last_request_finished_at = time.time() + self._latency_hist.observe(latency_ms / 1000.0) + alpha = 0.2 + if self._ewma_latency_ms is None: + self._ewma_latency_ms = latency_ms + else: + self._ewma_latency_ms = ( + alpha * latency_ms + (1.0 - alpha) * self._ewma_latency_ms + ) + + def _record_error(self, error: Exception) -> None: + self._error_count += 1 + self._last_error = f"{type(error).__name__}: {error}" + + def stats_snapshot(self) -> Dict[str, Any]: + scheduler_stats = None + if self._scheduler is not None: + scheduler_stats = self._scheduler.stats_snapshot() + return { + "worker": { + "pid": os.getpid(), + "host": self._host, + "port": self._port, + "listeners": dict(self._listener_modes), + "metadata": self._metadata, + "started_at": self._started_at, + "uptime_s": time.time() - self._started_at, + "batching_enabled": self.batching_enabled, + }, + "requests": { + "active_connections": self._active_connections, + "active_connection_ids": sorted(self._active_connection_ids), + "in_flight": self._in_flight, + "total_connections": self._total_connections, + "total_requests": self._total_requests, + "error_count": self._error_count, + "last_error": self._last_error, + "last_request_started_at": self._last_request_started_at, + "last_request_finished_at": self._last_request_finished_at, + "last_request_latency_ms": self._last_request_latency_ms, + "ewma_latency_ms": self._ewma_latency_ms, + "latency_histogram": self._latency_hist.snapshot(), + }, + "scheduler": scheduler_stats, + } + + def _process_request( + self, connection: _server.ServerConnection, request: _server.Request + ) -> _server.Response | None: + if request.path == "/healthz": + return connection.respond(http.HTTPStatus.OK, "OK\n") + if request.path == "/v2/health/ready": + return connection.respond(http.HTTPStatus.OK, "OK\n") + if request.path == "/metrics": + from wall_x._vendor.harrix.serving.monitoring.metrics import ( + render_worker_metrics, + ) + + return connection.respond( + http.HTTPStatus.OK, render_worker_metrics(self.stats_snapshot()) + ) + if request.path in ("/worker/status", "/stats"): + return connection.respond( + http.HTTPStatus.OK, + json.dumps(self.stats_snapshot(), default=str, indent=2, sort_keys=True) + + "\n", + ) + if request.path == "/weights/status": + ws = getattr(self._policy, "weights_status", None) + payload = ws() if callable(ws) else {} + return connection.respond( + http.HTTPStatus.OK, + json.dumps(payload, default=str) + "\n", + ) + return None diff --git a/wall_x/_vendor/harrix/serving/scheduler.py b/wall_x/_vendor/harrix/serving/scheduler.py deleted file mode 100644 index 8d36fdb3..00000000 --- a/wall_x/_vendor/harrix/serving/scheduler.py +++ /dev/null @@ -1,160 +0,0 @@ -import asyncio -import logging -import time -import uuid -from collections import deque -from dataclasses import dataclass -from typing import Any, Dict, List - -logger = logging.getLogger(__name__) - - -@dataclass -class Request: - request_id: str - obs: Dict[str, Any] - future: asyncio.Future - timestamp: float - - -class RequestScheduler: - def __init__( - self, - policy, - max_batch_size: int = 8, - max_wait_time_ms: float = 100, - max_queue_size: int = 128, - timeout_ms: float = 5000, - ): - self.policy = policy - self.max_batch_size = max_batch_size - self.max_wait_time = max_wait_time_ms / 1000.0 - self.max_queue_size = max_queue_size - self.timeout = timeout_ms / 1000.0 - - self.queue = deque() - self.queue_lock = asyncio.Lock() - self.queue_not_empty = asyncio.Condition(self.queue_lock) - - self.running = False - self.batch_task = None - - async def start(self): - self.running = True - self.batch_task = asyncio.create_task(self._batch_loop()) - logger.info( - f"RequestScheduler started: max_batch={self.max_batch_size}, max_wait={self.max_wait_time*1000}ms" - ) - - async def stop(self): - self.running = False - # Process all remaining requests before shutting down - async with self.queue_lock: - remaining = len(self.queue) - if remaining: - logger.info( - f"Graceful shutdown: processing {remaining} remaining request(s)" - ) - while True: - async with self.queue_lock: - if len(self.queue) == 0: - break - batch = [] - while self.queue and len(batch) < self.max_batch_size: - batch.append(self.queue.popleft()) - if batch: - await self._process_batch(batch) - # Now stop the batch loop - async with self.queue_lock: - self.queue_not_empty.notify_all() - if self.batch_task: - await self.batch_task - logger.info("RequestScheduler stopped") - - async def add_request(self, obs: Dict[str, Any]) -> Dict[str, Any]: - request_id = str(uuid.uuid4()) - future = asyncio.Future() - request = Request( - request_id=request_id, obs=obs, future=future, timestamp=time.monotonic() - ) - - async with self.queue_lock: - if len(self.queue) >= self.max_queue_size: - raise RuntimeError( - f"Queue full: {len(self.queue)}/{self.max_queue_size}" - ) - self.queue.append(request) - self.queue_not_empty.notify() - - try: - result = await asyncio.wait_for(future, timeout=self.timeout) - return result - except asyncio.TimeoutError: - logger.error(f"Request {request_id} timeout after {self.timeout}s") - raise - - async def _batch_loop(self): - while self.running: - try: - batch = await self._collect_batch() - if batch: - await self._process_batch(batch) - except Exception as e: - logger.error(f"Unexpected error in batch loop: {e}", exc_info=True) - - async def _collect_batch(self) -> List[Request]: - async with self.queue_not_empty: - while self.running and len(self.queue) == 0: - await self.queue_not_empty.wait() - - if not self.running: - return [] - - batch = [] - deadline = time.monotonic() + self.max_wait_time - - while len(batch) < self.max_batch_size: - if len(self.queue) > 0: - batch.append(self.queue.popleft()) - - if len(batch) >= self.max_batch_size: - break - - if len(self.queue) == 0: - remaining = deadline - time.monotonic() - if remaining <= 0: - break - try: - await asyncio.wait_for( - self.queue_not_empty.wait(), timeout=remaining - ) - except asyncio.TimeoutError: - break - - return batch - - async def _process_batch(self, batch: List[Request]): - if not batch: - return - - start_time = time.monotonic() - logger.info(f"Processing batch of {len(batch)} requests") - - try: - obs_list = [req.obs for req in batch] - results = await asyncio.to_thread(self.policy.infer_batch, obs_list) - - for req, result in zip(batch, results): - if not req.future.done(): - req.future.set_result(result) - - infer_time = time.monotonic() - start_time - logger.info( - f"Batch processed in {infer_time*1000:.1f}ms, throughput: {len(batch)/infer_time:.1f} req/s" - ) - - except Exception as e: - logger.error(f"Batch processing failed: {e}", exc_info=True) - for req in batch: - if not req.future.done(): - req.future.set_exception(e) diff --git a/wall_x/_vendor/harrix/serving/websocket_policy_server.py b/wall_x/_vendor/harrix/serving/websocket_policy_server.py deleted file mode 100644 index 54ae810e..00000000 --- a/wall_x/_vendor/harrix/serving/websocket_policy_server.py +++ /dev/null @@ -1,184 +0,0 @@ -import asyncio -import http -import logging -import time -import traceback -from typing import Dict, Any, List - -try: - import msgpack - import msgpack_numpy as m - - m.patch() -except ImportError: - logging.warning( - "msgpack-numpy not installed. Install with: pip install msgpack-numpy" - ) - msgpack = None - -import websockets.asyncio.server as _server -import websockets.frames - -logger = logging.getLogger(__name__) - - -class BasePolicy: - """Base class for policies that can be served.""" - - def infer(self, obs: Dict) -> Dict: - """Infer actions from observations.""" - raise NotImplementedError - - def infer_batch(self, obs_list: List[Dict]) -> List[Dict]: - """Batch inference. Default implementation calls infer() per sample.""" - return [self.infer(obs) for obs in obs_list] - - def reset(self) -> None: - """Reset the policy to its initial state.""" - pass - - @property - def metadata(self) -> Dict[str, Any]: - """Return metadata about the policy.""" - return {} - - -class WebsocketPolicyServer: - """Serves a policy using the websocket protocol. - - Implements a websocket server that: - 1. Sends policy metadata on connection - 2. Receives observations - 3. Returns predicted actions (single or batched) - 4. Tracks timing information - - When batching parameters are provided, dynamic batching is enabled: - requests from multiple clients are queued and processed in batches. - """ - - def __init__( - self, - policy: BasePolicy, - host: str = "0.0.0.0", - port: int = 8000, - metadata: Dict | None = None, - # Dynamic batching parameters (None = disabled) - max_batch_size: int | None = None, - max_wait_time_ms: float | None = None, - max_queue_size: int = 100, - timeout_ms: float = 30000, - ) -> None: - self._policy = policy - self._host = host - self._port = port - self._metadata = metadata or {} - - # Dynamic batching - self._scheduler = None - if max_batch_size is not None: - from .scheduler import RequestScheduler - - self._scheduler = RequestScheduler( - policy=policy, - max_batch_size=max_batch_size, - max_wait_time_ms=( - max_wait_time_ms if max_wait_time_ms is not None else 0 - ), - max_queue_size=max_queue_size, - timeout_ms=timeout_ms, - ) - - logging.getLogger("websockets.server").setLevel(logging.INFO) - - @property - def batching_enabled(self) -> bool: - return self._scheduler is not None - - def serve_forever(self) -> None: - asyncio.run(self.run()) - - async def run(self): - # Start the scheduler if batching is enabled - if self._scheduler is not None: - await self._scheduler.start() - - try: - async with _server.serve( - self._handler, - self._host, - self._port, - compression=None, - max_size=None, - ping_interval=None, - ping_timeout=None, - process_request=_health_check, - ) as server: - mode_str = "batched" if self.batching_enabled else "single" - logger.info( - f"Server started on {self._host}:{self._port} (mode={mode_str})" - ) - await server.serve_forever() - finally: - if self._scheduler is not None: - await self._scheduler.stop() - - async def _handler(self, websocket: _server.ServerConnection): - logger.info(f"Connection from {websocket.remote_address} opened") - - if msgpack is None: - await websocket.close( - code=websockets.frames.CloseCode.INTERNAL_ERROR, - reason="msgpack-numpy not installed on server", - ) - return - - # Send metadata to client - await websocket.send(msgpack.packb(self._metadata)) - - prev_total_time = None - while True: - try: - start_time = time.monotonic() - obs = msgpack.unpackb(await websocket.recv()) - - infer_time = time.monotonic() - - if self._scheduler is not None: - # Dynamic batching path - action = await self._scheduler.add_request(obs) - else: - # Single inference path - action = await asyncio.to_thread(self._policy.infer, obs) - - infer_time = time.monotonic() - infer_time - - action["server_timing"] = { - "infer_ms": infer_time * 1000, - } - if prev_total_time is not None: - action["server_timing"]["prev_total_ms"] = prev_total_time * 1000 - - await websocket.send(msgpack.packb(action)) - prev_total_time = time.monotonic() - start_time - - except websockets.ConnectionClosed: - logger.info(f"Connection from {websocket.remote_address} closed") - break - except Exception as e: - logger.error(f"Error handling request: {e}") - await websocket.send(traceback.format_exc()) - await websocket.close( - code=websockets.frames.CloseCode.INTERNAL_ERROR, - reason="Internal server error. Traceback included in previous frame.", - ) - raise - - -def _health_check( - connection: _server.ServerConnection, request: _server.Request -) -> _server.Response | None: - if request.path == "/healthz": - return connection.respond(http.HTTPStatus.OK, "OK\n") - if request.path == "/v2/health/ready": - return connection.respond(http.HTTPStatus.OK, "OK\n") - return None diff --git a/wall_x/_vendor/harrix/utils/ckpt_load.py b/wall_x/_vendor/harrix/utils/ckpt_load.py index 4e003500..53808524 100644 --- a/wall_x/_vendor/harrix/utils/ckpt_load.py +++ b/wall_x/_vendor/harrix/utils/ckpt_load.py @@ -7,82 +7,11 @@ from __future__ import annotations import os -from typing import Callable, Optional import torch from safetensors.torch import load_file -def _noop_log(_msg: str, **_kw) -> None: - pass - - -def _align_checkpoint_tensor( - param: torch.Tensor, - target: torch.Tensor, - name: str, - log_fn: Callable, -) -> torch.Tensor | None: - """Crop or pad a checkpoint tensor to match the current model parameter.""" - if param.shape == target.shape: - return param - if param.ndim != target.ndim: - log_fn( - f"Skipping '{name}': ndim mismatch " - f"checkpoint={param.ndim} model={target.ndim}" - ) - return None - - overlap = tuple( - slice(0, min(src, dst)) for src, dst in zip(param.shape, target.shape) - ) - - if all(src >= dst for src, dst in zip(param.shape, target.shape)): - aligned = param[overlap].contiguous() - log_fn( - f"Cropped '{name}': checkpoint {tuple(param.shape)} " - f"-> model {tuple(aligned.shape)}" - ) - return aligned - - if all(src <= dst for src, dst in zip(param.shape, target.shape)): - aligned = target.detach().clone() - aligned[overlap] = param[overlap] - log_fn( - f"Padded '{name}': checkpoint {tuple(param.shape)} " - f"-> model {tuple(aligned.shape)} (tail keeps model init)" - ) - return aligned - - aligned = target.detach().clone() - aligned[overlap] = param[overlap] - log_fn( - f"Partially aligned '{name}': checkpoint {tuple(param.shape)} " - f"-> model {tuple(aligned.shape)} (non-overlap keeps model init)" - ) - return aligned - - -def reshape_compatible_state_dict( - state_dict: dict, model_sd: dict, log_fn: Optional[Callable] = None -) -> dict: - """Align checkpoint tensors to the target model shapes via crop / pad.""" - log_fn = log_fn or _noop_log - out = {} - for name, param in state_dict.items(): - if name not in model_sd: - log_fn(f"Not used parameter: {name}") - continue - target = model_sd[name] - if param.shape == target.shape: - out[name] = param - continue - aligned = _align_checkpoint_tensor(param, target, name, log_fn) - if aligned is not None: - out[name] = aligned - return out - - def load_state_dict(checkpoint_path: str, model_class) -> dict: """Load a state dict from a checkpoint directory. @@ -121,51 +50,3 @@ def read_global_step(checkpoint_path: str) -> int | None: return None payload = torch.load(p) return int(payload["global_step"]) - - -def _dir_has_weights(path: str) -> bool: - return os.path.exists( - os.path.join(path, "pytorch_model_fsdp.bin") - ) or os.path.exists(os.path.join(path, "model.safetensors")) - - -def resolve_checkpoint_dir(checkpoint_path: str) -> str: - """Return a directory that directly contains model weights. - - Training saves under a root such as ``libero6/`` with step subdirs - ``libero6/0/``, ``libero6/3/``, etc. Inference callers may pass either the - root or a concrete step directory. - """ - if os.path.isfile(checkpoint_path): - checkpoint_path = os.path.dirname(checkpoint_path) - - if _dir_has_weights(checkpoint_path): - return checkpoint_path - - if not os.path.isdir(checkpoint_path): - raise FileNotFoundError(f"checkpoint path does not exist: {checkpoint_path}") - - candidates: list[tuple[int, float, str]] = [] - for entry in os.listdir(checkpoint_path): - sub = os.path.join(checkpoint_path, entry) - if not os.path.isdir(sub) or not _dir_has_weights(sub): - continue - step = read_global_step(sub) - sort_step = step if step is not None else -1 - candidates.append((sort_step, os.path.getmtime(sub), sub)) - - if not candidates: - return checkpoint_path - - candidates.sort() - resolved = candidates[-1][2] - if resolved != checkpoint_path: - import logging - - logging.getLogger(__name__).info( - "Resolved checkpoint root %s -> %s (global_step=%s)", - checkpoint_path, - resolved, - read_global_step(resolved), - ) - return resolved diff --git a/wall_x/_vendor/harrix/utils/normalizer.py b/wall_x/_vendor/harrix/utils/normalizer.py index 2494b83e..75b4cfb8 100644 --- a/wall_x/_vendor/harrix/utils/normalizer.py +++ b/wall_x/_vendor/harrix/utils/normalizer.py @@ -1,9 +1,15 @@ """Build action/proprio normalizers and resolve the effective norm key. -Public inference artifacts must carry their own normalization data. This module -uses checkpoint-local ``norm_stats.json`` first, then checkpoint-side normalizer -state dicts, and finally an explicit ``customized_action_statistic_dof`` path. -It does not fall back to internal default action statistics. +Priority: + 1. ``norm_stats.json`` for LeRobot checkpoints. + 2. Checkpoint-side ``normalizer_action.pth`` / ``normalizer_propri.pth``. + 3. An explicit ``customized_action_statistic_dof`` path. + +Public inference artifacts must carry their own normalization data. Unlike the +internal build, this module does NOT fall back to a bundled default +``action_statistic_dof`` table (it is not shipped in the OSS package). If the +requested norm key is not present, a prefix match or the only available key is +used as a compatibility fallback. """ from __future__ import annotations @@ -13,13 +19,10 @@ import os import torch +import torch.nn as nn from wall_x.data.backends.lerobot.utils import NormStats -from wall_x.model.core.action.normalizer import Normalizer, pad_normalizer_to_dim -from wall_x._vendor.harrix.utils.train_config import ( - resolve_agent_pos_config, - resolve_dof_config, -) +from wall_x.model.core.action.normalizer import Normalizer logger = logging.getLogger(__name__) @@ -32,39 +35,87 @@ def _load_norm_stats(norm_stats_path: str, action_key: str) -> NormStats: return NormStats(min=q01, max=q99, delta=q99 - q01) -def _load_custom_action_stats(train_config: dict) -> dict | None: - custom = train_config.get("customized_action_statistic_dof", None) - if not custom: - return None - with open(custom, "r") as f: - return json.load(f) - - -def _normalizer_from_stats(action_stats: dict, train_config: dict, key: str) -> Normalizer: - return Normalizer( - action_stats, - train_config[key], - min_key=train_config.get("min_key", "min"), - delta_key=train_config.get("delta_key", "delta"), - ) +def _pad_normalizer_to_dim(normalizer, target_dim, label): + """Tail-pad every dataset entry to ``target_dim`` (virtual ``action_padding`` tail). + + Vendored locally so the public normalizer does not depend on a helper that is + absent from the OSS ``wall_x.model.core.action.normalizer`` module. + """ + target_dim = int(target_dim) + for name in list(normalizer.min.keys()): + current_dim = int(normalizer.min[name].numel()) + if current_dim == target_dim: + continue + if current_dim > target_dim: + raise ValueError( + f"Normalizer {label}[{name!r}] dim ({current_dim}) exceeds configured " + f"{label} dim ({target_dim})" + ) + pad_dim = target_dim - current_dim + dtype = normalizer.min[name].dtype + normalizer.min[name] = nn.Parameter( + torch.cat( + [normalizer.min[name].flatten(), torch.zeros(pad_dim, dtype=dtype)] + ), + requires_grad=False, + ) + normalizer.delta[name] = nn.Parameter( + torch.cat( + [normalizer.delta[name].flatten(), torch.ones(pad_dim, dtype=dtype)] + ), + requires_grad=False, + ) + logger.info( + "Padded %s normalizer[%r] from %d to %d dims", + label, + name, + current_dim, + target_dim, + ) -def _missing_normalizer_error(checkpoint_path: str, train_config: dict) -> FileNotFoundError: - custom = train_config.get("customized_action_statistic_dof", None) +def _missing_default_stats_error(checkpoint_path: str, custom) -> FileNotFoundError: return FileNotFoundError( - "Public inference requires normalization data. Expected one of: " + "Public inference requires per-dataset normalization data in the " + "checkpoint; the internal default action_statistic_dof table is not " + "shipped in the OSS build. Provide one of: " f"{os.path.join(checkpoint_path, 'norm_stats.json')}; checkpoint-side " - "normalizer_action.pth and normalizer_propri.pth; or an explicit " - f"customized_action_statistic_dof path. Current customized_action_statistic_dof={custom!r}." + "normalizer_action.pth / normalizer_propri.pth; or an explicit " + "data.normalizer_config.customized_action_statistic_dof path " + f"(current customized_action_statistic_dof={custom!r})." ) def build_normalizers( checkpoint_path: str, - train_config: dict, + cfg, norm_key: str, ) -> tuple[Normalizer, Normalizer, str]: - """Return action/proprio normalizers and the resolved norm key.""" + """Return action/proprio normalizers and the resolved norm key. + + ``cfg`` is the typed ``TrainConfig``. The returned key may differ from + the requested key if fallback resolution was needed. + """ + # Stats options live under data.normalizer_config; migrated configs may + # carry data-level copies. + norm_cfg = getattr(cfg.data, "normalizer_config", None) or {} + + def _norm_opt(key, default=None): + value = norm_cfg.get(key) + if value is None: + value = getattr(cfg.data, key, None) + return default if value is None else value + + # Custom action statistics are the only supported default source in the OSS + # build; there is no bundled internal action_statistic_dof fallback. + custom = _norm_opt("customized_action_statistic_dof") + if custom: + with open(custom, "r") as f: + action_statistic_dof = json.load(f) + else: + action_statistic_dof = None + + # LeRobot norm_stats.json takes precedence. norm_stats_path = os.path.join(checkpoint_path, "norm_stats.json") if os.path.exists(norm_stats_path): propri_stats = _load_norm_stats(norm_stats_path, "observation.state") @@ -72,26 +123,42 @@ def build_normalizers( normalizer_propri = Normalizer.from_lerobot_norm_stats(propri_stats, norm_key) normalizer_action = Normalizer.from_lerobot_norm_stats(action_stats, norm_key) else: + # Checkpoint-side .pth normalizers take precedence over defaults. action_pth = os.path.join(checkpoint_path, "normalizer_action.pth") propri_pth = os.path.join(checkpoint_path, "normalizer_propri.pth") - custom_stats = _load_custom_action_stats(train_config) - if custom_stats is None and (not os.path.exists(action_pth) or not os.path.exists(propri_pth)): - raise _missing_normalizer_error(checkpoint_path, train_config) - + min_key = _norm_opt("min_key", "min") + delta_key = _norm_opt("delta_key", "delta") if os.path.exists(action_pth): normalizer_action = Normalizer.from_ckpt(action_pth) else: - normalizer_action = _normalizer_from_stats(custom_stats, train_config, "dof_config") - + if action_statistic_dof is None: + raise _missing_default_stats_error(checkpoint_path, custom) + normalizer_action = Normalizer( + action_statistic_dof, + cfg.task.dof_config, + min_key=min_key, + delta_key=delta_key, + ) if os.path.exists(propri_pth): normalizer_propri = Normalizer.from_ckpt(propri_pth) else: - normalizer_propri = _normalizer_from_stats(custom_stats, train_config, "agent_pos_config") - - action_dim = sum(resolve_dof_config(train_config).values()) - propri_dim = sum(resolve_agent_pos_config(train_config).values()) - pad_normalizer_to_dim(normalizer_action, action_dim, "action") - pad_normalizer_to_dim(normalizer_propri, propri_dim, "propri") + if action_statistic_dof is None: + raise _missing_default_stats_error(checkpoint_path, custom) + normalizer_propri = Normalizer( + action_statistic_dof, + cfg.task.agent_pos_config, + min_key=min_key, + delta_key=delta_key, + ) + + # Public checkpoints carry real-dim-only stats (e.g. libero's 7-dim action / + # 8-dim proprio) while the model tensors include a virtual action_padding + # tail. Tail-pad the stats to the full configured dim so virtual-tail + # normalization lines up. Order-independent: only the total dim matters. + action_dim = sum(cfg.task.dof_config.values()) + propri_dim = sum(cfg.task.agent_pos_config.values()) + _pad_normalizer_to_dim(normalizer_action, action_dim, "action") + _pad_normalizer_to_dim(normalizer_propri, propri_dim, "propri") resolved = _resolve_norm_key(norm_key, normalizer_action, normalizer_propri) return normalizer_action, normalizer_propri, resolved @@ -102,7 +169,10 @@ def _resolve_norm_key( normalizer_action: Normalizer, normalizer_propri: Normalizer, ) -> str: - """Resolve a requested norm key against normalizer keys.""" + """Resolve a requested norm key against normalizer keys. + + For example, ``libero_all`` can resolve to ``libero_all_v3``. + """ available = sorted( set(normalizer_action.min.keys()) & set(normalizer_propri.min.keys()) ) diff --git a/wall_x/_vendor/harrix/utils/train_config.py b/wall_x/_vendor/harrix/utils/train_config.py index 435c28af..c93d388b 100644 --- a/wall_x/_vendor/harrix/utils/train_config.py +++ b/wall_x/_vendor/harrix/utils/train_config.py @@ -1,30 +1,63 @@ """Load checkpoint-side train config and construct runtime config objects. -Inference adapters use this module to load the train config, apply checkpoint -overrides for moved processor/tokenizer files, build model config, and build the +Inference adapters use this module to load the train config through wall_x's +single inference-side loader (typed schema only; migrate legacy files with +``scripts/migrate_legacy_yml.py``), build the HF model config, and build the data config required by preprocessing. """ from __future__ import annotations -import os +import dataclasses import logging - -import yaml +import os +from dataclasses import dataclass logger = logging.getLogger(__name__) - -_TASK_INFERENCE_KEYS = ( +# Task keys mirrored into the data: section, matching the training-side dict +# shape (``TrainConfig.build_data_loader_dict``). +_TASK_DATA_SECTION_KEYS = ( "dof_config", "agent_pos_config", - "ar_dof_config", "action_horizon", "action_horizon_flow", - "noise_scheduler", + "ar_dof_config", "use_state_string_representation", ) +_QWEN_IMAGE_FACTOR = 28 +_QWEN_MIN_PIXELS = 3136 +_QWEN_MAX_PIXELS = 12845056 + + +@dataclass +class _LegacyVisionParams: + """Compatibility projection for images without x2robot_utils >= 0.3.0.""" + + resolution: dict + image_factor: int = _QWEN_IMAGE_FACTOR + min_pixels: int = _QWEN_MIN_PIXELS + max_pixels: int = _QWEN_MAX_PIXELS + use_relative_action: bool = False + + +def _build_legacy_vision_params(train_config: dict) -> _LegacyVisionParams: + data = train_config.get("data", {}) or {} + return _LegacyVisionParams( + resolution=dict(data.get("resolution", {})), + image_factor=int(data.get("image_factor", _QWEN_IMAGE_FACTOR)), + min_pixels=int(data.get("min_pixels", _QWEN_MIN_PIXELS)), + max_pixels=int(data.get("max_pixels", _QWEN_MAX_PIXELS)), + use_relative_action=bool(data.get("use_relative_action", False)), + ) + + +# Virtual (padding) dof keys that must sort AFTER the real LeRobot dims so the +# normalizer's virtual-tail padding lines up. Public checkpoints save +# ``action_padding`` at the head of the layout; inference must move it to the +# tail or ``normalize_data_with_virtual_tail`` falls back to whole-tensor +# normalization and mismatches the real-dim stats. _VIRTUAL_TAIL_KEYS = frozenset(("action_padding",)) @@ -40,6 +73,7 @@ def _move_virtual_keys_to_tail(layout: dict | None) -> dict | None: def _canonicalize_task_layouts(target: dict) -> None: + """In-place: move virtual padding dims to the tail across every layout view.""" if os.environ.get("WALLX_CANONICALIZE_VIRTUAL_DOF_ORDER", "1") == "0": return task = target.get("task") @@ -58,314 +92,95 @@ def _canonicalize_task_layouts(target: dict) -> None: data[key] = data_value -def resolve_use_state_string_representation(train_config: dict) -> bool: - """Read ``use_state_string_representation`` with task YAML as authority.""" - task = train_config.get("task") or {} - if task.get("use_state_string_representation") is not None: - return bool(task["use_state_string_representation"]) - if train_config.get("use_state_string_representation") is not None: - return bool(train_config["use_state_string_representation"]) - data = train_config.get("data") or {} - if data.get("use_state_string_representation") is not None: - return bool(data["use_state_string_representation"]) - return False - - -def resolve_state_bins(train_config: dict, default: int = 256) -> int: - """Read discretization bin count from flat or nested train config.""" - if train_config.get("state_bins") is not None: - return int(train_config["state_bins"]) - data = train_config.get("data") or {} - if data.get("state_bins") is not None: - return int(data["state_bins"]) - return default - - -def resolve_agent_pos_config(train_config: dict) -> dict: - """Return ``agent_pos_config`` from flat or ``task``-nested train YAML.""" - task = train_config.get("task") or {} - agent_pos_config = train_config.get("agent_pos_config") or task.get( - "agent_pos_config" - ) - if not agent_pos_config: - raise KeyError( - "agent_pos_config missing from train config " - "(expected top-level or task.agent_pos_config)" - ) - return dict(agent_pos_config) - - -def resolve_dof_config(train_config: dict) -> dict: - """Return ``dof_config`` from flat or ``task``-nested train YAML.""" - task = train_config.get("task") or {} - dof_config = train_config.get("dof_config") or task.get("dof_config") - if not dof_config: - raise KeyError( - "dof_config missing from train config (expected top-level or task.dof_config)" - ) - return dict(dof_config) - - -def resolve_cam_names_from_train_config(train_config: dict) -> list[str] | None: - """Infer model camera keys from ``data.key_mappings.camera`` (e.g. LIBERO).""" - data = train_config.get("data") or {} - key_mappings = data.get("key_mappings") or {} - camera = key_mappings.get("camera") or {} - if not camera: - return None - names: list[str] = [] - for value in camera.values(): - name = str(value) - if name not in names: - names.append(name) - return names or None - - -def resolve_camera_label(cam_name: str, camera_name_mapping: dict | None = None) -> str: - """Match training ``get_wallx_normal_text`` camera display names.""" - mapping = camera_name_mapping or {} - return mapping.get(cam_name, cam_name) - - -def resolve_max_length(train_config: dict, default: int = 768) -> int: - """Read ``max_length`` from flat or nested train YAML (collator default 768).""" - data = train_config.get("data") or {} - if train_config.get("max_length") is not None: - return int(train_config["max_length"]) - if data.get("max_length") is not None: - return int(data["max_length"]) - raw = train_config.get("_raw_data") or {} - if raw.get("max_length") is not None: - return int(raw["max_length"]) - return default - - -# Fields ``load_wallx_processors`` and model construction read from a flat dict. -_ACTION_TOKENIZER_KEYS = ( - "action_tokenizer_type", - "action_tokenizer_path", - "action_tokenizer_checkpoint_path", - "action_tokenizer_config_dir", - "action_tokenizer", -) - -_INFERENCE_MODEL_KEYS = ( - "processor_path", - "pretrained_path", - "config_path", - *_ACTION_TOKENIZER_KEYS, - "ar_loss_weight", - "attn_deterministic", - "flow_loss_weight", - "use_ema", -) - - -def _merge_model_fields(target: dict, source: dict, *, overwrite: bool = False) -> None: - """Copy model-related keys from ``source`` into flat ``target``.""" - if not source: - return - for key in _INFERENCE_MODEL_KEYS: - value = source.get(key) - if value is None: - continue - if overwrite or target.get(key) in (None, ""): - target[key] = value - - -def _mirror_task_fields(target: dict, task: dict) -> None: - """Copy typed task fields into the flat/data legacy inference mirrors.""" - if not isinstance(task, dict): - return - data = target.setdefault("data", {}) - if not isinstance(data, dict): - data = {} - target["data"] = data - for key in _TASK_INFERENCE_KEYS: - value = task.get(key) - if value is None: - continue - target[key] = value - data[key] = value - - -def _load_ckpt_config_overlay(checkpoint_path: str) -> dict: - """Read ``config.yml`` saved beside a training checkpoint.""" - ckpt_yml = os.path.join(checkpoint_path, "config.yml") - if not os.path.exists(ckpt_yml): - return {} - with open(ckpt_yml, "r") as f: - raw = yaml.load(f, Loader=yaml.FullLoader) or {} - - overlay: dict = {} - _merge_model_fields(overlay, raw.get("model") or {}, overwrite=True) - raw_yaml = raw.get("_raw_yaml") or {} - _merge_model_fields(overlay, raw_yaml.get("model") or {}, overwrite=False) - _merge_model_fields(overlay, raw, overwrite=False) - return overlay - - -def strip_action_tokenizer_fields(train_config: dict) -> None: - """Drop action-tokenizer fields so flow inference can skip AR tokenizer setup.""" - for key in _ACTION_TOKENIZER_KEYS: - train_config.pop(key, None) - model = train_config.get("model") - if isinstance(model, dict): - for key in _ACTION_TOKENIZER_KEYS: - model.pop(key, None) - - def load_train_config_with_ckpt_overlay( train_config_path: str, checkpoint_path: str, -) -> dict: - """Load train YAML and apply checkpoint-local processor/tokenizer overlays. - - This lets a checkpoint remain portable when the original training-machine - processor or action-tokenizer paths are no longer available. - """ - with open(train_config_path, "r") as f: - train_config = yaml.load(f, Loader=yaml.FullLoader) - - preprocessor_file = os.path.join(checkpoint_path, "preprocessor_config.json") - if os.path.exists(preprocessor_file): - train_config["processor_path"] = checkpoint_path - - orig_action_tok = train_config.get("action_tokenizer_path", None) - if orig_action_tok is not None and not os.path.exists(orig_action_tok): - tokenizer_file = os.path.join(checkpoint_path, "tokenizer.json") - tokenizer_config_file = os.path.join(checkpoint_path, "tokenizer_config.json") - if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file): - train_config["action_tokenizer_path"] = checkpoint_path - - ckpt_overlay = _load_ckpt_config_overlay(checkpoint_path) - _merge_model_fields(train_config, ckpt_overlay, overwrite=False) - if isinstance(train_config.get("model"), dict): - _merge_model_fields(train_config["model"], ckpt_overlay, overwrite=False) - - return train_config - - -def normalize_train_config_for_inference( - train_config: dict, - train_config_path: str, -) -> dict: - """Flatten typed TrainConfig YAML into the legacy dict inference expects. +): + """Load the checkpoint train config through the single wall_x loader. - Serving and ``update_model_config`` still read top-level ``dof_config`` and - nested ``data.*`` fields. New training yamls store those under ``task:``. + Returns ``(cfg, train_config)``: the typed ``TrainConfig`` (ckpt overlays + applied; the single source consumers should read) plus its dict view + shaped like the training-side ``TrainConfig.build_data_loader_dict`` + output -- kept solely for the dict-contract consumers listed above. """ - try: - from wall_x.config.loader import load_config - - typed_cfg = load_config(train_config_path) - except Exception as e: - logger.debug( - "Train config is not typed TrainConfig schema (%s): %s", - train_config_path, - e, - ) - _mirror_task_fields(train_config, train_config.get("task") or {}) - _canonicalize_task_layouts(train_config) - return train_config - - import dataclasses - - normalized = typed_cfg.build_data_loader_dict() - model_dict = dataclasses.asdict(typed_cfg.model) - _merge_model_fields(normalized, model_dict, overwrite=True) - - normalized["data"] = dict(normalized.get("data") or {}) - if isinstance(train_config.get("model"), dict): - _merge_model_fields(normalized, train_config["model"], overwrite=False) - for key, value in train_config.items(): - if key in _INFERENCE_MODEL_KEYS or key == "qwen_vl_act_config_path": - if value is not None: - normalized[key] = value - elif key == "data" and isinstance(value, dict): - normalized["data"].update(value) - - _mirror_task_fields(normalized, dataclasses.asdict(typed_cfg.task)) - _canonicalize_task_layouts(normalized) - return normalized - - -def register_data_backend(train_config: dict) -> None: - """Register the data backend before typed data config construction.""" - from wall_x.data._registry import _set_data_backend - - data_section = train_config.get("data", {}) - dataset_type = train_config.get("dataset_type") or data_section.get( - "dataset_type", "lerobot" - ) - _set_data_backend(dataset_type) - - -def build_model_config( - config_class, - checkpoint_path: str, - train_config: dict, - train_config_path: str | None = None, -): + from wall_x.config.loader import dump_config, load_train_config_from_checkpoint + + cfg = load_train_config_from_checkpoint(checkpoint_path, train_config_path) + train_config = dump_config(cfg) + task_dict = dataclasses.asdict(cfg.task) + train_config.update(task_dict) + data_section = train_config.setdefault("data", {}) + for key in _TASK_DATA_SECTION_KEYS: + value = task_dict.get(key) + if value is not None: + data_section.setdefault(key, value) + # Public checkpoints save the virtual ``action_padding`` dim at the head of + # the layout; move it to the tail so the normalizer's virtual-tail masking + # lines up. Canonicalize BOTH the typed ``cfg.task`` (what the normalizer and + # preprocessing read) and the dict view (kept for dict-contract callers). + if os.environ.get("WALLX_CANONICALIZE_VIRTUAL_DOF_ORDER", "1") != "0": + for _key in ("dof_config", "agent_pos_config", "ar_dof_config"): + _canon = _move_virtual_keys_to_tail(getattr(cfg.task, _key, None)) + if _canon is not None: + setattr(cfg.task, _key, _canon) + _canonicalize_task_layouts(train_config) + return cfg, train_config + + +def build_model_config(config_class, checkpoint_path: str, cfg, train_config: dict): """Build the HF model config and inject train_config-derived fields. - ``config_class`` is supplied by the variant adapter. If the checkpoint does - not include ``config.json``, this falls back to - ``train_config["qwen_vl_act_config_path"]``. + ``config_class`` is supplied by the variant adapter. A checkpoint-local + ``config.json`` wins; otherwise ``cfg.model.config_path`` is used. """ ckpt_config_path = os.path.join(checkpoint_path, "config.json") if os.path.exists(ckpt_config_path): resolved = ckpt_config_path else: - resolved = train_config.get("qwen_vl_act_config_path") - if resolved is None or not os.path.exists(resolved): + resolved = getattr(cfg.model, "config_path", None) + if not resolved or not os.path.exists(resolved): raise ValueError( f"cannot load model config: checkpoint file {ckpt_config_path} " - f"does not exist, and fallback qwen_vl_act_config_path={resolved!r} " + f"does not exist, and fallback cfg.model.config_path={resolved!r} " "does not exist either" ) - train_config["qwen_vl_act_config_path"] = resolved if resolved.endswith(".json"): model_config = config_class.from_json_file(resolved) else: model_config = config_class.from_pretrained(resolved) - legacy_train_config = ( - normalize_train_config_for_inference(train_config, train_config_path) - if train_config_path - else train_config - ) - model_config.update_model_config(legacy_train_config) + model_config.update_model_config(train_config) model_config._attn_implementation = "sdpa" model_config.vision_config._attn_implementation = "flash_attention_2" return model_config -def build_data_config(train_config_path: str, train_config: dict): +def build_data_config(cfg, train_config: dict): """Build the data config used by image resizing and preprocessing. - New checkpoints use the typed schema. Older checkpoints may still use a flat - legacy schema, so this falls back to the active data backend. + x2robot_v2 / lerobot: ``load_x2r_data_config`` returns None for these + backends, so image params + resolution come from the train/infer-shared + projection in x2robot_utils (>= 0.3.0). The ``load_x2r_data_config`` call + sites are rewritten to ``load_trainer_data_config`` at OSS export time -- + see ``scripts/export_opensource.py::patch_vendored_harrix``. """ - from wall_x.config.loader import load_config - from wall_x.data import data_backend - - try: - typed_cfg = load_config(train_config_path) - backend = data_backend() - if backend.supports("load_trainer_data_config"): - return backend.load_trainer_data_config(typed_cfg) - except Exception as e: - logger.warning( - "typed TrainConfig loading failed; falling back to raw data config: %s", - e, - ) - - backend = data_backend() - if backend.supports("load_trainer_data_config_from_yaml_dict"): - return backend.load_trainer_data_config_from_yaml_dict(train_config) - raise RuntimeError( - f"active data backend {backend!r} cannot build a trainer data config" - ) + if cfg.data.dataset_type in ("x2robot_v2", "lerobot"): + try: + from wall_x._vendor.x2robot_utils.vision_params import build_vision_params + except ModuleNotFoundError as error: + if error.name != "x2robot_utils.vision_params": + raise + logger.warning( + "x2robot_utils.vision_params is unavailable; using Harrix's " + "equivalent inference-only vision parameter projection. " + "Upgrade x2robot_utils to >= 0.3.0 to use the shared source." + ) + return _build_legacy_vision_params(train_config) + + return build_vision_params(train_config) + + from wall_x.trainer.adapters.base_adapter import load_x2r_data_config + + return load_x2r_data_config(cfg) diff --git a/wall_x/_vendor/x2robot_utils/geometry.py b/wall_x/_vendor/x2robot_utils/geometry.py index 3816ce12..0125da4c 100644 --- a/wall_x/_vendor/x2robot_utils/geometry.py +++ b/wall_x/_vendor/x2robot_utils/geometry.py @@ -1,9 +1,9 @@ """Rotation / pose geometry utilities (pure numpy + numba). Shared between wall-x and internal_dataset_backend: -- ``euler_to_matrix_zyx_6d_nb``: ZYX Euler -> flattened top-2-rows of R (N, 6) -- ``so3_to_euler_zyx_batch_nb``: 6D rotation -> ZYX Euler (canonicalized) -- ``compose_state_and_delta_to_abs_{rpy,6d}``: state + delta -> absolute pose +- ``euler_to_matrix_zyx_6d_nb``: ZYX Euler → flattened top-2-rows of R (N, 6) +- ``so3_to_euler_zyx_batch_nb``: 6D rotation → ZYX Euler (canonicalized) +- ``compose_state_and_delta_to_abs_{rpy,6d}``: state ⊕ delta → absolute pose """ import numpy as np @@ -12,7 +12,7 @@ @jit(nopython=True) def euler_to_matrix_zyx_6d_nb(eulers): - """Euler angles (N, 3) -> flattened top two rows of rotation matrix (N, 6).""" + """Euler angles (N, 3) → flattened top two rows of rotation matrix (N, 6).""" N = eulers.shape[0] R6 = np.empty((N, 6), dtype=np.float64) for i in prange(N): @@ -70,7 +70,7 @@ def euler_to_matrix_zyx_batch_nb(eulers): @jit(nopython=True) def matrix_to_euler_zyx_batch_nb(Rs): - """R = Rz(yaw) * Ry(pitch) * Rx(roll) -> (roll, pitch, yaw).""" + """R = Rz(yaw) * Ry(pitch) * Rx(roll) → (roll, pitch, yaw).""" N = Rs.shape[0] eulers = np.empty((N, 3), dtype=np.float64) for i in prange(N): @@ -137,7 +137,7 @@ def so3_to_matrix_batch_nb(batch_so3): @jit(nopython=True) def canonicalize_euler_zyx_batch_nb(rpy_batch): - """Canonicalize ZYX Euler angles so each component falls in (-pi, pi].""" + """Canonicalize ZYX Euler angles so each component falls in (-π, π].""" N = rpy_batch.shape[0] out = np.empty_like(rpy_batch) two_pi = 2.0 * np.pi @@ -179,10 +179,10 @@ def so3_to_euler_zyx_batch_nb(batch_so3): @jit(nopython=True) def compose_state_and_delta_to_abs_rpy(delta, state): - """Compose a delta (ZYX rpy or 6D) with an absolute state -> absolute rpy(ZYX). + """Compose a delta (ZYX rpy or 6D) with an absolute state → absolute rpy(ZYX). - delta: (N, 3) deltarpy or (N, 6) delta6D. state: (3,) rpy or (6,) 6D. - Output: (N, 3) rpy canonicalized into (-pi, pi]. + delta: (N, 3) Δrpy or (N, 6) Δ6D. state: (3,) rpy or (6,) 6D. + Output: (N, 3) rpy canonicalized into (-π, π]. """ if delta.shape[-1] == 3: R_delta = euler_to_matrix_zyx_batch_nb(delta) @@ -242,7 +242,7 @@ def compose_state_and_delta_to_abs_rpy(delta, state): @jit(nopython=True) def compose_state_and_delta_to_abs_6d(delta, state): - """Compose a 6D delta with a 6D state -> absolute 6D rotation. + """Compose a 6D delta with a 6D state → absolute 6D rotation. delta: (N, 6). state: (6,). Output: (N, 6). """ diff --git a/wall_x/_vendor/x2robot_utils/grounding.py b/wall_x/_vendor/x2robot_utils/grounding.py index 7e7eb5d3..0852cd01 100644 --- a/wall_x/_vendor/x2robot_utils/grounding.py +++ b/wall_x/_vendor/x2robot_utils/grounding.py @@ -26,7 +26,7 @@ def process_match(match): if model_type == "qwen2_5": new_x = max(0, min(round(x * scale_w), resized_width - 1)) new_y = max(0, min(round(y * scale_h), resized_height - 1)) - elif model_type in ["qwen2"]: + elif model_type in ["qwen2", "qwen3", "qwen3_5"]: new_x = max(0, min(999.999, (x / orig_width) * 1000)) new_y = max(0, min(999.999, (y / orig_height) * 1000)) else: @@ -40,7 +40,7 @@ def process_match(match): new_y1 = max(0, min(round(y1 * scale_h), resized_height - 1)) new_x2 = max(0, min(round(x2 * scale_w), resized_width - 1)) new_y2 = max(0, min(round(y2 * scale_h), resized_height - 1)) - elif model_type in ["qwen2"]: + elif model_type in ["qwen2", "qwen3", "qwen3_5"]: new_x1 = max(0, min(999.999, (x1 / orig_width) * 1000)) new_y1 = max(0, min(999.999, (y1 / orig_height) * 1000)) new_x2 = max(0, min(999.999, (x2 / orig_width) * 1000)) @@ -78,7 +78,7 @@ def extract_grounding_points(text: str) -> List[List[float]]: def reverse_grounding_points( text: str, orig_height, orig_width, resized_height, resized_width, model_type ) -> str: - """Inverse of ``process_grounding_points`` - map resized coords back to original.""" + """Inverse of ``process_grounding_points`` — map resized coords back to original.""" point_pattern = re.compile(r"<(point|box|bbox)>(.*?)") def reverse_match(match): @@ -95,7 +95,7 @@ def reverse_match(match): if model_type == "qwen2_5": orig_x = max(0, min(orig_width - 1, round(x / scale_w))) orig_y = max(0, min(orig_height - 1, round(y / scale_h))) - elif model_type in ["qwen2"]: + elif model_type in ["qwen2", "qwen3", "qwen3_5"]: orig_x = max( 0, min(orig_width - 1, round((x / 1000) * orig_width)) ) @@ -113,7 +113,7 @@ def reverse_match(match): orig_y1 = max(0, min(orig_height - 1, round(y1 / scale_h))) orig_x2 = max(0, min(orig_width - 1, round(x2 / scale_w))) orig_y2 = max(0, min(orig_height - 1, round(y2 / scale_h))) - elif model_type in ["qwen2"]: + elif model_type in ["qwen2", "qwen3", "qwen3_5"]: orig_x1 = max( 0, min(orig_width - 1, round((x1 / 1000) * orig_width)) ) diff --git a/wall_x/_vendor/x2robot_utils/vision_params.py b/wall_x/_vendor/x2robot_utils/vision_params.py new file mode 100644 index 00000000..da36c810 --- /dev/null +++ b/wall_x/_vendor/x2robot_utils/vision_params.py @@ -0,0 +1,58 @@ +"""Vision preprocessing parameters -- the single source of truth shared by +training and inference. + +These are the parameters the model-side image preprocessing (Qwen +``smart_resize``) needs: per-camera resolution, the resize constants, and the +relative-action marker. The training side (when each data backend builds +samples) and the inference side (when the harrix adapter preprocesses camera +frames) must use the same semantics, otherwise the train/infer distributions +diverge. + +The constants are pinned as literals here rather than forwarded from +``qwen_vl_utils``: a third-party library's defaults drift across versions +(``qwen_vl_utils`` 0.0.14 removed ``IMAGE_FACTOR``), whereas these values are a +reproducibility anchor consistent with the Qwen2.5-VL processor spec. A +checkpoint's ``train_config`` may override any field. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +# Qwen2.5-VL smart_resize standard constants (image_factor=28, +# min_pixels=28**2*4, max_pixels=28**2*16384). +QWEN_IMAGE_FACTOR = 28 +QWEN_MIN_PIXELS = 3136 +QWEN_MAX_PIXELS = 12845056 + + +@dataclass +class VisionParams: + """Projection of the image preprocessing parameters. + + ``resolution`` maps each camera name to its edge length; ``use_relative_action`` + is the action-space marker, consumed by the inference adapter alongside the + preprocessing config, so it is folded into this projection. + """ + + resolution: dict + image_factor: int = QWEN_IMAGE_FACTOR + min_pixels: int = QWEN_MIN_PIXELS + max_pixels: int = QWEN_MAX_PIXELS + use_relative_action: bool = False + + +def build_vision_params(train_config: dict) -> VisionParams: + """Build a :class:`VisionParams` from a checkpoint's (flat) ``train_config``. + + Reads the same-named fields under ``train_config["data"]``, falling back to + the Qwen standard constants when absent. + """ + data = train_config.get("data", {}) or {} + return VisionParams( + resolution=dict(data.get("resolution", {})), + image_factor=int(data.get("image_factor", QWEN_IMAGE_FACTOR)), + min_pixels=int(data.get("min_pixels", QWEN_MIN_PIXELS)), + max_pixels=int(data.get("max_pixels", QWEN_MAX_PIXELS)), + use_relative_action=bool(data.get("use_relative_action", False)), + ) diff --git a/wall_x/config/__init__.py b/wall_x/config/__init__.py index 329780a9..10887c31 100644 --- a/wall_x/config/__init__.py +++ b/wall_x/config/__init__.py @@ -6,6 +6,7 @@ ConstantSchedulerConfig, CosineSchedulerConfig, DMuonConfig, + NaiveFSDP2MuonConfig, OptimizerConfig, SchedulerConfig, StepSchedulerConfig, @@ -14,8 +15,11 @@ from .infra_config import ( CheckpointConfig, DebugConfig, + DMuonBackendConfig, DistributedConfig, + FSDPBackendConfig, LoggingConfig, + ShardedBackendConfig, ) from .loader import load_config, save_config from .model_config import ModelConfig, QActModelConfig @@ -43,11 +47,15 @@ "OptimizerConfig", "AdamWConfig", "DMuonConfig", + "NaiveFSDP2MuonConfig", "SchedulerConfig", "CosineSchedulerConfig", "ConstantSchedulerConfig", "StepSchedulerConfig", # Infra configs + "ShardedBackendConfig", + "FSDPBackendConfig", + "DMuonBackendConfig", "DistributedConfig", "LoggingConfig", "CheckpointConfig", diff --git a/wall_x/config/data_config.py b/wall_x/config/data_config.py index 9ed7f94f..f6749cdc 100644 --- a/wall_x/config/data_config.py +++ b/wall_x/config/data_config.py @@ -6,7 +6,7 @@ """ from dataclasses import dataclass, field -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional from .registry import register_data_config @@ -45,6 +45,7 @@ class LeRobotDataConfig(DataConfig): """ dataset_type: str = "lerobot" + dataset_config_path: Optional[str] = None lerobot_config: Optional[Dict[str, Any]] = None key_mappings: Optional[Dict[str, Any]] = None norm_stats_path: Optional[str] = None @@ -55,6 +56,19 @@ class LeRobotDataConfig(DataConfig): use_fast_tokenizer: bool = False padding_side: str = "left" noise_scheduler: Optional[Dict[str, Any]] = None + # --- qwen preprocess knobs (read by prompt/state encoding) --- + state_bins: int = 256 + use_embodied_system_prompt_ratio: float = 0.0 + rgb_keys: Optional[List[str]] = None + lowdim_keys: Optional[List[str]] = None + skip_interval: int = 2 + image_history_length: int = 1 + obs_history_stride: int = 0 + event_level: bool = False + image_horizon: int = 16 + future_image_length: int = 16 + future_image_interval: int = 2 + action_history_length: int = 1 __all__ = [ diff --git a/wall_x/config/hyperparams_config.py b/wall_x/config/hyperparams_config.py index 10a57b68..7259185b 100644 --- a/wall_x/config/hyperparams_config.py +++ b/wall_x/config/hyperparams_config.py @@ -32,6 +32,10 @@ class OptimizerConfig: action_expert_learning_rate: Optional[float] = None action_lr_keywords: Optional[List[str]] = None + @property + def use_dedicated_params(self) -> bool: + return False + @register_optimizer_config("adamw") @dataclass @@ -63,6 +67,31 @@ class DMuonConfig(OptimizerConfig): ns_backend: str = "gram" ns_coefficients: str = "default" nesterov: bool = True + owner_strategy: str = "lpt" + owner_cost_model: str = "optimizer" + owner_pack_small_params: bool = True + hsdp_column_balance: bool = True + record_step_profile: bool = False + first_step_progress_log_limit: int = 8 + replicate_async: bool = True + group_prepare_ahead: bool = True + forward_prefetch_depth: int = 1 + + @property + def use_dedicated_params(self) -> bool: + return True + + +@register_optimizer_config("naive_fsdp2_muon") +@dataclass +class NaiveFSDP2MuonConfig(DMuonConfig): + """Temporary FSDP2 baseline: every rank runs full-matrix Muon locally.""" + + optimizer_type: str = "naive_fsdp2_muon" + + @property + def use_dedicated_params(self) -> bool: + return False @dataclass @@ -91,9 +120,10 @@ class CosineSchedulerConfig(SchedulerConfig): @register_scheduler_config("constant") @dataclass class ConstantSchedulerConfig(SchedulerConfig): - """Constant learning rate with no decay.""" + """Constant learning rate, optionally after linear warmup.""" scheduler_type: str = "constant" + num_warmup_steps: int = 0 @register_scheduler_config("step") @@ -108,10 +138,26 @@ class StepSchedulerConfig(SchedulerConfig): @dataclass class TrainHyperParams: + # 训练算法判别器（依附于它所多态的 hyperparams section，与 model_type↔model 对称）： + # 选 TrainHyperParams 子类 + DDP trainer 路径。 + # "sl"(默认)=监督学习（model_type=rlt 时为 encoder warmup）; "td3"=off-policy TD3。 + # loader._resolve_hyperparams_cls 按 (algorithm, model_type) 惰性路由到 RL 子类。 + algorithm: str = "sl" num_epoch: int = 1 batch_size_per_gpu: int = 1 gradient_accumulation_steps: int = 1 seed: int = 42 + # opt-in 严格确定性：调 torch.use_deterministic_algorithms 让 same-seed + + # same-yaml 两次 run 在 pi0 路径下 step 0 forward byte-equal。caller 必须 + # 同步在 launch 脚本里 export CUBLAS_WORKSPACE_CONFIG=:4096:8（库代码不 + # silently setdefault env，避免 fork 出来的子进程 env 难追溯）。warn_only + # 内置 True：bf16 训练里总有几个 op 没确定实现，warn 比 crash 实用。 + # default False 保兼容（公共 utility，pi0 以外的 model 可能含非 deterministic + # op，强开会抛 RuntimeError）。 + # 这是确定性的**唯一开关**：除 torch.use_deterministic_algorithms 外，还派生 + # model_config.attn_deterministic（pi0 / qwen2_5 / qwen3 / qwen3_5 的 attention + # 内核）——旧的独立 ``attn_deterministic`` yaml 字段已下线（MR !556 review）。 + deterministic_algorithms: bool = False optimizer: OptimizerConfig = field(default_factory=AdamWConfig) scheduler: SchedulerConfig = field(default_factory=CosineSchedulerConfig) @@ -122,6 +168,7 @@ class TrainHyperParams: "CosineSchedulerConfig", "DMuonConfig", "LRGroupConfig", + "NaiveFSDP2MuonConfig", "OptimizerConfig", "SchedulerConfig", "StepSchedulerConfig", diff --git a/wall_x/config/infra_config.py b/wall_x/config/infra_config.py index 64c9ceca..b07a9a67 100644 --- a/wall_x/config/infra_config.py +++ b/wall_x/config/infra_config.py @@ -1,35 +1,164 @@ -"""Infrastructure config: distributed runtime, logging, checkpoints, and debug.""" +"""Infrastructure config: distributed runtime, logging, checkpoints, debug, and serving.""" -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Optional +@dataclass +class ModelServiceConfig: + """Model-service push config (used by CkptPusher — trainer pushes new actor + weights back to serving).""" + + model_service_url: str = "" + push_max_retries: int = 3 + + +@dataclass +class RLDataServiceConfig: + """RLDataSource ctor kwargs for RL off-policy ingest (see ADR-0007). + + Collects RLDataSource's 5 kwargs under the yaml ``cfg.serving.data`` block so + ``train_rl.py::_build_trainer`` can pass them straight to ``RLDataSource()``. + (Episode physical location comes from the data service's discover ``data_uri``, + not configured here — the trainer does not couple to storage layout.)""" + + data_service_url: str = "" + tag: str = "" + trainer_encoder_version: str = "" + trainer_pi0_model_hash: str = "" + trainer_data_pipeline_version: str = "" + + +@dataclass +class ServingConfig: + """Unified config for the two serving counterparts the trainer talks to. + + ``model`` = model service (trainer → serving actor-weight push); ``data`` = + data service (robot → DS upload, trainer ← DS ingest). Both are serving + components the training process talks to over HTTP, so they are grouped under + ``serving:``.""" + + model: ModelServiceConfig = field(default_factory=ModelServiceConfig) + data: RLDataServiceConfig = field(default_factory=RLDataServiceConfig) + + +@dataclass +class ShardedBackendConfig: + """Shared 1D/2D sharding layout options.""" + + sharding_strategy: str = "full_shard" + hsdp_replicate_size: Optional[int] = None + + @property + def reshard_after_forward(self) -> bool: + return self.sharding_strategy not in ("shard_grad_op", "_hybrid_shard_zero2") + + +@dataclass +class FSDPBackendConfig(ShardedBackendConfig): + """PyTorch FSDP2-specific distributed backend options.""" + + backward_prefetch: str = "backward_pre" + cpu_offload: bool = False + use_orig_params: bool = True + limit_all_gathers: bool = True + forward_prefetch: bool = False + sync_module_states: bool = True + save_policy: str = "full" + reduce_dtype: str = "bf16" + + +@dataclass +class DMuonBackendConfig(ShardedBackendConfig): + """Pure DMuon distributed backend options.""" + + sharding_strategy: str = "shard_grad_op" + + +@dataclass +class DDPBackendConfig: + """PyTorch DistributedDataParallel backend options.""" + + find_unused_parameters: bool = False + broadcast_buffers: bool = True + bucket_cap_mb: int = 25 + + @dataclass class DistributedConfig: - # FSDP - use_fsdp: bool = False - fsdp_sharding_strategy: str = "full_shard" - fsdp_backward_prefetch: str = "backward_pre" - fsdp_cpu_offload: bool = False - fsdp_use_orig_params: bool = True - fsdp_limit_all_gathers: bool = True - fsdp_forward_prefetch: bool = False - fsdp_sync_module_states: bool = True - fsdp_save_policy: str = "full" - fsdp_hsdp_replicate_size: Optional[int] = None + # Public distributed backend. DMuon means pure DMuon; FSDP2+DMuon dedicated + # params is intentionally not a separate backend. DDP is kept as a + # lightweight debug/small-model backend. + backend: str = "fsdp2" + sp_size: int = 1 + fsdp: FSDPBackendConfig = field(default_factory=FSDPBackendConfig) + dmuon: DMuonBackendConfig = field(default_factory=DMuonBackendConfig) + ddp: DDPBackendConfig = field(default_factory=DDPBackendConfig) # Mixed precision use_mixed_precision: bool = True bf16: bool = True - fsdp_reduce_dtype: str = "bf16" use_amp: bool = False use_gradient_checkpointing: bool = False use_gradient_checkpointing_offload: bool = False + # Partial recompute: number of leading transformer layers to checkpoint. + # None = checkpoint all layers (legacy behavior when use_gradient_checkpointing=True). + # N in [0, num_layers] = checkpoint only the first N layers, keep the rest's + # activations resident in HBM. Clamped to [0, num_layers] at apply time. + # Ignored when use_gradient_checkpointing=False. + gradient_checkpoint_num_layers: Optional[int] = None use_selective_recompute: bool = False # DDP fallback find_unused_parameters: bool = False broadcast_buffers: bool = True bucket_cap_mb: int = 25 + @property + def use_fsdp(self) -> bool: + """Compatibility property for code paths that require FSDP2 wrappers.""" + + return self.backend == "fsdp2" + + +@dataclass +class Pi0DistributedConfig(DistributedConfig): + """PI0/PI05-only distributed knobs (split out of the base config so other + models' schema stays clean — MR !556 review). + + Consumed only by ``pi0_pytorch.convert_to_fsdp`` / + ``distribution_strategy``; built for ``model_type in {pi0, pi05}`` (see + ``loader._build_distributed_config``). For other models these keys are + absent from ``dataclasses.asdict(distributed)``, so ``build_strategy``'s + ``if k in config`` forward skips them and the strategy uses its defaults. + """ + + # Checkpoint the SigLIP vision encoder (paligemma.vision_tower). True + # (default) preserves legacy behavior — embed_image goes through + # torch.utils.checkpoint, SigLIP forward recomputes during backward. False + # skips ckpt on the vision call site: vision activation stays in HBM + # (+~1-2 GiB per rank for batched 3-cam bs/gpu=32) and the long + # `paligemma.vision_tower` recompute block disappears from the backward + # timeline. Independent of gradient_checkpoint_num_layers (which only + # controls the 18-layer joint transformer). Ignored when + # use_gradient_checkpointing=False. + gradient_checkpoint_vision_tower: bool = True + # 把 root 上的 non-layer module（vision_tower / embed_tokens / + # multi_modal_projector / suffix adapter / final norms）从 FSDP2 nested + # shard 中拆出走 DDP 风格（params + grad 全副本，每 step backward 后 + # all_reduce 同步 grad）。motivation: 这些模块的 unshard all_gather 在 + # step 开头占 nccl schedule 大头，去掉换 DDP 通信简化 timeline。代价： + # 每 rank +~1.8 GB params bf16 + ~11 GB AdamW 状态（无 ZeRO-1 时）。 + replicate_non_layer_modules: bool = False + # 在 vision_tower 的 FSDP pre_forward 里显式 prefetch joint_layer[0] + # 的 all_gather。默认 FSDP2 prefetch 是 i→i+1 chain，joint_layer[0] 没有 + # 上游帮它 prefetch，必须等进入它自己的 forward 才发 all_gather → 卡在 + # vision_tower 之后。开启后 joint_layer[0] all_gather 在 copy stream 上 + # 跟 vision_tower compute 并行，timeline 上的"vision_tower 完 → 等 + # joint_layer[0] all_gather"那段 gap 消失。HBM 代价：joint_layer[0] 的 + # unsharded params 多驻留 (vision_tower forward 时段)，~30 MB bf16 量级 + # 可忽略。要求 vision_tower 走 FSDP shard 路径（与 replicate_non_layer_modules + # 互斥；后者为 True 时此 toggle no-op）。 + fsdp_forward_prefetch_joint0_with_vision: bool = True + @dataclass class LoggingConfig: @@ -42,11 +171,17 @@ class LoggingConfig: save_interval: int = 1000 val_interval: int = 4000 epoch_save_interval: int = 1 + enable_validation: bool = True gc_interval_steps: int = 1000 ignore_until_interval: int = 0 + # 前几步 step_time / samples_per_sec 不写 wandb：cuDNN benchmark warm-up、 + # dataloader 首次预取、lazy init 让 step 0~N 远比稳态慢，写进去会把 y 轴拉爆。 + # 控制台 line 仍然显示原值，便于现场看到 warm-up 发生。 + perf_log_skip_steps: int = 5 + param_norm_interval: int = 100 # Rolling-window for smoothing per-step training metrics displayed on the # console (and reused by tqdm). 1 = raw per-step (historical behavior). - # 10 = DZ-style 10-step rolling average - diffusion losses are dominated + # 10 = DZ-style 10-step rolling average — diffusion losses are dominated # by timestep-sampling noise per step; the smoothing only changes display, # not training. Independent of log_interval (which buffers for wandb). loss_log_smooth_window: int = 1 @@ -72,6 +207,11 @@ class CheckpointConfig: resume_rng: Optional[str] = None resume_data: Optional[str] = None resume_step: Optional[str] = None + # When True, passes try_harder=True to reshape_compatible_state_dict so that + # action-preprocessor weights are padded/sliced to match the current model shape. + try_harder: bool = False + # CkptPusher: path to the JSON log of already-pushed weight versions. + weight_version_persistence_path: str = "" def get_resume_path(self, component: str) -> Optional[str]: """Return the resume path for one checkpoint component.""" @@ -88,6 +228,47 @@ class DebugConfig: profile_wait_iters: int = 1 profile_warmup_iters: int = 1 profile_active_iters: int = 3 + # Number of profile cycles to record. Default 1 = capture once and stop. + # torch.profiler.schedule's own default is 0 = repeat forever, which was + # the silent behavior when this field was missing (the yaml key was being + # dropped by the typed loader). Each cycle dumps a ~500MB trace per worker + # and adds a multi-second stall, so 1 is the right default for ad-hoc runs. + profile_repeat: int = 1 + # When True, traces are gzip-compressed by torch.profiler's + # tensorboard_trace_handler. Output files become *.pt.trace.json.gz + # (~10× smaller). Perfetto / chrome://tracing accept .gz directly. + profile_use_gzip: bool = False + # torch.profiler.profile() kwargs. Defaults tuned for "post-hoc analysis + # via visualize_wallx" rather than the legacy "tensorboard + py-spy stack" + # workflow: + # - with_stack=False saves a lot of trace size and CPU overhead BUT it + # drops the `nn.Module: ` annotations (which `with_stack` + # inferred indirectly from python frames). wall-x compensates by + # emitting explicit record_function zones — see wall_x/utils/profile_zones.py. + # - with_modules=True has no effect in eager mode (PyTorch limits it + # to TorchScript), so the explicit-zone approach is the only path. + # - record_shapes keeps GEMM input shapes (visualize uses it for FLOPS) + # - profile_memory off by default — overlaps with enable_memory_snapshot + profile_with_stack: bool = False + profile_with_modules: bool = False # no-op in eager mode — keep off + profile_record_shapes: bool = True + profile_memory: bool = False + # When True, only rank 0 records a trace. Avoids dumping N copies of the + # same timeline on multi-GPU runs (16-rank × ~80MB each = 1.3GB per cycle). + profile_only_rank0: bool = True + # CUDA memory snapshot. When True, rank 0 records the allocator history + # starting at step 0 and dumps a .pickle at step `memory_snapshot_until_step`. + # Output drops into `memory_snapshot_dir/memory_snapshot.pickle`; open it at + # https://pytorch.org/memory_viz to see per-allocation timeline + peak + + # which Python frame owned each block. + enable_memory_snapshot: bool = False + memory_snapshot_dir: str = "./memory_snapshot" + # Recording starts in trainer __init__ (covers model load + FSDP wrap + + # optimizer creation); dump fires right after global_step == until_step. + # 3 is enough to see steady-state per-step alloc pattern without making + # the .pickle huge. + memory_snapshot_until_step: int = 3 + memory_snapshot_max_entries: int = 100000 show_time_details: bool = False visualize_sample: bool = False save_debug_batch_path: Optional[str] = None @@ -96,3 +277,8 @@ class DebugConfig: enable_mfu: bool = False # Optional FLOPs profiling runs an extra forward at step 0. enable_mfu_profile: bool = False + # Peak per-GPU bf16 (dense) throughput in TFLOP/s, used as the MFU + # denominator. Leave null to auto-detect from torch.cuda.get_device_name(); + # set explicitly when your GPU/precision isn't in the built-in table + # (otherwise MFU on an unlisted card is meaningless). + mfu_peak_tflops: Optional[float] = None diff --git a/wall_x/config/loader.py b/wall_x/config/loader.py index c22bcd5b..1e8bc5dc 100644 --- a/wall_x/config/loader.py +++ b/wall_x/config/loader.py @@ -3,7 +3,8 @@ import dataclasses import os import shutil -from typing import Any, Type, TypeVar +import warnings +from typing import Any, Optional, Type, TypeVar import yaml @@ -20,6 +21,8 @@ DebugConfig, DistributedConfig, LoggingConfig, + Pi0DistributedConfig, + ServingConfig, ) from .model_config import ModelConfig, QActModelConfig from .registry import ( @@ -68,28 +71,118 @@ def _ensure_config_plugins_loaded() -> None: def load_config(config_path: str, cli_args: Any = None) -> TrainConfig: """Load a ``TrainConfig`` from a YAML file.""" - _ensure_config_plugins_loaded() - with open(config_path, "r") as f: raw = yaml.load(f, Loader=_TrainConfigSafeLoader) if raw is None: raise ValueError(f"Config file is empty: {config_path}") + config = _build_train_config(raw, source=config_path, cli_args=cli_args) + + # Register the active data backend now that cfg is fully resolved + # (post-CLI-override, post-validate). This is the single source of + # truth for "which dataset backend is this run using". + from wall_x.data._registry import _set_data_backend + + _set_data_backend(config.data.dataset_type) + + return config + + +def load_train_config_from_checkpoint( + checkpoint_path: str, + train_config_path: Optional[str] = None, +) -> TrainConfig: + """The single inference-side entry point for loading a train config. + + Locates the yaml (explicit ``train_config_path`` wins, else + ``/config.yml(.yaml)``), builds the typed ``TrainConfig`` through the + same pipeline as training, applies checkpoint-local processor/tokenizer + overlays, and registers the data backend. + """ + yml_path = train_config_path + if yml_path is None: + for fname in ("config.yml", "config.yaml"): + candidate = os.path.join(checkpoint_path, fname) + if os.path.exists(candidate): + yml_path = candidate + break + if yml_path is None: + raise FileNotFoundError( + f"No config.yml/config.yaml in {checkpoint_path}; " + "pass train_config_path explicitly." + ) + elif not os.path.exists(yml_path): + raise FileNotFoundError(f"train_config_path not found: {yml_path}") + + with open(yml_path, "r") as f: + raw = yaml.load(f, Loader=_TrainConfigSafeLoader) + if raw is None: + raise ValueError(f"Config file is empty: {yml_path}") + + config = _build_train_config(raw, source=yml_path) + _apply_checkpoint_overlays(config, checkpoint_path) + + from wall_x.data._registry import _set_data_backend + + _set_data_backend(config.data.dataset_type) + + return config + + +def _apply_checkpoint_overlays(config: TrainConfig, checkpoint_path: str) -> None: + """Point processor/tokenizer paths at checkpoint-local copies so the + checkpoint stays portable when the training-machine paths are gone. + No-op for model configs without these fields (VGA/pi0).""" + if not os.path.isdir(checkpoint_path): + return + + model = config.model + if hasattr(model, "processor_path") and os.path.exists( + os.path.join(checkpoint_path, "preprocessor_config.json") + ): + model.processor_path = checkpoint_path + + orig_tok = getattr(model, "action_tokenizer_path", None) + if orig_tok is not None and not os.path.exists(orig_tok): + if os.path.exists( + os.path.join(checkpoint_path, "tokenizer.json") + ) and os.path.exists(os.path.join(checkpoint_path, "tokenizer_config.json")): + model.action_tokenizer_path = checkpoint_path + + +def _build_train_config(raw: dict, *, source: str, cli_args: Any = None) -> TrainConfig: + """Build a validated ``TrainConfig`` from a typed-schema raw dict. + Data backend registration is left to the public entry points.""" + _ensure_config_plugins_loaded() + model_type = raw.get("model_type") if model_type is None: - raise ValueError(f"Missing required field 'model_type' in {config_path}") + raise ValueError(f"Missing required field 'model_type' in {source}") + + hyperparams = _build_hyperparams( + raw.get("hyperparams", {}), + model_type=model_type, + ) + distributed = _build_distributed_config( + raw.get("distributed", {}), + optimizer_type=hyperparams.optimizer.optimizer_type, + model_type=model_type, + ) config = TrainConfig( model_type=model_type, - task=_build_dataclass(TaskConfig, raw.get("task", {})), + task=_build_dataclass( + TaskConfig, raw.get("task", {}), where=f"task: section of {source}" + ), model=_build_model_config(model_type, raw.get("model", {})), data=_build_data_config(raw.get("data", {})), - hyperparams=_build_hyperparams(raw.get("hyperparams", {})), - distributed=_build_dataclass(DistributedConfig, raw.get("distributed", {})), + hyperparams=hyperparams, + distributed=distributed, logging=_build_dataclass(LoggingConfig, raw.get("logging", {})), checkpoint=_build_dataclass(CheckpointConfig, raw.get("checkpoint", {})), debug=_build_dataclass(DebugConfig, raw.get("debug", {})), + serving=_build_dataclass(ServingConfig, raw.get("serving", {})), _raw_data=raw.get("data", {}), _raw_yaml=raw, dataset_path=raw.get("dataset_path"), @@ -100,14 +193,48 @@ def load_config(config_path: str, cli_args: Any = None) -> TrainConfig: _validate(config) - # Register the active data backend now that cfg is fully resolved - # (post-CLI-override, post-validate). This is the single source of - # truth for "which dataset backend is this run using". - from wall_x.data._registry import _set_data_backend + return config - _set_data_backend(config.data.dataset_type) - return config +def dump_config(config: TrainConfig) -> dict: + """Serialize the effective ``TrainConfig`` back into the typed yaml schema. + + Inverse of :func:`load_config`: a yaml-safe dict in the same schema as + hand-written configs, no ``_raw_*`` private fields, load/dump roundtrip + stable. Checkpoints persist this form -- the *effective* config (defaults + materialized, CLI overrides applied) -- so reloads never re-derive values + from future defaults. + """ + distributed = dataclasses.asdict(config.distributed) + # These fields double as legacy DDP aliases and would warn on every + # reload; omit them when they equal the default (non-default values are + # kept so the roundtrip stays lossless). + _dist_defaults = { + f.name: f.default for f in dataclasses.fields(type(config.distributed)) + } + for legacy_alias in ( + "find_unused_parameters", + "broadcast_buffers", + "bucket_cap_mb", + ): + if distributed.get(legacy_alias) == _dist_defaults.get(legacy_alias): + distributed.pop(legacy_alias, None) + + out = { + "model_type": config.model_type, + "task": dataclasses.asdict(config.task), + "model": dataclasses.asdict(config.model), + "data": dataclasses.asdict(config.data), + "hyperparams": dataclasses.asdict(config.hyperparams), + "distributed": distributed, + "logging": dataclasses.asdict(config.logging), + "checkpoint": dataclasses.asdict(config.checkpoint), + "debug": dataclasses.asdict(config.debug), + "serving": dataclasses.asdict(config.serving), + } + if config.dataset_path is not None: + out["dataset_path"] = config.dataset_path + return _sanitize_for_yaml(out) def save_config(config: TrainConfig, save_dir: str) -> str: @@ -115,10 +242,13 @@ def save_config(config: TrainConfig, save_dir: str) -> str: os.makedirs(save_dir, exist_ok=True) config_path = os.path.join(save_dir, "config.yml") - data = _sanitize_for_yaml(dataclasses.asdict(config)) with open(config_path, "w") as f: - yaml.dump( - data, f, default_flow_style=False, allow_unicode=True, sort_keys=False + yaml.safe_dump( + dump_config(config), + f, + default_flow_style=False, + allow_unicode=True, + sort_keys=False, ) dataset_config_path = getattr(config.data, "dataset_config_path", None) @@ -130,19 +260,35 @@ def save_config(config: TrainConfig, save_dir: str) -> str: def _sanitize_for_yaml(obj: Any) -> Any: - """Convert dataclass output to YAML-safe containers.""" + """Convert dataclass/config output to YAML-safe containers so it can be + written with ``yaml.safe_dump`` (no ``!!python/*`` tags). tuples/sets -> + lists; numpy scalars -> python scalars; any other non-primitive leaf -> str + (config consumers read these as plain scalars/sequences).""" if isinstance(obj, dict): - return {k: _sanitize_for_yaml(v) for k, v in obj.items()} - elif isinstance(obj, (list, tuple)): + return { + (k if isinstance(k, str) else _sanitize_for_yaml(k)): _sanitize_for_yaml(v) + for k, v in obj.items() + } + if isinstance(obj, (list, tuple, set)): return [_sanitize_for_yaml(v) for v in obj] - return obj - - -def _build_dataclass(cls: Type[T], raw: dict) -> T: - """Build a dataclass from a dict, ignoring unknown keys.""" - if not raw: - return cls() - + if isinstance(obj, (str, bool, int, float)) or obj is None: + return obj + item = getattr(obj, "item", None) # e.g. numpy scalar -> python scalar + if callable(item): + try: + return obj.item() + except Exception: + pass + return str(obj) + + +def _build_dataclass(cls: Type[T], raw: dict, *, where: Optional[str] = None) -> T: + """Build a dataclass from a dict, ignoring unknown keys. + + Fields declared without a default are required: missing ones raise a + ValueError instead of letting the dataclass default paper over them. + """ + raw = raw or {} field_names = {f.name for f in dataclasses.fields(cls)} field_types = {f.name: f.type for f in dataclasses.fields(cls)} filtered = {} @@ -160,9 +306,201 @@ def _build_dataclass(cls: Type[T], raw: dict) -> T: else: filtered[k] = v + missing = [ + f.name + for f in dataclasses.fields(cls) + if f.default is dataclasses.MISSING + and f.default_factory is dataclasses.MISSING + and f.name not in filtered + ] + if missing: + loc = f" ({where})" if where else "" + raise ValueError(f"{cls.__name__} is missing required keys {missing}{loc}") + return cls(**filtered) +def _as_bool(value: Any, *, field: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, int): + return bool(value) + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + raise ValueError(f"distributed.{field} must be a boolean, got {value!r}") + + +def _warn_distributed_alias(old: str, new: str) -> None: + warnings.warn( + "[config deprecation] " + f"distributed.{old} is deprecated; use distributed.{new}. " + "The alias will be removed in the next Wall-X config cleanup release.", + FutureWarning, + stacklevel=3, + ) + + +def _set_nested_alias( + raw: dict, + *, + old_key: str, + section: str, + new_key: str, +) -> None: + if old_key not in raw: + return + _warn_distributed_alias(old_key, f"{section}.{new_key}") + nested = raw.setdefault(section, {}) + if not isinstance(nested, dict): + raise ValueError(f"distributed.{section} must be a mapping") + value = raw[old_key] + if new_key in nested and nested[new_key] != value: + raise ValueError( + f"distributed.{old_key} conflicts with distributed." f"{section}.{new_key}" + ) + nested[new_key] = value + + +def _normalize_distributed_config( + raw: dict | None, + *, + optimizer_type: str | None = None, +) -> dict: + raw = dict(raw or {}) + backend = raw.get("backend") + if backend is not None: + backend = str(backend) + + dmuon_strategy = raw.get("dmuon_strategy") + if dmuon_strategy is not None: + strategy = str(dmuon_strategy) + if strategy == "pure": + _warn_distributed_alias("dmuon_strategy: pure", "backend: dmuon") + if backend is not None and backend != "dmuon": + raise ValueError( + "distributed.dmuon_strategy=pure conflicts with " + f"distributed.backend={backend!r}" + ) + backend = "dmuon" + elif strategy == "fsdp": + raise ValueError( + "distributed.dmuon_strategy=fsdp has been removed. " + "Use distributed.backend=dmuon for pure DMuon, or use a " + "non-DMuon optimizer with distributed.backend=fsdp2." + ) + else: + raise ValueError( + "distributed.dmuon_strategy only accepted legacy value 'pure'; " + f"got {strategy!r}" + ) + + if backend is None and optimizer_type == "dmuon": + _warn_distributed_alias( + "implicit backend from optimizer_type=dmuon", + "backend: dmuon", + ) + backend = "dmuon" + + if "use_fsdp" in raw: + _warn_distributed_alias("use_fsdp", "backend") + if not _as_bool(raw["use_fsdp"], field="use_fsdp"): + if backend is not None and backend != "ddp": + raise ValueError( + "distributed.use_fsdp=false conflicts with " + f"distributed.backend={backend!r}" + ) + backend = "ddp" + elif backend is None: + backend = "fsdp2" + + if backend is None: + backend = "fsdp2" + if backend not in {"fsdp2", "dmuon", "ddp"}: + raise ValueError( + f"Unknown distributed.backend={backend!r}; expected 'fsdp2', " + "'dmuon', or 'ddp'." + ) + raw["backend"] = backend + + active_section = "dmuon" if backend == "dmuon" else "fsdp" + if backend != "ddp": + _set_nested_alias( + raw, + old_key="fsdp_sharding_strategy", + section=active_section, + new_key="sharding_strategy", + ) + _set_nested_alias( + raw, + old_key="fsdp_hsdp_replicate_size", + section=active_section, + new_key="hsdp_replicate_size", + ) + + fsdp_aliases = { + "fsdp_backward_prefetch": "backward_prefetch", + "fsdp_cpu_offload": "cpu_offload", + "fsdp_use_orig_params": "use_orig_params", + "fsdp_limit_all_gathers": "limit_all_gathers", + "fsdp_forward_prefetch": "forward_prefetch", + "fsdp_sync_module_states": "sync_module_states", + "fsdp_save_policy": "save_policy", + "fsdp_reduce_dtype": "reduce_dtype", + } + for old_key, new_key in fsdp_aliases.items(): + if old_key in raw: + if backend == "dmuon": + _warn_distributed_alias( + old_key, + f"fsdp.{new_key} (ignored by backend=dmuon)", + ) + continue + if backend == "ddp": + _warn_distributed_alias( + old_key, + f"fsdp.{new_key} (ignored by backend=ddp)", + ) + continue + _set_nested_alias(raw, old_key=old_key, section="fsdp", new_key=new_key) + + ddp_aliases = { + "find_unused_parameters": "find_unused_parameters", + "broadcast_buffers": "broadcast_buffers", + "bucket_cap_mb": "bucket_cap_mb", + } + for old_key, new_key in ddp_aliases.items(): + if old_key in raw: + if backend != "ddp": + _warn_distributed_alias( + old_key, + f"ddp.{new_key} (ignored by backend={backend})", + ) + continue + _set_nested_alias(raw, old_key=old_key, section="ddp", new_key=new_key) + + return raw + + +def _build_distributed_config( + raw: dict | None, + *, + optimizer_type: str | None = None, + model_type: str | None = None, +) -> DistributedConfig: + # PI0/PI05 carry extra distributed knobs (vision-tower ckpt / non-layer + # replicate / joint0 prefetch) in a subclass so other models' schema stays + # clean (MR !556 review). build_strategy forwards whatever keys are present. + cls = Pi0DistributedConfig if model_type in ("pi0", "pi05") else DistributedConfig + return _build_dataclass( + cls, + _normalize_distributed_config(raw, optimizer_type=optimizer_type), + ) + + def _build_model_config(model_type: str, raw: dict) -> ModelConfig: """Build the registered model config for ``model_type``.""" cls = get_model_config(model_type) @@ -234,10 +572,32 @@ def _build_scheduler_config(raw: dict) -> SchedulerConfig: return _build_dataclass(cls, raw) -def _build_hyperparams(raw: dict) -> TrainHyperParams: - """Build ``TrainHyperParams`` with polymorphic optimizer/scheduler config.""" +def _resolve_hyperparams_cls(algorithm: str, model_type: str): + """Resolve the TrainHyperParams subclass for ``algorithm``. + + Only ``algorithm="sl"`` (default, supervised training) is supported in + this release; any other value is a config error. + """ + if algorithm and algorithm != "sl": + raise ValueError( + f"hyperparams.algorithm={algorithm!r} is not available in this " + f"release; only 'sl' (supervised training) is supported." + ) + return TrainHyperParams + + +def _build_hyperparams(raw: dict, model_type: str = "") -> TrainHyperParams: + """Build ``TrainHyperParams`` (incl. RL subclasses) with polymorphic + optimizer/scheduler config + (algorithm, model_type) routing. + + The discriminator ``hyperparams.algorithm`` lives in the section it + multiplexes (symmetric with ``model_type`` ↔ ``model``); it is read from + the raw dict before the typed object is built. + """ + algorithm = (raw or {}).get("algorithm", "sl") + cls = _resolve_hyperparams_cls(algorithm, model_type) if not raw: - return TrainHyperParams() + return cls() raw = dict(raw) # shallow copy to avoid mutating caller's dict optimizer_raw = raw.pop("optimizer", {}) @@ -246,7 +606,7 @@ def _build_hyperparams(raw: dict) -> TrainHyperParams: optimizer = _build_optimizer_config(optimizer_raw) scheduler = _build_scheduler_config(scheduler_raw) - hp = _build_dataclass(TrainHyperParams, raw) + hp = _build_dataclass(cls, raw) hp.optimizer = optimizer hp.scheduler = scheduler return hp @@ -255,7 +615,10 @@ def _build_hyperparams(raw: dict) -> TrainHyperParams: def _apply_cli_overrides(config: TrainConfig, args: Any) -> None: """Apply CLI overrides to the loaded config.""" if getattr(args, "fsdp_sharding_strategy", None) is not None: - config.distributed.fsdp_sharding_strategy = args.fsdp_sharding_strategy + if config.distributed.backend == "dmuon": + config.distributed.dmuon.sharding_strategy = args.fsdp_sharding_strategy + elif config.distributed.backend == "fsdp2": + config.distributed.fsdp.sharding_strategy = args.fsdp_sharding_strategy if getattr(args, "debug", False): config.logging.log_name = "debug" @@ -269,11 +632,84 @@ def _apply_cli_overrides(config: TrainConfig, args: Any) -> None: config.logging.wandb_offline = args.wandb_offline in ("true", "True", "1") +def _validate_sp(config: TrainConfig) -> None: + """Validate the sequence-parallel (distributed.sp_size) constraints.""" + sp_size = getattr(config.distributed, "sp_size", 1) + if sp_size < 1: + raise ValueError(f"distributed.sp_size must be >= 1, got {sp_size}") + if sp_size == 1: + return + backend = config.distributed.backend + if backend != "fsdp2": + raise ValueError( + "distributed.sp_size>1 (sequence parallel) requires backend=fsdp2, " + f"got backend={backend!r}." + ) + if config.distributed.fsdp.sharding_strategy in { + "hybrid_shard", + "_hybrid_shard_zero2", + }: + raise ValueError( + "distributed.sp_size>1 is not supported with HSDP; " + "use full_shard or shard_grad_op." + ) + m = config.model + cross_view_prob = getattr(m, "cross_view_masking_prob", 0.0) or 0.0 + if getattr(m, "use_sightcone_matching", False) or cross_view_prob > 0: + raise ValueError( + "distributed.sp_size>1 is mutually exclusive with sightcone/" + "cross-view masking; set sp_size=1 or disable masking." + ) + + def _validate(config: TrainConfig) -> None: """Validate required fields and model-specific config.""" if not config.model_type: raise ValueError("model_type is required") + backend = config.distributed.backend + if backend not in {"fsdp2", "dmuon", "ddp"}: + raise ValueError( + "distributed.backend must be 'fsdp2', 'dmuon', or 'ddp', " + f"got {backend!r}" + ) + optimizer_type = config.hyperparams.optimizer.optimizer_type + if optimizer_type == "naive_fsdp2_muon" and backend != "fsdp2": + raise ValueError( + "optimizer_type=naive_fsdp2_muon requires distributed.backend=fsdp2." + ) + if backend == "dmuon" and optimizer_type != "dmuon": + raise ValueError( + "distributed.backend=dmuon requires " + "hyperparams.optimizer.optimizer_type=dmuon." + ) + if backend == "fsdp2" and optimizer_type == "dmuon": + raise ValueError( + "FSDP2 + DMuon dedicated params is no longer supported. " + "Use distributed.backend=dmuon for pure DMuon." + ) + if backend == "ddp" and optimizer_type == "dmuon": + raise ValueError( + "distributed.backend=ddp is not compatible with optimizer_type=dmuon. " + "Use distributed.backend=dmuon for DMuon." + ) + valid_sharding = { + "full_shard", + "shard_grad_op", + "hybrid_shard", + "_hybrid_shard_zero2", + } + if backend != "ddp": + backend_section = "dmuon" if backend == "dmuon" else "fsdp" + sharding = getattr(config.distributed, backend_section).sharding_strategy + if sharding not in valid_sharding: + raise ValueError( + f"distributed.{backend_section}.sharding_strategy={sharding!r} is invalid; " + f"expected one of {sorted(valid_sharding)}." + ) + + _validate_sp(config) + if not config.task.dof_config: raise ValueError("task.dof_config is required (cannot be empty)") diff --git a/wall_x/config/registry.py b/wall_x/config/registry.py index d34debfd..3ecfdb25 100644 --- a/wall_x/config/registry.py +++ b/wall_x/config/registry.py @@ -11,6 +11,7 @@ from collections.abc import Callable from typing import Any, TypeVar + T = TypeVar("T") _DATA_CONFIG_REGISTRY: dict[str, type[Any]] = {} diff --git a/wall_x/config/task_config.py b/wall_x/config/task_config.py index 1cd02b82..f1be3987 100644 --- a/wall_x/config/task_config.py +++ b/wall_x/config/task_config.py @@ -1,6 +1,6 @@ """Task config shared by model and data code.""" -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any, Dict, Optional @@ -8,10 +8,10 @@ class TaskConfig: """Robot task definition shared by model construction and data slicing.""" - dof_config: Dict[str, int] = field(default_factory=dict) - agent_pos_config: Dict[str, int] = field(default_factory=dict) + dof_config: Dict[str, int] + agent_pos_config: Dict[str, int] + action_horizon: int ar_dof_config: Optional[Dict[str, int]] = None - action_horizon: int = 32 action_horizon_flow: int = 32 noise_scheduler: Optional[Dict[str, Any]] = None use_state_string_representation: bool = False diff --git a/wall_x/config/train_config.py b/wall_x/config/train_config.py index 979622e3..7f335159 100644 --- a/wall_x/config/train_config.py +++ b/wall_x/config/train_config.py @@ -11,6 +11,7 @@ DebugConfig, DistributedConfig, LoggingConfig, + ServingConfig, ) from .model_config import ModelConfig from .task_config import TaskConfig @@ -20,8 +21,8 @@ class TrainConfig: """Top-level Wall-X training config.""" - model_type: str = "qwen2_5" - task: TaskConfig = field(default_factory=TaskConfig) + model_type: str + task: TaskConfig model: ModelConfig = field(default_factory=ModelConfig) data: DataConfig = field(default_factory=DataConfig) hyperparams: TrainHyperParams = field(default_factory=TrainHyperParams) @@ -29,6 +30,7 @@ class TrainConfig: logging: LoggingConfig = field(default_factory=LoggingConfig) checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig) debug: DebugConfig = field(default_factory=DebugConfig) + serving: ServingConfig = field(default_factory=ServingConfig) # Raw YAML sections preserved verbatim for backend APIs that read fields # not captured by the typed DataConfig dataclass. _raw_data: Dict[str, Any] = field(default_factory=dict) diff --git a/wall_x/data/__init__.py b/wall_x/data/__init__.py index b4661af3..bd53c4f0 100644 --- a/wall_x/data/__init__.py +++ b/wall_x/data/__init__.py @@ -1,19 +1,19 @@ """Public data facade for backend selection and dataset construction.""" -from wall_x.data import backends # noqa: F401 (side-effect registration) from wall_x.data._bundle import DataBundle from wall_x.data._protocol import BuildContext, DatasetBackend from wall_x.data._registry import ( MissingOperationError, + data_backend, available_backends, backend_for, build_data, - data_backend, has_data_backend, register, register_module, resolve_dataset_type, ) +from wall_x.data import backends # noqa: F401 (side-effect registration) __all__ = [ "DataBundle", diff --git a/wall_x/data/_protocol.py b/wall_x/data/_protocol.py index 50694eb6..bc6abed5 100644 --- a/wall_x/data/_protocol.py +++ b/wall_x/data/_protocol.py @@ -21,6 +21,8 @@ class BuildContext: tokenizer: Optional[Any] = None processor: Optional[Any] = None tokenizer_mixin: Optional[Any] = None + train_action_tokenizer: Optional[Any] = None + val_action_tokenizer: Optional[Any] = None normalizer_action: Optional[Any] = None normalizer_propri: Optional[Any] = None model_config: Optional[Any] = None diff --git a/wall_x/data/_registry.py b/wall_x/data/_registry.py index 23bfe017..5e036e46 100644 --- a/wall_x/data/_registry.py +++ b/wall_x/data/_registry.py @@ -96,7 +96,7 @@ def __init__(self, backend_name: str, op_name: str) -> None: hint = f"Supported by: {impls}." else: hint = ( - f"No registered backend implements {op_name!r} - check the " + f"No registered backend implements {op_name!r} — check the " f"spelling or add it to the backend module." ) super().__init__( @@ -179,7 +179,7 @@ def resolve_dataset_type(cfg_or_yaml: Any, default: str = "lerobot") -> str: def build_data(cfg: Any, ctx: BuildContext) -> DataBundle: """Dispatch ``backend.build(cfg, ctx)`` to the backend named by cfg. - ``cfg`` is expected to be a typed :class:`TrainConfig` - that's what + ``cfg`` is expected to be a typed :class:`TrainConfig` — that's what ``wall_x.config.load_config`` returns and what every trainer entry point passes in. Raw yaml dicts are not supported here on purpose: the typed schema is the contract that gives backends their @@ -198,10 +198,10 @@ def build_data(cfg: Any, ctx: BuildContext) -> DataBundle: def _set_data_backend(name: str) -> None: - """Internal - only config loaders should call this. + """Internal — only config loaders should call this. Strict semantics: first call wins; same-value re-set is a no-op; - different-value re-set raises. The error is loud on purpose - it means + different-value re-set raises. The error is loud on purpose — it means two config loaders disagreed about which backend to use, which is almost always a bug (e.g. business code calling this directly, or two yamls being loaded in the same process). @@ -222,7 +222,7 @@ def _set_data_backend(name: str) -> None: def _reset_data_backend() -> None: - """Internal - clear the active backend. Tests only.""" + """Internal — clear the active backend. Tests only.""" global _DATA_BACKEND _DATA_BACKEND = None @@ -236,7 +236,7 @@ def data_backend() -> _BackendProxy: """Return the active backend proxy. Raises if no config has been loaded. Direct construction of ``TrainConfig(...)`` does **not** register a - backend - that's intentional. Business code must go through + backend — that's intentional. Business code must go through :func:`wall_x.config.load_config`, which registers the backend before returning. """ diff --git a/wall_x/data/backends/lerobot/build.py b/wall_x/data/backends/lerobot/build.py index 36edd4fb..493ae488 100644 --- a/wall_x/data/backends/lerobot/build.py +++ b/wall_x/data/backends/lerobot/build.py @@ -40,12 +40,7 @@ def load_lerobot_normalizers(cfg: Any): ) dataset_name = str(lerobot_config["repo_id"]) - norm_stats = load_norm_stats( - norm_stats_path, - key_mappings, - dof_config=dict(cfg.task.dof_config or {}), - agent_pos_config=dict(cfg.task.agent_pos_config or {}), - ) + norm_stats = load_norm_stats(norm_stats_path, key_mappings) normalizer_action, normalizer_propri = create_normalizers_from_lerobot_norm_stats( norm_stats, dataset_name, @@ -105,8 +100,8 @@ def get_val_dataloader(self): # train and val splits (flipped by ``_train()`` / ``_eval()``). # The val DataLoader's DistributedSampler caches total_size sized # to the val split but ``__iter__`` reads ``len(self.dataset)`` - # live - if a preceding train_loop left the pointer at train, that - # live len is ~20x total_size and DistributedSampler asserts. + # live — if a preceding train_loop left the pointer at train, that + # live len is ~20× total_size and DistributedSampler asserts. # Rebuild each time so ``_eval()`` runs and a fresh sampler is # snapped to the current (val) split length. Mirrors the train-side # rebuild-on-every-epoch pattern. @@ -155,7 +150,7 @@ def load_trainer_data_config_from_yaml_dict(yaml_dict: Dict[str, Any]) -> Lerobo def _build_flat_config(cfg: Any) -> Tuple[Dict[str, Any], Dict[str, Any]]: - """Map typed TrainConfig -> (flat_config, lerobot_config) for legacy entry. + """Map typed TrainConfig → (flat_config, lerobot_config) for legacy entry. ``load_lerobot_data`` expects a 2509-style flat dict plus a separate ``lerobot_config`` carrying ``repo_id`` / ``root``. This function is @@ -167,7 +162,6 @@ def _build_flat_config(cfg: Any) -> Tuple[Dict[str, Any], Dict[str, Any]]: data = cfg.data hp = cfg.hyperparams raw = getattr(cfg, "_raw_yaml", {}) or {} - raw_data = dict(getattr(cfg, "_raw_data", {}) or {}) lerobot_cfg = dict(data.lerobot_config or {}) if "repo_id" not in lerobot_cfg: @@ -183,27 +177,17 @@ def _build_flat_config(cfg: Any) -> Tuple[Dict[str, Any], Dict[str, Any]]: "seed": hp.seed, "resolution": data.resolution, } - if raw_data.get("max_length") is not None: - data_section["max_length"] = raw_data["max_length"] if data.priority_order is not None: data_section["priority_order"] = data.priority_order if data.camera_name_mapping is not None: data_section["camera_name_mapping"] = data.camera_name_mapping - data_section.setdefault( - "use_state_string_representation", - cfg.task.use_state_string_representation, - ) - data_section.setdefault( - "state_bins", - raw_data.get("state_bins", raw.get("state_bins", 256)), - ) # Dof/agent_pos totals for the collator's zero-pad step. When resuming - # from a checkpoint trained on a larger action space, task.dof_config - # should include an ``action_padding`` key that absorbs the diff; the - # collator right-pads action/agent_pos tensors to these totals with - # dof_mask/agent_pos_mask zeroed on padded dims so loss doesn't flow - # through them. + # from a checkpoint trained on a larger action space (e.g. bus2602 26-dim + # dof while libero data is 7-dim), task.dof_config should include an + # ``action_padding`` key that absorbs the diff; the collator right-pads + # action/agent_pos tensors to these totals with dof_mask/agent_pos_mask + # zeroed on padded dims so loss doesn't flow through them. dof_total = int(sum((cfg.task.dof_config or {}).values())) agent_pos_total = int(sum((cfg.task.agent_pos_config or {}).values())) @@ -222,15 +206,6 @@ def _build_flat_config(cfg: Any) -> Tuple[Dict[str, Any], Dict[str, Any]]: "noise_scheduler": data.noise_scheduler or {}, "dof_total_dim": dof_total, "agent_pos_total_dim": agent_pos_total, - "dof_config": dict(cfg.task.dof_config or {}), - "agent_pos_config": dict(cfg.task.agent_pos_config or {}), - "use_state_string_representation": cfg.task.use_state_string_representation, - "state_bins": int( - raw_data.get("state_bins") - or data_section.get("state_bins") - or raw.get("state_bins") - or 256 - ), "data": data_section, } return flat, lerobot_cfg @@ -244,7 +219,7 @@ def load_lerobot_v2( The third return value ``train_num`` is a snapshot of ``len(train_dataloader)`` at construction time. It matches ``wrapper.global_train_iters.value`` initially but does not track - subsequent rebuilds inside ``set_epoch`` - callers doing dynamic + subsequent rebuilds inside ``set_epoch`` — callers doing dynamic resampling should read from the mp.Value, not from this snapshot. """ from wall_x.data.backends.lerobot.loader import load_lerobot_data @@ -269,7 +244,7 @@ def load_lerobot_v2( # PreprocessedDataset.get_*_dataloader returns (dataloader, sampler). # Build val first, train second, so the inner ``_dataset`` pointer is - # left at the train split when we finish - workers fork from that + # left at the train split when we finish — workers fork from that # state on first iteration. val_dataloader, _ = inner.get_val_dataloader() val_num = len(val_dataloader) if val_dataloader is not None else 0 @@ -308,7 +283,7 @@ def load_lerobot_v2( def build(cfg, ctx): - """Backend Protocol entry - returns a ``DataBundle``. + """Backend Protocol entry — returns a ``DataBundle``. Wraps ``load_lerobot_v2`` (which returns the trainer-facing triple) into the unified ``DataBundle`` shape every backend exposes. diff --git a/wall_x/data/backends/lerobot/config.py b/wall_x/data/backends/lerobot/config.py index 0b79d07c..5a6c8075 100644 --- a/wall_x/data/backends/lerobot/config.py +++ b/wall_x/data/backends/lerobot/config.py @@ -1,7 +1,6 @@ -from dataclasses import dataclass, field from typing import Any, Dict, Optional - -from qwen_vl_utils.vision_process import IMAGE_FACTOR, MAX_PIXELS, MIN_PIXELS +from dataclasses import dataclass, field +from qwen_vl_utils.vision_process import MIN_PIXELS, MAX_PIXELS, IMAGE_FACTOR @dataclass diff --git a/wall_x/data/backends/lerobot/loader.py b/wall_x/data/backends/lerobot/loader.py index c000ec84..26138388 100644 --- a/wall_x/data/backends/lerobot/loader.py +++ b/wall_x/data/backends/lerobot/loader.py @@ -4,78 +4,26 @@ import logging import os -from typing import Protocol, SupportsIndex, TypeVar - import numpy as np import torch +from torch.utils.data import DistributedSampler, random_split from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from typing import Protocol, SupportsIndex, TypeVar from qwen_vl_utils.vision_process import smart_resize -from torch.utils.data import DistributedSampler, random_split -from transformers import AutoProcessor - -from wall_x._vendor.x2robot_utils.geometry import ( - canonicalize_euler_zyx_batch_nb, - euler_to_matrix_zyx_batch_nb, - matrix_to_euler_zyx_batch_nb, - so3_to_matrix_batch_nb, -) from wall_x.data.backends.lerobot.config import LerobotConfig -from wall_x.data.backends.lerobot.rotation_layout import ( - LAYOUT_SKIP_KEYS, - maybe_convert_euler_to_6d, -) -from wall_x.data.backends.lerobot.rotation_layout import ( - euler_layout_dim as _euler_layout_dim, -) -from wall_x.data.backends.lerobot.rotation_layout import ( - layout_uses_6d_rotation as _layout_uses_6d_rotation, -) from wall_x.data.backends.lerobot.utils import ( - get_wallx_normal_text, - load_norm_stats, - preprocesser_call, process_grounding_points, + get_wallx_normal_text, replace_action_token, + preprocesser_call, + load_norm_stats, ) +from transformers import AutoProcessor + T_co = TypeVar("T_co", covariant=True) logger = logging.getLogger(__name__) -RELATIVE_KEYWORD = "relative" -ROTATION_KEYWORD = "rotation" -RELATIVE_SKIP_KEYS = LAYOUT_SKIP_KEYS - - -def _compute_delta_from_state_and_abs_rot( - rotations: np.ndarray, state: np.ndarray -) -> np.ndarray: - """Relative rotation: R_rel = R_abs @ R_state^T.""" - if rotations.shape[-1] == 3: - rotations_matrix = euler_to_matrix_zyx_batch_nb(rotations) - out_is_euler = True - elif rotations.shape[-1] == 6: - rotations_matrix = so3_to_matrix_batch_nb(rotations) - out_is_euler = False - else: - raise ValueError( - f"Only 3D euler or 6D rotation supported, got {rotations.shape[-1]}D" - ) - - if state.shape[-1] == 3: - state_matrix = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0] - elif state.shape[-1] == 6: - state_matrix = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] - else: - raise ValueError( - f"Only 3D euler or 6D rotation supported, got {state.shape[-1]}D" - ) - - r_rel = np.matmul(rotations_matrix, state_matrix.T) - if out_is_euler: - d_euler = matrix_to_euler_zyx_batch_nb(r_rel) - return canonicalize_euler_zyx_batch_nb(d_euler) - return r_rel[:, :2, :].reshape(r_rel.shape[0], 6) - # Abstract class for dataset class Dataset(Protocol[T_co]): @@ -139,86 +87,6 @@ def __init__( self._state_key_mapping = self.key_mappings self._action_key_mapping = self.key_mappings - task_cfg = self.config.get("task") or {} - self._dof_config = self.config.get("dof_config") or task_cfg.get( - "dof_config", {} - ) - self._agent_pos_config = self.config.get("agent_pos_config") or task_cfg.get( - "agent_pos_config", {} - ) - self._use_relative_action = any( - RELATIVE_KEYWORD in key for key in self._dof_config - ) - self._convert_action_euler_to_6d = _layout_uses_6d_rotation(self._dof_config) - self._convert_state_euler_to_6d = _layout_uses_6d_rotation( - self._agent_pos_config - ) - if self._convert_action_euler_to_6d or self._convert_state_euler_to_6d: - logger.info( - "LeRobot loader: Euler->6D rotation enabled " - "(action=%s, state=%s; raw action dim=%s -> %s)", - self._convert_action_euler_to_6d, - self._convert_state_euler_to_6d, - ( - _euler_layout_dim(self._dof_config) - if self._convert_action_euler_to_6d - else "-" - ), - ( - sum( - d - for k, d in self._dof_config.items() - if k not in RELATIVE_SKIP_KEYS - ) - if self._convert_action_euler_to_6d - else "-" - ), - ) - - def _maybe_convert_euler_to_6d(self, vec, layout_config: dict, enabled: bool): - converted = maybe_convert_euler_to_6d(vec, layout_config, enabled) - if ( - enabled - and layout_config - and isinstance(vec, torch.Tensor) - and converted is not vec - ): - return torch.as_tensor(converted, dtype=vec.dtype, device=vec.device) - return converted - - def _to_relative_action(self, action, agent_pos): - """Convert absolute action horizon to deltas w.r.t. current agent_pos.""" - action = np.asarray(action, dtype=np.float64) - agent_pos = np.asarray(agent_pos, dtype=np.float64) - if action.ndim == 1: - action = action[np.newaxis, :] - if agent_pos.ndim > 1: - agent_pos = agent_pos.reshape(-1) - - parts = [] - cur = 0 - for key, dim in self._dof_config.items(): - if key in RELATIVE_SKIP_KEYS: - continue - action_clip = action[:, cur : cur + dim] - agent_pos_clip = agent_pos[cur : cur + dim] - if RELATIVE_KEYWORD not in key: - parts.append(action_clip) - elif ROTATION_KEYWORD in key: - parts.append( - _compute_delta_from_state_and_abs_rot( - action_clip.astype(np.float64), - agent_pos_clip.astype(np.float64), - ) - ) - else: - parts.append(action_clip - agent_pos_clip[np.newaxis, :]) - cur += dim - - if not parts: - return action - return np.concatenate(parts, axis=1).astype(np.float32) - def _vision_preprocess(self, frames): processed_frames = [] for key in self.hf_dataset.meta.camera_keys: @@ -261,19 +129,6 @@ def __getitem__(self, index): image_inputs, h, w, resize_h, resize_w = self._vision_preprocess(data) agent_pos = data[self._state_key_mapping["state"]] action = data[self._action_key_mapping["action"]] - agent_pos = self._maybe_convert_euler_to_6d( - agent_pos, self._agent_pos_config, self._convert_state_euler_to_6d - ) - action = self._maybe_convert_euler_to_6d( - action, self._dof_config, self._convert_action_euler_to_6d - ) - if self._use_relative_action: - device = action.device if isinstance(action, torch.Tensor) else None - action = torch.as_tensor( - self._to_relative_action(action, agent_pos), - dtype=torch.float32, - device=device, - ) frame_index = data["frame_index"] instruction_info = {"instruction": data["task"]} generate_subtask_ratio = self.data_config.generate_subtask_ratio @@ -417,10 +272,6 @@ def __init__(self, config, dataload_config, stats, lerobot_config): self.beta_alpha = float(self.beta_alpha) self.beta_beta = float(self.beta_beta) self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) - self.use_state_string_representation = bool( - self.config.get("use_state_string_representation", False) - ) - self.state_bins = int(self.config.get("state_bins", 256)) self.load_processor() def load_processor(self): @@ -516,7 +367,7 @@ def __call__(self, batch): additional_inputs = {} # Tail-pad widths when dof_config / agent_pos_config (sum) is larger - # than the lerobot action/state - typical when resuming a ckpt that + # than the lerobot action/state — typical when resuming a ckpt that # was pretrained on a bigger action space. Extra columns are filled # with zeros and their mask set to 0 so loss is not propagated. dof_total = int(self.config.get("dof_total_dim", 0) or 0) @@ -622,14 +473,6 @@ def __call__(self, batch): truncation=True, return_tensors="pt", max_length=self.dataload_config.get("max_length", 768), - norm_state=( - additional_inputs["proprioception"] - if self.use_state_string_representation - and "proprioception" in additional_inputs - else None - ), - agent_pos_mask=additional_inputs.get("agent_pos_mask"), - state_bins=self.state_bins, ) action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") @@ -717,17 +560,7 @@ def load_lerobot_data( assert ( norm_stats_path is not None ), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats" - task_cfg = config.get("task") or {} - dof_config = config.get("dof_config") or task_cfg.get("dof_config", {}) - agent_pos_config = config.get("agent_pos_config") or task_cfg.get( - "agent_pos_config", {} - ) - norm_stats = load_norm_stats( - norm_stats_path, - key_mappings, - dof_config=dof_config, - agent_pos_config=agent_pos_config, - ) + norm_stats = load_norm_stats(norm_stats_path, key_mappings) delta_timestamps = { # action chunk diff --git a/wall_x/data/backends/lerobot/rotation_layout.py b/wall_x/data/backends/lerobot/rotation_layout.py deleted file mode 100644 index ba8a780f..00000000 --- a/wall_x/data/backends/lerobot/rotation_layout.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Layout helpers when config expects 6D rotation but LeRobot stores 3D Euler.""" - -from __future__ import annotations - -import numpy as np - -from wall_x._vendor.x2robot_utils.geometry import euler_to_matrix_zyx_6d_nb - -LAYOUT_SKIP_KEYS = frozenset( - {"velocity_decomposed", "height", "head_actions", "action_padding"} -) -ROTATION_KEYWORD = "rotation" -ROTATION_6D_KEYWORD = "6D" - - -def layout_uses_6d_rotation(layout_config: dict) -> bool: - for key, dim in layout_config.items(): - if key in LAYOUT_SKIP_KEYS: - continue - if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6: - return True - return False - - -def euler_layout_dim(layout_config: dict) -> int: - """Vector width in LeRobot when rotation slices are still 3D Euler.""" - total = 0 - for key, dim in layout_config.items(): - if key in LAYOUT_SKIP_KEYS: - continue - if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6: - total += 3 - else: - total += int(dim) - return total - - -def convert_euler_to_6d(vec: np.ndarray, layout_config: dict) -> np.ndarray: - """Rewrite [pos, euler(3), tail...] to [pos, rot6d(6), tail...] per layout.""" - vec = np.asarray(vec, dtype=np.float64) - single = vec.ndim == 1 - if single: - vec = vec[np.newaxis, :] - - out_rows = [] - for row in vec: - parts: list[np.ndarray] = [] - raw_cur = 0 - for key, dim in layout_config.items(): - if key in LAYOUT_SKIP_KEYS: - continue - dim = int(dim) - if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6: - euler = row[raw_cur : raw_cur + 3] - rot6d = euler_to_matrix_zyx_6d_nb(euler.reshape(1, 3)).reshape(6) - parts.append(rot6d) - raw_cur += 3 - else: - parts.append(row[raw_cur : raw_cur + dim]) - raw_cur += dim - out_rows.append(np.concatenate(parts, axis=0)) - - out = np.stack(out_rows, axis=0) - return out[0] if single else out - - -def maybe_convert_norm_stats_vector( - values, - layout_config: dict, - enabled: bool | None = None, -): - """Convert a 1D norm-stat vector (q01/q99/mean/std) from Euler layout to 6D.""" - if enabled is None: - enabled = layout_uses_6d_rotation(layout_config) - if not enabled or not layout_config: - return values - arr = np.asarray(values, dtype=np.float64) - if arr.ndim != 1: - return values - raw_dim = euler_layout_dim(layout_config) - if arr.shape[0] != raw_dim: - return values - return convert_euler_to_6d(arr, layout_config).astype(np.float32) - - -def maybe_convert_euler_to_6d( - vec: np.ndarray, layout_config: dict, enabled: bool -) -> np.ndarray: - if not enabled or not layout_config: - return vec - raw_dim = euler_layout_dim(layout_config) - arr = np.asarray(vec) - if arr.shape[-1] != raw_dim: - return vec - return convert_euler_to_6d(arr, layout_config).astype(np.float32) diff --git a/wall_x/data/backends/lerobot/utils.py b/wall_x/data/backends/lerobot/utils.py index e0028401..7bf3d93d 100644 --- a/wall_x/data/backends/lerobot/utils.py +++ b/wall_x/data/backends/lerobot/utils.py @@ -6,7 +6,6 @@ from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple, Union -import numpy as np import torch from transformers import BatchFeature @@ -20,63 +19,21 @@ class NormStats: delta: torch.Tensor -def load_norm_stats( - norm_stats_path, - key_mappings, - dof_config: dict | None = None, - agent_pos_config: dict | None = None, -): - from wall_x.data.backends.lerobot.rotation_layout import ( - layout_uses_6d_rotation, - maybe_convert_norm_stats_vector, - ) - +def load_norm_stats(norm_stats_path, key_mappings): with open(norm_stats_path, "r") as f: norm_stats = json.load(f) action_key = key_mappings["action"] state_key = key_mappings["state"] - - action_q01 = maybe_convert_norm_stats_vector( - norm_stats["norm_stats"][action_key]["q01"], dof_config or {} - ) - action_q99 = maybe_convert_norm_stats_vector( - norm_stats["norm_stats"][action_key]["q99"], dof_config or {} - ) - if layout_uses_6d_rotation(dof_config or {}) and len(action_q01) != len( - norm_stats["norm_stats"][action_key]["q01"] - ): - logger.info( - "Converted action norm stats Euler->6D (%d -> %d dims)", - len(norm_stats["norm_stats"][action_key]["q01"]), - len(action_q01), - ) - - q01 = torch.tensor(action_q01) - q99 = torch.tensor(action_q99) + q01 = torch.tensor(norm_stats["norm_stats"][action_key]["q01"]) + q99 = torch.tensor(norm_stats["norm_stats"][action_key]["q99"]) delta = q99 - q01 action_norm_stats = NormStats( min=q01, max=q99, delta=delta, ) - - state_q01 = maybe_convert_norm_stats_vector( - norm_stats["norm_stats"][state_key]["q01"], agent_pos_config or {} - ) - state_q99 = maybe_convert_norm_stats_vector( - norm_stats["norm_stats"][state_key]["q99"], agent_pos_config or {} - ) - if layout_uses_6d_rotation(agent_pos_config or {}) and len(state_q01) != len( - norm_stats["norm_stats"][state_key]["q01"] - ): - logger.info( - "Converted state norm stats Euler->6D (%d -> %d dims)", - len(norm_stats["norm_stats"][state_key]["q01"]), - len(state_q01), - ) - - q01 = torch.tensor(state_q01) - q99 = torch.tensor(state_q99) + q01 = torch.tensor(norm_stats["norm_stats"][state_key]["q01"]) + q99 = torch.tensor(norm_stats["norm_stats"][state_key]["q99"]) delta = q99 - q01 state_norm_stats = NormStats( min=q01, @@ -463,10 +420,6 @@ def preprocesser_call( truncation: Optional[bool] = None, max_length: Optional[int] = None, return_tensors: str = "pt", - norm_state=None, - agent_pos_mask=None, - state_bins: int = 256, - state_drop_prob: float = 0.0, ) -> BatchFeature: """Unified preprocessing function for Wall-X model handling text, image and video inputs. @@ -523,7 +476,7 @@ def preprocesser_call( ) else: raise RuntimeError( - "processor has no video_processor attribute - this code path " + "processor has no video_processor attribute — this code path " "requires transformers>=5.2. Upgrade transformers or pass " "videos=None." ) @@ -536,34 +489,6 @@ def preprocesser_call( if not isinstance(text, list): text = [text] - # Discretize normalized proprioception into the <|propri|> prompt slot. - if norm_state is not None: - norm_state = ( - norm_state.cpu().numpy() - if isinstance(norm_state, torch.Tensor) - else norm_state - ) - discretized = ( - np.digitize(norm_state, bins=np.linspace(-1, 1, state_bins + 1)[:-1]) - 1 - ) - discretized = discretized[:, 0, :] - if agent_pos_mask is not None: - mask = ( - agent_pos_mask[:, 0, :].cpu().numpy().astype(bool) - if isinstance(agent_pos_mask, torch.Tensor) - else agent_pos_mask[:, 0, :].astype(bool) - ) - else: - mask = np.ones(discretized.shape, dtype=bool) - for i in range(len(text)): - if "<|propri|>" not in text[i]: - continue - if state_drop_prob > 0 and random.random() < state_drop_prob: - text[i] = text[i].replace("<|propri|>", "") - else: - state_str = " ".join(map(str, discretized[i, mask[i]])) - text[i] = text[i].replace("<|propri|>", state_str) - # Process image placeholder tokens in text if image_grid_thw is not None: merge_length = processor.image_processor.merge_size**2 diff --git a/wall_x/distributed/__init__.py b/wall_x/distributed/__init__.py new file mode 100644 index 00000000..69e2ac74 --- /dev/null +++ b/wall_x/distributed/__init__.py @@ -0,0 +1,24 @@ +# Sequence Parallelism and Data Parallelism Support for wall-x +from .parallel_state import ( + # SP functions + init_sequence_parallel, + get_sp_group, + get_sp_world_size, + get_sp_rank, + is_sp_initialized, + destroy_sp_group, + # DP functions (new) + get_dp_group, + get_dp_world_size, + get_dp_rank, + # Utility functions + is_parallel_initialized, + destroy_parallel_groups, +) +from .communication_op import ( + sp_shard, + sp_gather, + sp_all_to_all_4d, + sp_all_gather_4d, + reduce_sp_partial_grads, +) diff --git a/wall_x/distributed/communication_op.py b/wall_x/distributed/communication_op.py new file mode 100644 index 00000000..f7accb3b --- /dev/null +++ b/wall_x/distributed/communication_op.py @@ -0,0 +1,400 @@ +""" +Sequence Parallelism Communication Operations + +SP communication operations: sharding, gathering, all-to-all, etc. +""" + +import torch +import torch.distributed as dist +from typing import Tuple +from .parallel_state import get_sp_group, get_sp_world_size, get_sp_rank + + +class _Shard(torch.autograd.Function): + """ + Differentiable shard operation. + + Forward: Take only this rank's portion of the input tensor. + Backward: All-gather gradients from all ranks. + """ + + @staticmethod + def forward( + ctx, input_: torch.Tensor, dim: int, padded_seq_len: int + ) -> torch.Tensor: + sp_world_size = get_sp_world_size() + sp_rank = get_sp_rank() + sp_group = get_sp_group() + + ctx.dim = dim + ctx.sp_world_size = sp_world_size + ctx.sp_group = sp_group + ctx.padded_seq_len = padded_seq_len + + # Shard: take only this rank's portion + elements_per_rank = padded_seq_len // sp_world_size + start_idx = sp_rank * elements_per_rank + end_idx = (sp_rank + 1) * elements_per_rank + + slices = [slice(None)] * input_.dim() + slices[dim] = slice(start_idx, end_idx) + return input_[tuple(slices)].contiguous() + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None, None]: + # Backward of shard is all-gather + grad_input = _gather_impl(grad_output, ctx.dim, ctx.sp_world_size, ctx.sp_group) + return grad_input, None, None + + +def sp_shard( + input_: torch.Tensor, + dim: int = 1, +) -> Tuple[torch.Tensor, int]: + """ + Differentiable shard of the input tensor along the specified dimension across SP ranks. + + Args: + input_: Input tensor, shape like [B, S, D] + dim: Dimension to shard, default 1 (sequence dimension) + + Returns: + tuple: (sharded tensor, original sequence length) + + Example: + # SP=4, input [B, 1000, D] + # output [B, 256, D] (padded to 1024, then 256 per rank) + """ + sp_world_size = get_sp_world_size() + + if sp_world_size == 1: + return input_, input_.shape[dim] + + original_seq_len = input_.shape[dim] + + # Calculate padding (ensure divisible by sp_size) + remainder = original_seq_len % sp_world_size + if remainder != 0: + padding_amount = sp_world_size - remainder + padded_seq_len = original_seq_len + padding_amount + + # Create padding (this is differentiable - just concat with zeros) + pad_shape = list(input_.shape) + pad_shape[dim] = padding_amount + padding = torch.zeros(pad_shape, dtype=input_.dtype, device=input_.device) + input_ = torch.cat([input_, padding], dim=dim) + else: + padded_seq_len = original_seq_len + + # Differentiable shard + output = _Shard.apply(input_, dim, padded_seq_len) + + return output, original_seq_len + + +def _gather_impl( + input_: torch.Tensor, + dim: int, + sp_world_size: int, + sp_group, +) -> torch.Tensor: + """ + Internal implementation of all-gather (non-differentiable). + """ + gathered_list = [torch.empty_like(input_) for _ in range(sp_world_size)] + dist.all_gather(gathered_list, input_, group=sp_group) + return torch.cat(gathered_list, dim=dim) + + +def _split_impl( + input_: torch.Tensor, + dim: int, + sp_world_size: int, + sp_rank: int, +) -> torch.Tensor: + """ + Internal implementation of split/scatter (non-differentiable). + Takes only the slice corresponding to the current rank. + """ + dim_size = input_.shape[dim] + assert ( + dim_size % sp_world_size == 0 + ), f"dim_size ({dim_size}) must be divisible by sp_size ({sp_world_size})" + + chunk_size = dim_size // sp_world_size + start_idx = sp_rank * chunk_size + end_idx = (sp_rank + 1) * chunk_size + + slices = [slice(None)] * input_.dim() + slices[dim] = slice(start_idx, end_idx) + return input_[tuple(slices)].contiguous() + + +class _AllGather(torch.autograd.Function): + """ + Differentiable all-gather operation. + + Forward: Gathers shards from all ranks along specified dimension. + Backward: Splits gradient and returns only this rank's portion. + """ + + @staticmethod + def forward(ctx, input_: torch.Tensor, dim: int) -> torch.Tensor: + sp_world_size = get_sp_world_size() + sp_group = get_sp_group() + sp_rank = get_sp_rank() + ctx.dim = dim + ctx.sp_world_size = sp_world_size + ctx.sp_rank = sp_rank + return _gather_impl(input_, dim, sp_world_size, sp_group) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]: + # Backward of all-gather is split: take only this rank's portion + grad_input = _split_impl(grad_output, ctx.dim, ctx.sp_world_size, ctx.sp_rank) + return grad_input, None + + +def sp_gather( + input_: torch.Tensor, + original_seq_len: int, + dim: int = 1, +) -> torch.Tensor: + """ + Differentiable gather of tensors from all SP ranks and remove padding. + + Args: + input_: Sharded tensor + original_seq_len: Original sequence length + dim: Gather dimension + + Returns: + Gathered tensor + """ + sp_world_size = get_sp_world_size() + + if sp_world_size == 1: + return input_ + + # Differentiable all-gather + gathered = _AllGather.apply(input_, dim) + + # Remove padding (this is differentiable - just slicing) + if gathered.shape[dim] > original_seq_len: + slices = [slice(None)] * gathered.dim() + slices[dim] = slice(0, original_seq_len) + gathered = gathered[tuple(slices)].contiguous() + + return gathered + + +def reduce_sp_partial_grads(named_parameters, is_sharded, sp_group=None) -> int: + """Sum partial grads of SP-sharded params across the sequence-parallel group. + + Under sequence parallelism (sp_size>1) the params selected by ``is_sharded`` + run on this rank's sequence shard, so their ``.grad`` is partial and must be + summed over the SP group. Params trained on the gathered full sequence hold + the identical full grad on every SP rank and must be excluded (summing would + over-count them by sp_size). The orthogonal DP reduction stays with FSDP's + reduce-scatter. The many small block-grad shards are coalesced into one flat + all-reduce per dtype (~1 collective) instead of one NCCL launch per param. + + Args: + named_parameters: iterable of (name, param), e.g. model.named_parameters(). + is_sharded: predicate on the param name; True => SP-sharded (partial grad). + sp_group: SP process group; defaults to get_sp_group(). + + Returns: + Number of params whose grad was summed (0 when SP is off or none matched). + """ + if get_sp_world_size() <= 1: + return 0 + from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors + + if sp_group is None: + sp_group = get_sp_group() + by_dtype = {} + for name, p in named_parameters: + if p.grad is None or not is_sharded(name): + continue + local = p.grad.to_local() if hasattr(p.grad, "to_local") else p.grad + by_dtype.setdefault(local.dtype, []).append(local) + for grads in by_dtype.values(): + flat = _flatten_dense_tensors(grads) + dist.all_reduce(flat, op=dist.ReduceOp.SUM, group=sp_group) + for t, synced in zip(grads, _unflatten_dense_tensors(flat, grads)): + t.copy_(synced) + return sum(len(v) for v in by_dtype.values()) + + +def _all_to_all_4d_impl( + input_: torch.Tensor, + scatter_dim: int, + gather_dim: int, + sp_world_size: int, + sp_group, +) -> torch.Tensor: + """ + Internal implementation of 4D all-to-all (non-differentiable). + + For seq->head (scatter_dim=2, gather_dim=1): + - Input: [B, S/sp, H, D] -> Output: [B, S, H/sp, D] + + For head->seq (scatter_dim=1, gather_dim=2): + - Input: [B, S, H/sp, D] -> Output: [B, S/sp, H, D] + """ + if sp_world_size == 1: + return input_ + + if scatter_dim == 2 and gather_dim == 1: + bs, shard_seqlen, h, d = input_.shape + seq_len = shard_seqlen * sp_world_size + shard_h = h // sp_world_size + sp = sp_world_size + input_ = input_.view(bs, shard_seqlen, sp, shard_h, d) + input_ = input_.permute( + 2, 0, 1, 3, 4 + ).contiguous() # sp, bs, shard_seqlen, shard_h, d + output = torch.empty_like(input_) + dist.all_to_all_single( + output, input_, group=sp_group + ) # sp, bs, shard_seqlen, shard_h, d + + output = output.permute( + 1, 0, 2, 3, 4 + ).contiguous() # bs, sp, shard_seqlen, shard_h, d + output = output.view(bs, seq_len, shard_h, d) # bs, seq_len, h, d + + elif scatter_dim == 1 and gather_dim == 2: + bs, seq_len, shard_h, d = input_.shape + h = shard_h * sp_world_size + shard_seqlen = seq_len // sp_world_size + sp = sp_world_size + input_ = input_.view(bs, sp, shard_seqlen, shard_h, d) + input_ = input_.permute( + 1, 0, 2, 3, 4 + ).contiguous() # sp, bs, shard_seqlen, shard_h, d + output = torch.empty_like(input_) + dist.all_to_all_single( + output, input_, group=sp_group + ) # sp, bs, shard_seqlen, shard_h, d + output = output.permute( + 1, 2, 0, 3, 4 + ).contiguous() # bs, shard_seqlen, sp, shard_h, d + output = output.view(bs, shard_seqlen, h, d) # bs, shard_seqlen, h, d + + return output + + +class _SeqToHeadAllToAll(torch.autograd.Function): + """ + Differentiable all-to-all: sequence sharding -> head sharding. + + Forward: [B, S/sp, H, D] -> [B, S, H/sp, D] + Backward: [B, S, H/sp, D] -> [B, S/sp, H, D] + """ + + @staticmethod + def forward(ctx, input_: torch.Tensor) -> torch.Tensor: + sp_world_size = get_sp_world_size() + sp_group = get_sp_group() + ctx.sp_world_size = sp_world_size + ctx.sp_group = sp_group + return _all_to_all_4d_impl( + input_, + scatter_dim=2, + gather_dim=1, + sp_world_size=sp_world_size, + sp_group=sp_group, + ) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: + # Backward of seq->head is head->seq + return _all_to_all_4d_impl( + grad_output, + scatter_dim=1, + gather_dim=2, + sp_world_size=ctx.sp_world_size, + sp_group=ctx.sp_group, + ) + + +class _HeadToSeqAllToAll(torch.autograd.Function): + """ + Differentiable all-to-all: head sharding -> sequence sharding. + + Forward: [B, S, H/sp, D] -> [B, S/sp, H, D] + Backward: [B, S/sp, H, D] -> [B, S, H/sp, D] + """ + + @staticmethod + def forward(ctx, input_: torch.Tensor) -> torch.Tensor: + sp_world_size = get_sp_world_size() + sp_group = get_sp_group() + ctx.sp_world_size = sp_world_size + ctx.sp_group = sp_group + return _all_to_all_4d_impl( + input_, + scatter_dim=1, + gather_dim=2, + sp_world_size=sp_world_size, + sp_group=sp_group, + ) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: + # Backward of head->seq is seq->head + return _all_to_all_4d_impl( + grad_output, + scatter_dim=2, + gather_dim=1, + sp_world_size=ctx.sp_world_size, + sp_group=ctx.sp_group, + ) + + +def sp_all_to_all_4d( + input_: torch.Tensor, + scatter_dim: int = 2, + gather_dim: int = 1, +) -> torch.Tensor: + """ + Differentiable all-to-all communication for 4D tensors. + + In attention computation: + - Input: [B, S/sp, H, D] (sequence sharded) + - Output: [B, S, H/sp, D] (head sharded) + + Args: + input_: Input tensor [B, S_local, H, D] + scatter_dim: Scatter dimension (usually the head dimension) + gather_dim: Gather dimension (usually the sequence dimension) + + Returns: + Rearranged tensor + """ + sp_world_size = get_sp_world_size() + if sp_world_size == 1: + return input_ + + return _SeqToHeadAllToAll.apply(input_) + + +def sp_all_gather_4d( + input_: torch.Tensor, + scatter_dim: int = 1, + gather_dim: int = 2, +) -> torch.Tensor: + """ + Differentiable reverse all-to-all for sp_all_to_all_4d. + + After attention computation: + - Input: [B, S, H/sp, D] (head sharded) + - Output: [B, S/sp, H, D] (sequence sharded) + """ + sp_world_size = get_sp_world_size() + if sp_world_size == 1: + return input_ + + return _HeadToSeqAllToAll.apply(input_) diff --git a/wall_x/distributed/parallel_state.py b/wall_x/distributed/parallel_state.py new file mode 100644 index 00000000..7ceea7d0 --- /dev/null +++ b/wall_x/distributed/parallel_state.py @@ -0,0 +1,192 @@ +"""Sequence-parallel and data-parallel process group state.""" + +import logging + +import torch.distributed as dist +from typing import Optional + +logger = logging.getLogger(__name__) + +# Global SP group state +_SP_GROUP: Optional[dist.ProcessGroup] = None +_SP_WORLD_SIZE: int = 1 +_SP_RANK: int = 0 + +# Global DP group state +_DP_GROUP: Optional[dist.ProcessGroup] = None +_DP_WORLD_SIZE: int = 1 +_DP_RANK: int = 0 + +# Track if parallel groups are initialized +_PARALLEL_INITIALIZED: bool = False + + +def init_sequence_parallel( + sp_size: int = 1, + backend: str = "nccl", +) -> None: + """ + Initialize Sequence Parallelism and Data Parallelism process groups. + + Args: + sp_size: SP group size (number of GPUs per group) + backend: Distributed backend, default "nccl" + + Example: + # 8 GPUs, sp_size=4 will create: + # + # 2 SP groups (consecutive partitioning): + # SP Group 0: [GPU 0, 1, 2, 3] + # SP Group 1: [GPU 4, 5, 6, 7] + # + # 4 DP groups (strided partitioning): + # DP Group 0: [GPU 0, 4] + # DP Group 1: [GPU 1, 5] + # DP Group 2: [GPU 2, 6] + # DP Group 3: [GPU 3, 7] + """ + global _SP_GROUP, _SP_WORLD_SIZE, _SP_RANK + global _DP_GROUP, _DP_WORLD_SIZE, _DP_RANK + global _PARALLEL_INITIALIZED + + if not dist.is_initialized(): + raise RuntimeError( + "torch.distributed must be initialized before init_sequence_parallel" + ) + + if _PARALLEL_INITIALIZED: + # Already initialized, skip + return + + world_size = dist.get_world_size() + rank = dist.get_rank() + + if sp_size == 1: + # No SP, DP covers the entire world + _SP_GROUP = None + _SP_WORLD_SIZE = 1 + _SP_RANK = 0 + _DP_GROUP = None # None means use the default world group + _DP_WORLD_SIZE = world_size + _DP_RANK = rank + _PARALLEL_INITIALIZED = True + if rank == 0: + logger.info( + "[Parallel] SP disabled, DP uses world group (size=%s)", world_size + ) + return + + assert ( + world_size % sp_size == 0 + ), f"world_size ({world_size}) must be divisible by sp_size ({sp_size})" + + num_sp_groups = world_size // sp_size + + # ============ Create SP groups (consecutive partitioning) ============ + # GPUs within an SP group process different sequence segments of the same sample + for i in range(num_sp_groups): + ranks = list(range(i * sp_size, (i + 1) * sp_size)) + group = dist.new_group(ranks, backend=backend) + if rank in ranks: + _SP_GROUP = group + _SP_WORLD_SIZE = sp_size + _SP_RANK = ranks.index(rank) + + # ============ Create DP groups (strided partitioning) ============ + # GPUs within a DP group process different data samples, need gradient synchronization + # Each DP group contains GPUs from the same position across different SP groups + num_dp_groups = sp_size + for i in range(num_dp_groups): + # Select ranks in strided manner: [i, i+sp_size, i+2*sp_size, ...] + ranks = list(range(i, world_size, sp_size)) + group = dist.new_group(ranks, backend=backend) + if rank in ranks: + _DP_GROUP = group + _DP_WORLD_SIZE = len(ranks) # = num_sp_groups + _DP_RANK = ranks.index(rank) + + _PARALLEL_INITIALIZED = True + + if rank == 0: + logger.info("[Parallel] Initialized SP+DP groups:") + logger.info( + " - SP: %s groups, each with %s GPUs (consecutive)", + num_sp_groups, + sp_size, + ) + logger.info( + " - DP: %s groups, each with %s GPUs (strided)", + num_dp_groups, + num_sp_groups, + ) + + +def is_sp_initialized() -> bool: + """Check if SP is initialized (sp_size > 1)""" + return _SP_WORLD_SIZE > 1 + + +def is_parallel_initialized() -> bool: + """Check if parallel groups are initialized""" + return _PARALLEL_INITIALIZED + + +def get_sp_group() -> Optional[dist.ProcessGroup]: + """Get the SP group that the current process belongs to""" + return _SP_GROUP + + +def get_sp_world_size() -> int: + """Get SP group size""" + return _SP_WORLD_SIZE + + +def get_sp_rank() -> int: + """Get the rank of current process within its SP group""" + return _SP_RANK + + +def get_dp_group() -> Optional[dist.ProcessGroup]: + """ + Get the DP group that the current process belongs to. + + Returns: + DP process group. Returns None if sp_size=1 (use default world group) + """ + return _DP_GROUP + + +def get_dp_world_size() -> int: + """Get DP group size (i.e., number of SP groups)""" + return _DP_WORLD_SIZE + + +def get_dp_rank() -> int: + """Get the rank of current process within its DP group""" + return _DP_RANK + + +def destroy_parallel_groups() -> None: + """Destroy all parallel groups""" + global _SP_GROUP, _SP_WORLD_SIZE, _SP_RANK + global _DP_GROUP, _DP_WORLD_SIZE, _DP_RANK + global _PARALLEL_INITIALIZED + + if _SP_GROUP is not None: + dist.destroy_process_group(_SP_GROUP) + if _DP_GROUP is not None: + dist.destroy_process_group(_DP_GROUP) + + _SP_GROUP = None + _SP_WORLD_SIZE = 1 + _SP_RANK = 0 + _DP_GROUP = None + _DP_WORLD_SIZE = 1 + _DP_RANK = 0 + _PARALLEL_INITIALIZED = False + + +# Backward compatibility +def destroy_sp_group() -> None: + """Destroy SP group (backward compatible interface)""" + destroy_parallel_groups() diff --git a/wall_x/model/core/action/__init__.py b/wall_x/model/core/action/__init__.py index 65582b14..aa2a816f 100644 --- a/wall_x/model/core/action/__init__.py +++ b/wall_x/model/core/action/__init__.py @@ -1,6 +1,5 @@ """Action model components: normalizer, processor, head, MoE.""" -from wall_x.model.core.action.head import SinusoidalPosEmb from wall_x.model.core.action.normalizer import ( Normalizer, create_normalizers, @@ -8,3 +7,4 @@ unnormalize_data_with_virtual_tail, ) from wall_x.model.core.action.processor import ActionProcessor +from wall_x.model.core.action.head import SinusoidalPosEmb diff --git a/wall_x/model/core/action/layout.py b/wall_x/model/core/action/layout.py new file mode 100644 index 00000000..5232fc19 --- /dev/null +++ b/wall_x/model/core/action/layout.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + + +@dataclass(frozen=True) +class TwoExpertLayout: + action_starts: torch.Tensor + action_ends: torch.Tensor + llm_offsets: torch.Tensor + action_offsets: torch.Tensor + expert_segments: tuple[tuple[int, int], tuple[int, int]] + + +def fixed_two_expert_action_ranges( + experts_indices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + if experts_indices.dim() != 2: + raise ValueError("fixed scatter MoE path expects [B, S] expert indices") + + action_mask = experts_indices == 1 + has_action = action_mask.any(dim=1) + starts_raw = torch.argmax(action_mask.to(torch.int64), dim=1) + ends_raw = experts_indices.size(1) - torch.argmax( + action_mask.flip(dims=[1]).to(torch.int64), dim=1 + ) + zeros = torch.zeros_like(starts_raw) + starts = torch.where(has_action, starts_raw, zeros) + ends = torch.where(has_action, ends_raw, starts) + action_lens = action_mask.sum(dim=1) + if not bool(((ends - starts) == action_lens).all().item()): + raise ValueError("fixed scatter MoE path requires contiguous action tokens") + return starts.contiguous(), ends.contiguous() + + +def resolve_two_expert_action_ranges( + experts_indices: Optional[torch.Tensor], + action_starts: Optional[torch.Tensor] = None, + action_ends: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + if experts_indices is None or experts_indices.dim() != 2: + raise ValueError("fixed scatter MoE path expects [B, S] expert indices") + + batch_size = experts_indices.size(0) + device = experts_indices.device + if action_starts is None and action_ends is None: + return fixed_two_expert_action_ranges(experts_indices) + if action_starts is None or action_ends is None: + raise ValueError("action_starts and action_ends must be provided together") + + action_starts = torch.as_tensor( + action_starts, dtype=torch.int64, device=device + ).reshape(-1) + action_ends = torch.as_tensor( + action_ends, dtype=torch.int64, device=device + ).reshape(-1) + + if action_starts.numel() != batch_size or action_ends.numel() != batch_size: + raise ValueError("fixed scatter MoE path expects [B] action ranges") + return action_starts.contiguous(), action_ends.contiguous() + + +def fixed_two_expert_segments( + batch_size: int, + seq_length: int, + action_starts: torch.Tensor, + action_ends: torch.Tensor, +) -> tuple[tuple[int, int], tuple[int, int]]: + starts_cpu = torch.as_tensor(action_starts, dtype=torch.int64).reshape(-1).cpu() + ends_cpu = torch.as_tensor(action_ends, dtype=torch.int64).reshape(-1).cpu() + if starts_cpu.numel() != int(batch_size) or ends_cpu.numel() != int(batch_size): + raise ValueError( + "fixed scatter MoE path expects one action range per batch sample" + ) + + action_lens = ends_cpu - starts_cpu + total_action = int(action_lens.sum().item()) + total_tokens = int(batch_size) * int(seq_length) + llm_end = total_tokens - total_action + return (0, llm_end), (llm_end, total_tokens) + + +def _cpu_int64_1d(values: torch.Tensor) -> torch.Tensor: + if not isinstance(values, torch.Tensor): + values = torch.as_tensor(values, dtype=torch.int64) + values = values.detach().reshape(-1) + if values.device.type != "cpu" or values.dtype != torch.int64: + values = values.to(device="cpu", dtype=torch.int64) + return values.contiguous() + + +def _device_int64_1d( + values: torch.Tensor, + fallback_cpu: torch.Tensor, + *, + device: torch.device, +) -> torch.Tensor: + if isinstance(values, torch.Tensor) and values.device == device: + values = values.reshape(-1) + if values.dtype != torch.int64: + values = values.to(dtype=torch.int64) + if not values.is_contiguous(): + values = values.contiguous() + return values + return fallback_cpu.to(device=device, non_blocking=True) + + +def prepare_two_expert_layout( + experts_indices: Optional[torch.Tensor], + *, + batch_size: Optional[int] = None, + seq_length: Optional[int] = None, + action_starts: Optional[torch.Tensor] = None, + action_ends: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, +) -> TwoExpertLayout: + if experts_indices is not None: + if experts_indices.dim() != 2: + raise ValueError("fixed scatter MoE path expects [B, S] expert indices") + batch_size = experts_indices.size(0) + seq_length = experts_indices.size(1) + if device is None: + device = experts_indices.device + if batch_size is None or seq_length is None: + raise ValueError("batch_size and seq_length are required") + if device is None: + raise ValueError("device is required") + + if action_starts is None or action_ends is None: + action_starts, action_ends = resolve_two_expert_action_ranges( + experts_indices, + action_starts=action_starts, + action_ends=action_ends, + ) + + starts_cpu = _cpu_int64_1d(action_starts) + ends_cpu = _cpu_int64_1d(action_ends) + if starts_cpu.numel() != batch_size or ends_cpu.numel() != batch_size: + raise ValueError("fixed scatter MoE path expects [B] action ranges") + if not bool( + ((starts_cpu >= 0) & (ends_cpu >= starts_cpu) & (ends_cpu <= seq_length)) + .all() + .item() + ): + raise ValueError( + "fixed scatter MoE action ranges must satisfy " + "0 <= start <= end <= seq_length" + ) + + action_lens = ends_cpu - starts_cpu + llm_counts = int(seq_length) - action_lens + zero = torch.zeros(1, dtype=torch.int64) + llm_offsets_cpu = torch.cat((zero, torch.cumsum(llm_counts, dim=0))).contiguous() + action_offsets_cpu = torch.cat( + (zero, torch.cumsum(action_lens, dim=0)) + ).contiguous() + total_llm = int(llm_offsets_cpu[-1].item()) + total_tokens = int(batch_size) * int(seq_length) + + device = torch.device(device) + return TwoExpertLayout( + action_starts=_device_int64_1d(action_starts, starts_cpu, device=device), + action_ends=_device_int64_1d(action_ends, ends_cpu, device=device), + llm_offsets=llm_offsets_cpu.to(device=device, non_blocking=True), + action_offsets=action_offsets_cpu.to(device=device, non_blocking=True), + expert_segments=((0, total_llm), (total_llm, total_tokens)), + ) diff --git a/wall_x/model/core/action/moe.py b/wall_x/model/core/action/moe.py index a57e77d1..59b2f15b 100644 --- a/wall_x/model/core/action/moe.py +++ b/wall_x/model/core/action/moe.py @@ -1,9 +1,29 @@ import torch import torch.nn as nn import torch.utils.checkpoint as cp + from transformers.activations import ACT2FN -from wall_x.model.core.ops import permute, unpermute +from wall_x.model.core.action.layout import ( + TwoExpertLayout, + fixed_two_expert_action_ranges, + fixed_two_expert_segments, + prepare_two_expert_layout, + resolve_two_expert_action_ranges, +) +from wall_x.model.core.ops import silu_and_mul, two_expert_pack, two_expert_unpack + + +__all__ = [ + "TwoExpertLayout", + "fixed_two_expert_action_ranges", + "fixed_two_expert_segments", + "prepare_two_expert_layout", + "resolve_two_expert_action_ranges", + "TokenTypeRouter", + "BlockSparseMLP", + "SparseMoeBlock", +] class TokenTypeRouter(nn.Module): @@ -41,15 +61,19 @@ def __init__(self, config, use_selective_recompute: bool = False): self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[self.hidden_act] + self.use_silu_and_mul = self.hidden_act in ("silu", "swish") # FIXME: The full MLP is recomputed for now; recomputing only activations can be optimized later. def _full_mlp(self, hidden_state): gate_up_out = self.gate_up_proj(hidden_state) - gate_out, up_out = gate_up_out.split( - [self.intermediate_size, self.intermediate_size], dim=-1 - ) - act_out = self.act_fn(gate_out) * up_out + if self.use_silu_and_mul: + act_out = silu_and_mul(gate_up_out, inplace=True) + else: + gate_out, up_out = gate_up_out.split( + [self.intermediate_size, self.intermediate_size], dim=-1 + ) + act_out = self.act_fn(gate_out) * up_out return self.down_proj(act_out) def forward(self, hidden_state): @@ -86,30 +110,43 @@ def __init__(self, config, num_experts: int, use_selective_recompute: bool = Fal self.dim_inputs = config.dim_inputs self.permuted = config.mot_opt + def _fixed_two_expert_action_ranges(self, experts_indices: torch.Tensor): + if self.num_experts != 2: + raise ValueError("fixed scatter MoE path requires num_experts == 2") + return fixed_two_expert_action_ranges(experts_indices) + def forward( self, hidden_states: torch.Tensor, experts_indices: torch.Tensor, - start_indices: torch.Tensor, - end_indices: torch.Tensor, + action_starts: torch.Tensor = None, + action_ends: torch.Tensor = None, + expert_segments: tuple[tuple[int, int], tuple[int, int]] = None, + two_expert_layout: TwoExpertLayout = None, ) -> torch.Tensor: + batch_size, seq_length = experts_indices.shape + if two_expert_layout is None: + two_expert_layout = prepare_two_expert_layout( + experts_indices, + batch_size=batch_size, + seq_length=seq_length, + action_starts=action_starts, + action_ends=action_ends, + device=hidden_states.device, + ) + expert_segments = two_expert_layout.expert_segments if self.permuted: permuted_inputs = hidden_states else: - batch_size, seq_length, hidden_dim = hidden_states.shape - - flat_hidden = hidden_states.reshape(-1, hidden_dim) - experts_indices = experts_indices.reshape(-1) - probs = torch.ones_like(experts_indices, dtype=torch.float32).reshape(-1, 1) - permuted_inputs, row_id_map = permute(flat_hidden, experts_indices) + permuted_inputs = two_expert_pack(hidden_states, two_expert_layout) # buffer final_output = torch.zeros_like(permuted_inputs) # Expert forward, including selective recompute for expert_idx, expert in enumerate(self.experts): - start, end = start_indices[expert_idx], end_indices[expert_idx] + start, end = expert_segments[expert_idx] if start == end: continue @@ -122,5 +159,6 @@ def forward( if self.permuted: return final_output else: - final_output = unpermute(final_output, row_id_map, probs) - return final_output.reshape(batch_size, seq_length, hidden_dim) + return two_expert_unpack( + final_output, batch_size, seq_length, two_expert_layout + ) diff --git a/wall_x/model/core/action/normalizer.py b/wall_x/model/core/action/normalizer.py index 20af594a..81f8870a 100644 --- a/wall_x/model/core/action/normalizer.py +++ b/wall_x/model/core/action/normalizer.py @@ -187,6 +187,34 @@ def from_lerobot_norm_stats( return instance + @classmethod + def from_vectors(cls, min_by_key, delta_by_key): + """Build a Normalizer directly from precomputed per-dataset min/delta. + + ``min_by_key`` / ``delta_by_key`` are ``{dataset_name: 1-D array/tensor}`` + with matching keys. Used by :func:`create_normalizers_from_v2_stats` so + the model-side normalizer carries the exact statistics the + internal_dataset_backend_v2 pipeline used. + """ + instance = cls.__new__(cls) + nn.Module.__init__(instance) + + instance.min = nn.ParameterDict() + instance.delta = nn.ParameterDict() + instance.min_key = "min" + instance.delta_key = "delta" + + for name, vec in min_by_key.items(): + instance.min[name] = nn.Parameter( + torch.as_tensor(vec, dtype=torch.float32), requires_grad=False + ) + for name, vec in delta_by_key.items(): + instance.delta[name] = nn.Parameter( + torch.as_tensor(vec, dtype=torch.float32), requires_grad=False + ) + + return instance + def __init__( self, action_statistic_dof, @@ -206,31 +234,36 @@ def __init__( action_statistic[robot_name] = {} all_dof_min = [] all_dof_delta = [] - for k in dof_config: + for config_key, dof_dim in dof_config.items(): + # Use a fallback stats key without mutating the model dof key. + stat_key = config_key + follow_key = config_key.replace("master_", "follow_", 1) if ( - k not in action_statistic_dof[robot_name] - and k.replace("master_", "follow_", 1) - in action_statistic_dof[robot_name] + stat_key not in action_statistic_dof[robot_name] + and follow_key in action_statistic_dof[robot_name] ): - k = k.replace("master_", "follow_", 1) - if k in action_statistic_dof[robot_name]: + stat_key = follow_key + + if stat_key in action_statistic_dof[robot_name]: if ( - min_key in action_statistic_dof[robot_name][k] - and delta_key in action_statistic_dof[robot_name][k] + min_key in action_statistic_dof[robot_name][stat_key] + and delta_key in action_statistic_dof[robot_name][stat_key] ): - all_dof_min.extend(action_statistic_dof[robot_name][k][min_key]) + all_dof_min.extend( + action_statistic_dof[robot_name][stat_key][min_key] + ) all_dof_delta.extend( - action_statistic_dof[robot_name][k][delta_key] + action_statistic_dof[robot_name][stat_key][delta_key] ) else: normalizer_missing_information.append( - f"Normalizer (Warning): min_key {min_key} or delta_key {delta_key} not in action_statistic_dof[{robot_name}][{k}], use default min 0.0 and delta 1.0" + f"Normalizer (Warning): min_key {min_key} or delta_key {delta_key} not in action_statistic_dof[{robot_name}][{stat_key}], use default min 0.0 and delta 1.0" ) - all_dof_min.extend([0.0] * dof_config[k]) - all_dof_delta.extend([1.0] * dof_config[k]) + all_dof_min.extend([0.0] * dof_dim) + all_dof_delta.extend([1.0] * dof_dim) else: - # k is a model key; action_statistic_dof may store raw keys - try fallback lookup - raw_key = _MODEL_KEY_TO_RAW_KEY.get(k) + # config_key is a model key; action_statistic_dof may store raw keys. + raw_key = _MODEL_KEY_TO_RAW_KEY.get(config_key) if ( raw_key is not None and raw_key in action_statistic_dof[robot_name] @@ -241,16 +274,16 @@ def __init__( all_dof_delta.extend(stat[delta_key]) else: normalizer_missing_information.append( - f"Normalizer (Warning): Action {k} not in action_statistic_dof for {robot_name}, use default min 0.0 and delta 1.0" + f"Normalizer (Warning): Action {config_key} not in action_statistic_dof for {robot_name}, use default min 0.0 and delta 1.0" ) - all_dof_min.extend([0.0] * dof_config[k]) - all_dof_delta.extend([1.0] * dof_config[k]) + all_dof_min.extend([0.0] * dof_dim) + all_dof_delta.extend([1.0] * dof_dim) else: normalizer_missing_information.append( - f"Normalizer (Warning): Action {k} not in action_statistic_dof for {robot_name}, use default min 0.0 and delta 1.0" + f"Normalizer (Warning): Action {config_key} not in action_statistic_dof for {robot_name}, use default min 0.0 and delta 1.0" ) - all_dof_min.extend([0.0] * dof_config[k]) - all_dof_delta.extend([1.0] * dof_config[k]) + all_dof_min.extend([0.0] * dof_dim) + all_dof_delta.extend([1.0] * dof_dim) all_dof_min = torch.tensor(all_dof_min) all_dof_delta = torch.tensor(all_dof_delta) action_statistic[robot_name][min_key] = all_dof_min @@ -344,48 +377,6 @@ def _pad_lerobot_stats(stats, target_dim, label): return type(stats)(min=min_stat, max=max_stat, delta=delta) -def pad_normalizer_to_dim(normalizer, target_dim, label): - """Tail-pad every dataset entry to ``target_dim`` (virtual ``action_padding`` tail).""" - target_dim = int(target_dim) - for name in list(normalizer.min.keys()): - current_dim = int(normalizer.min[name].numel()) - if current_dim == target_dim: - continue - if current_dim > target_dim: - raise ValueError( - f"Normalizer {label}[{name!r}] dim ({current_dim}) exceeds configured " - f"{label} dim ({target_dim})" - ) - - pad_dim = target_dim - current_dim - dtype = normalizer.min[name].dtype - normalizer.min[name] = nn.Parameter( - torch.cat( - [ - normalizer.min[name].flatten(), - torch.zeros(pad_dim, dtype=dtype), - ] - ), - requires_grad=False, - ) - normalizer.delta[name] = nn.Parameter( - torch.cat( - [ - normalizer.delta[name].flatten(), - torch.ones(pad_dim, dtype=dtype), - ] - ), - requires_grad=False, - ) - logger.info( - "Padded %s normalizer[%r] from %d to %d dims", - label, - name, - current_dim, - target_dim, - ) - - def create_normalizers_from_lerobot_norm_stats( norm_stats, dataset_name, action_dim, propri_dim ): @@ -405,20 +396,21 @@ def create_normalizers(config, action_statistic_dof=None): Public Wall-X builds intentionally do not bundle private default stats. """ + norm_cfg = config.get("normalizer_config") or {} if action_statistic_dof is None: - custom_path = config.get("customized_action_statistic_dof") + custom_path = norm_cfg.get("customized_action_statistic_dof") if custom_path: with open(custom_path, "r") as f: action_statistic_dof = json.load(f) else: raise ValueError( "create_normalizers requires action statistics. Provide " - "`customized_action_statistic_dof` in the config or pass " - "`action_statistic_dof` from the checkpoint/dataset." + "`customized_action_statistic_dof` in data.normalizer_config or " + "pass `action_statistic_dof` from the checkpoint/dataset." ) - min_key = config.get("min_key", "min") - delta_key = config.get("delta_key", "delta") + min_key = norm_cfg.get("min_key", "min") + delta_key = norm_cfg.get("delta_key", "delta") # Required keys for constructing action/proprio normalizers. missing_required = [] @@ -445,3 +437,148 @@ def create_normalizers(config, action_statistic_dof=None): delta_key=delta_key, ) return normalizer_action, normalizer_propri, action_statistic_dof + + +NORM_KEY_SEP = "/" +V2_DEFAULT_MIN_KEY = "q001" +V2_DEFAULT_DELTA_KEY = "delta_001" + + +def make_norm_key(dataset_type, end_effector_type=None): + """Compose the model/serving normalizer key from a ``(dataset_type, + end_effector_type)`` pair. + + Returns the bare ``dataset_type`` when there is no end_effector_type + (flat/legacy stats, e.g. lerobot / x2robot_v1). + """ + if end_effector_type is None or end_effector_type == "": + return dataset_type + return f"{dataset_type}{NORM_KEY_SEP}{end_effector_type}" + + +def _list_v2_end_effector_types(ds_stats, min_key, delta_key): + """List the end_effector_type levels under one dataset_type. + + Returns ``[ee_type, ...]`` for the two-level layout + ``stats[dataset_type][end_effector_type][action_key]`` (one or many), or + ``[None]`` when the layout is flat (one-level). + """ + if not isinstance(ds_stats, dict) or not ds_stats: + return [None] + + def _is_stat_leaf(d): + return isinstance(d, dict) and ( + min_key in d or delta_key in d or "min" in d or "q001" in d + ) + + children = list(ds_stats.values()) + if any(_is_stat_leaf(c) for c in children): + return [None] # flat / legacy one-level layout + return list(ds_stats.keys()) + + +def normalizer_is_identity(normalizer): + """True if every dataset entry is the identity transform (min==0, delta==1).""" + min_dict = getattr(normalizer, "min", None) + delta_dict = getattr(normalizer, "delta", None) + if not min_dict or not delta_dict: + return False + for k in min_dict.keys(): + if not torch.allclose(min_dict[k], torch.zeros_like(min_dict[k])): + return False + if k in delta_dict and not torch.allclose( + delta_dict[k], torch.ones_like(delta_dict[k]) + ): + return False + return True + + +def create_normalizers_from_v2_stats(config): + """Build wall-x ``Normalizer`` whose min/delta come from internal_dataset_backend_v2. + + Single source of truth for ``dataset_type == x2robot_v2`` checkpoints: the + statistics are resolved by v2's own ``MinMaxNormalizer`` (two-level + ``dataset_type/end_effector_type`` lookup, ``master_``->``follow_`` fallback, + ``key_dims`` truncation), so the model-side normalizer stays bit-exact with + what the v2 dataset pipeline applied during training. The wall-x + ``normalize_data`` / ``unnormalize_data`` math is identical to v2's + (``(x-min)/delta*2-1`` / ``(x+1)/2*delta+min``), so feeding v2-derived + min/delta reproduces training normalization exactly. + + The resulting ``Normalizer`` is keyed by the composite + ``make_norm_key(dataset_type, end_effector_type)`` (= ``dataset_type/ee_type`` + for two-level stats, bare ``dataset_type`` for flat stats). A dataset_type + with several end_effector_types yields one entry each, so a single checkpoint + can serve any of them; serving selects via ``InferConfig.end_effector_type`` + or an unambiguous bare ``norm_key`` (see ``_resolve_norm_key``). + + Expected ``config`` keys: ``dof_config`` (action keys+dims, order = model + action layout), ``agent_pos_config`` (proprio keys+dims), ``stats_file`` (or + ``customized_action_statistic_dof``), ``min_key`` (default ``q001``), + ``delta_key`` (default ``delta_001``); plus any extra v2 ``normalizer_config`` + fields, which are forwarded to ``build_action_normalizer`` verbatim. Returns + ``(normalizer_action, normalizer_propri, None)`` to match + :func:`create_normalizers`. + """ + from internal_dataset_backend_v2.common.normalizer import build_action_normalizer + + stats_file = config.get("stats_file") or config.get( + "customized_action_statistic_dof" + ) + if not stats_file: + raise ValueError( + "create_normalizers_from_v2_stats requires 'stats_file' " + "(or 'customized_action_statistic_dof')." + ) + dof_config = config["dof_config"] + agent_pos_config = config["agent_pos_config"] + min_key = config.get("min_key", V2_DEFAULT_MIN_KEY) + delta_key = config.get("delta_key", V2_DEFAULT_DELTA_KEY) + + # Forward any extra v2 normalizer_config fields verbatim so the bridge uses + # the exact same normalizer config the dataset pipeline used; force the + # stats source / keys we resolved above. + reserved = { + "dof_config", + "agent_pos_config", + "stats_file", + "customized_action_statistic_dof", + } + v2_cfg = {k: v for k, v in config.items() if k not in reserved} + v2_cfg.setdefault("type", "min_max") + v2_cfg["stats_file"] = stats_file + v2_cfg["min_key"] = min_key + v2_cfg["delta_key"] = delta_key + + # key_dims feeds v2's authoritative fallback dim. Action+proprio keys that + # collide must share the same dim (true for the v2 configs in use). + key_dims = {**dof_config, **agent_pos_config} + v2norm = build_action_normalizer(v2_cfg, key_dims=key_dims) + + action_keys = list(dof_config.keys()) + propri_keys = list(agent_pos_config.keys()) + + with open(stats_file, "r") as f: + raw_stats = json.load(f) + + get_vecs = getattr(v2norm, "get_min_delta_vectors", None) or v2norm._get_params + + action_min, action_delta = {}, {} + propri_min, propri_delta = {}, {} + for dataset_type, ds_stats in raw_stats.items(): + # A dataset_type may carry several end_effector_types (e.g. the same data + # collected with different hands). Build one entry per (dataset_type, + # end_effector_type), keyed by the composite ``dataset_type/ee_type`` so a + # single checkpoint can serve any of them. Two-level stats always use the + # composite key (even with a single ee_type); flat stats stay bare. + eets = _list_v2_end_effector_types(ds_stats, min_key, delta_key) + for eet in eets: + a_min, a_delta = get_vecs(action_keys, dataset_type, end_effector_type=eet) + p_min, p_delta = get_vecs(propri_keys, dataset_type, end_effector_type=eet) + key = make_norm_key(dataset_type, eet) + action_min[key], action_delta[key] = a_min, a_delta + propri_min[key], propri_delta[key] = p_min, p_delta + + normalizer_action = Normalizer.from_vectors(action_min, action_delta) + normalizer_propri = Normalizer.from_vectors(propri_min, propri_delta) + return normalizer_action, normalizer_propri, None diff --git a/wall_x/model/core/action/processor.py b/wall_x/model/core/action/processor.py index 358352ff..6b3c0524 100644 --- a/wall_x/model/core/action/processor.py +++ b/wall_x/model/core/action/processor.py @@ -165,14 +165,36 @@ def forward(self, action_chunk, dataset_names, sample_time, dof_mask=None): # 1. add noise to action_chunk noise = torch.randn_like(action_chunk) - time_expanded = sample_time.unsqueeze(-1).unsqueeze(-1) + if sample_time.dim() == 1: + per_token_time = False + time_expanded = sample_time.unsqueeze(-1).unsqueeze(-1) + elif sample_time.dim() == 2: + per_token_time = True + if sample_time.shape != action_chunk.shape[:2]: + raise ValueError( + "`sample_time` with shape [B, T] must match " + f"`action_chunk` prefix shape {tuple(action_chunk.shape[:2])}, " + f"got {tuple(sample_time.shape)}." + ) + time_expanded = sample_time.unsqueeze(-1) + else: + raise ValueError( + "`sample_time` must have shape [B] or [B, T], " + f"got {tuple(sample_time.shape)}." + ) noisy_action = ( 1 - time_expanded ) * noise + time_expanded * action_chunk # denoise from 0 to 1; integration does not need a negative sign flow = action_chunk - noise # used to compute loss # 2. sinusoidal positional encoding for timesteps - time_embed = self.time_embed(sample_time).to(torch.float32) + if per_token_time: + time_embed = self.time_embed(sample_time.reshape(-1)).reshape( + sample_time.shape[0], sample_time.shape[1], -1 + ) + else: + time_embed = self.time_embed(sample_time) + time_embed = time_embed.to(torch.float32) self.noise = noise self.noisy_action = noisy_action # for new x-pred @@ -186,6 +208,11 @@ def forward(self, action_chunk, dataset_names, sample_time, dof_mask=None): self.time_expanded = time_expanded # for new x-pred if not self.config.use_adarms: + if per_token_time: + raise NotImplementedError( + "Per-token `sample_time` is only supported when " + "`use_adarms=True`." + ) time_embed = ( time_embed.unsqueeze(1) .repeat(1, action_embed.shape[1], 1) @@ -196,6 +223,7 @@ def forward(self, action_chunk, dataset_names, sample_time, dof_mask=None): action_time_embed = self.w3(self.act_fn(concat_embed)) adarms_cond = None else: + time_embed = time_embed.to(dtype=self.time_mlp_in.weight.dtype) time_embed = self.time_mlp_in(time_embed) time_embed = self.act_fn(time_embed) time_embed = self.time_mlp_out(time_embed) diff --git a/wall_x/model/core/attention/__init__.py b/wall_x/model/core/attention/__init__.py index cd791d1f..8e5c7b60 100644 --- a/wall_x/model/core/attention/__init__.py +++ b/wall_x/model/core/attention/__init__.py @@ -1,9 +1,9 @@ """Attention mechanisms: joint attention, VLA attention, mask builders, backend selector.""" -from wall_x.model.core.attention.mask import ( +from wall_x.model.core.ops.attention_mask import ( find_first_last_ones, - update_joint_attention_flash_mask, - update_joint_attention_mask_2d, update_position_ids, + update_joint_attention_mask_2d, + update_joint_attention_flash_mask, ) from wall_x.model.core.attention.selector import AttentionsSelectorMixin diff --git a/wall_x/model/core/attention/joint.py b/wall_x/model/core/attention/joint.py index aae7433d..2c61fd1d 100644 --- a/wall_x/model/core/attention/joint.py +++ b/wall_x/model/core/attention/joint.py @@ -1,21 +1,26 @@ -from typing import Optional, Tuple - import torch import torch.nn as nn -from flash_attn import flash_attn_func -from transformers.cache_utils import Cache -from transformers.modeling_flash_attention_utils import ( - is_flash_attn_greater_or_equal_2_10, -) -from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig +from typing import Optional, Tuple + from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( Qwen2_5_VLRotaryEmbedding, repeat_kv, ) + +from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig +from transformers.cache_utils import Cache from transformers.utils import logging -from wall_x.model.core.attention.mask import find_first_last_ones -from wall_x.model.core.ops import m_rope, permute, unpermute +from wall_x.model.core.ops.attention_mask import find_first_last_ones +from wall_x.model.core.action.layout import TwoExpertLayout, prepare_two_expert_layout + +from wall_x.model.core.ops import m_rope, two_expert_pack, two_expert_unpack + +from flash_attn import flash_attn_func + +from transformers.modeling_flash_attention_utils import ( + is_flash_attn_greater_or_equal_2_10, +) try: from flash_mask.flash_mask_interface import flash_mask_attn_func @@ -119,11 +124,11 @@ def forward( cache_position: Optional[torch.LongTensor] = None, token_types: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - start_indices: Optional[torch.Tensor] = None, - end_indices: Optional[torch.Tensor] = None, - probs: Optional[torch.Tensor] = None, - row_id_map: Optional[torch.Tensor] = None, orig_shape: Optional[Tuple[int]] = None, + action_starts: Optional[torch.Tensor] = None, + action_ends: Optional[torch.Tensor] = None, + expert_segments: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None, + two_expert_layout: Optional[TwoExpertLayout] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if token_types is None: raise ValueError("token_types must not be empty") @@ -136,13 +141,13 @@ def forward( bsz, q_len, _ = orig_shape query_states, key_states, value_states = self._generate_qkv_mot_opt( hidden_states, - token_types, - start_indices, - end_indices, - probs, - row_id_map, bsz, q_len, + action_starts, + action_ends, + expert_segments, + two_expert_layout, + token_types, ) else: bsz, q_len, _ = hidden_states.size() @@ -262,7 +267,12 @@ def forward( if self.config.mot_opt: output = self._generate_output_mot_opt( - attn_output, token_types, start_indices, end_indices + attn_output, + action_starts, + action_ends, + expert_segments, + two_expert_layout, + token_types, ) else: output = self._generate_output(attn_output, masks) @@ -350,24 +360,20 @@ def _generate_qkv(self, hidden_states, masks): def _generate_qkv_mot_opt( self, hidden_states: torch.Tensor, - experts_indices: torch.Tensor, - start_indices: torch.Tensor, - end_indices: torch.Tensor, - probs: torch.Tensor, - row_id_map: torch.Tensor, batch_size: int, seq_length: int, + action_starts: torch.Tensor, + action_ends: torch.Tensor, + expert_segments: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None, + two_expert_layout: Optional[TwoExpertLayout] = None, + token_types: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ - Generate Q, K, V based on expert-sharded segments (start_indices / end_indices), + Generate Q, K, V based on fixed two-expert LLM/action segments, then restore them to the original sequence order. Args: hidden_states: [total_tokens, hidden_dim], tokens already permuted and grouped by experts - experts_indices: [B, S], expert index for each token - start_indices: start token index for each expert (in the permuted token space) - end_indices: end token index for each expert (in the permuted token space) - probs: probability vector for each token (used for unpermute) batch_size, seq_length: original batch size and sequence length Returns: @@ -378,6 +384,16 @@ def _generate_qkv_mot_opt( total_tokens, hidden_dim = hidden_states.shape device, dtype = hidden_states.device, hidden_states.dtype + if two_expert_layout is None: + two_expert_layout = prepare_two_expert_layout( + token_types, + batch_size=batch_size, + seq_length=seq_length, + action_starts=action_starts, + action_ends=action_ends, + device=device, + ) + expert_segments = two_expert_layout.expert_segments # Initialize Q/K/V buffers in the permuted token space q_buffer = torch.zeros(total_tokens, hidden_dim, device=device, dtype=dtype) @@ -396,7 +412,7 @@ def _generate_qkv_mot_opt( # === Each expert processes its own token slice === for expert_idx, qkv_proj in enumerate(self.qkv_proj_experts): - start, end = start_indices[expert_idx], end_indices[expert_idx] + start, end = expert_segments[expert_idx] if start == end: continue @@ -416,11 +432,15 @@ def _generate_qkv_mot_opt( k_buffer[start:end] = k_out v_buffer[start:end] = v_out - # === Restore tokens to the original order === - # unpermute (using the same unpermute operation) - q_unpermuted = unpermute(q_buffer, row_id_map, probs) - k_unpermuted = unpermute(k_buffer, row_id_map, probs) - v_unpermuted = unpermute(v_buffer, row_id_map, probs) + q_unpermuted = two_expert_unpack( + q_buffer, batch_size, seq_length, two_expert_layout + ) + k_unpermuted = two_expert_unpack( + k_buffer, batch_size, seq_length, two_expert_layout + ) + v_unpermuted = two_expert_unpack( + v_buffer, batch_size, seq_length, two_expert_layout + ) # === Reshape to final form === query_states = q_unpermuted.view( @@ -442,8 +462,8 @@ def _apply_rotary_pos_embed( query_states, key_states = m_rope( query_states.contiguous(), key_states.contiguous(), - cos[..., : (cos.size(3) // 2)].contiguous().float(), - sin[..., : (sin.size(3) // 2)].contiguous().float(), + cos.contiguous().float(), + sin.contiguous().float(), self.rope_scaling["mrope_section"], ) return query_states, key_states @@ -481,12 +501,14 @@ def _generate_output(self, attn_output, masks): def _generate_output_mot_opt( self, attn_output: torch.Tensor, - experts_indices: torch.Tensor, - start_indices: torch.Tensor, - end_indices: torch.Tensor, + action_starts: torch.Tensor, + action_ends: torch.Tensor, + expert_segments: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None, + two_expert_layout: Optional[TwoExpertLayout] = None, + token_types: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ - Expert-sharded version of attn_output processing based on start_indices / end_indices. + Expert-sharded version of attn_output processing based on fixed LLM/action segments. Rearranges the [B, S, H] attn_output according to expert order (permute), applies the o_proj projection for each expert individually, and keeps the final output in expert order ([Tokens, Hidden]) @@ -494,21 +516,26 @@ def _generate_output_mot_opt( Args: attn_output: [B, S, hidden_dim] - experts_indices: [B, S], expert index for each token - start_indices, end_indices: start and end token indices for each expert - (in the permuted token space) Returns: output_buffer: [TotalTokens, hidden_dim], arranged in expert order """ - _, _, hidden_dim = attn_output.shape + batch_size, seq_length, hidden_dim = attn_output.shape device, dtype = attn_output.device, attn_output.dtype + if two_expert_layout is None: + two_expert_layout = prepare_two_expert_layout( + token_types, + batch_size=batch_size, + seq_length=seq_length, + action_starts=action_starts, + action_ends=action_ends, + device=device, + ) + expert_segments = two_expert_layout.expert_segments - # === 1. Flatten and reorder by expert assignment === - flat_attn_output = attn_output.view(-1, hidden_dim) # [B*S, H] - flat_expert_indices = experts_indices.reshape(-1) # [B*S] - permuted_inputs, _ = permute(flat_attn_output, flat_expert_indices) + # === 1. Pack by fixed LLM/action layout, without sorting === + permuted_inputs = two_expert_pack(attn_output, two_expert_layout) total_tokens = permuted_inputs.shape[0] # === 2. Initialize output buffer (still in permuted token space) === @@ -518,7 +545,7 @@ def _generate_output_mot_opt( # === 3. Each expert processes its own token segment independently === for expert_idx, o_proj in enumerate(self.o_proj_experts): - start, end = start_indices[expert_idx], end_indices[expert_idx] + start, end = expert_segments[expert_idx] if start == end: continue @@ -558,11 +585,11 @@ def forward( position_embeddings: Optional[ Tuple[torch.Tensor, torch.Tensor] ] = None, # necessary, but kept here for BC - start_indices: Optional[torch.Tensor] = None, - end_indices: Optional[torch.Tensor] = None, - probs: Optional[torch.Tensor] = None, - row_id_map: Optional[torch.Tensor] = None, orig_shape: Optional[Tuple[int]] = None, + action_starts: Optional[torch.Tensor] = None, + action_ends: Optional[torch.Tensor] = None, + expert_segments: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None, + two_expert_layout: Optional[TwoExpertLayout] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if token_types is None: @@ -575,13 +602,13 @@ def forward( bsz, q_len, _ = orig_shape query_states, key_states, value_states = self._generate_qkv_mot_opt( hidden_states, - token_types, - start_indices, - end_indices, - probs, - row_id_map, bsz, q_len, + action_starts, + action_ends, + expert_segments, + two_expert_layout, + token_types, ) else: bsz, q_len, _ = hidden_states.size() @@ -654,7 +681,12 @@ def forward( if self.config.mot_opt: output = self._generate_output_mot_opt( - attn_output, token_types, start_indices, end_indices + attn_output, + action_starts, + action_ends, + expert_segments, + two_expert_layout, + token_types, ) else: output = self._generate_output(attn_output, masks) @@ -679,11 +711,11 @@ def forward( position_embeddings: Optional[ Tuple[torch.Tensor, torch.Tensor] ] = None, # necessary, but kept here for BC - start_indices: Optional[torch.Tensor] = None, - end_indices: Optional[torch.Tensor] = None, - probs: Optional[torch.Tensor] = None, - row_id_map: Optional[torch.Tensor] = None, orig_shape: Optional[Tuple[int]] = None, + action_starts: Optional[torch.Tensor] = None, + action_ends: Optional[torch.Tensor] = None, + expert_segments: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None, + two_expert_layout: Optional[TwoExpertLayout] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if token_types is None: raise ValueError("token_types must not be empty") @@ -696,13 +728,13 @@ def forward( bsz, q_len, _ = orig_shape query_states, key_states, value_states = self._generate_qkv_mot_opt( hidden_states, - token_types, - start_indices, - end_indices, - probs, - row_id_map, bsz, q_len, + action_starts, + action_ends, + expert_segments, + two_expert_layout, + token_types, ) else: bsz, q_len, _ = hidden_states.size() @@ -777,7 +809,12 @@ def forward( if self.config.mot_opt: output = self._generate_output_mot_opt( - attn_output, token_types, start_indices, end_indices + attn_output, + action_starts, + action_ends, + expert_segments, + two_expert_layout, + token_types, ) else: output = self._generate_output(attn_output, masks) @@ -803,11 +840,11 @@ def forward( position_embeddings: Optional[ Tuple[torch.Tensor, torch.Tensor] ] = None, # necessary, but kept here for BC - start_indices: Optional[torch.Tensor] = None, - end_indices: Optional[torch.Tensor] = None, - probs: Optional[torch.Tensor] = None, - row_id_map: Optional[torch.Tensor] = None, orig_shape: Optional[Tuple[int]] = None, + action_starts: Optional[torch.Tensor] = None, + action_ends: Optional[torch.Tensor] = None, + expert_segments: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None, + two_expert_layout: Optional[TwoExpertLayout] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if token_types is None: raise ValueError("token_types must not be empty") @@ -820,13 +857,13 @@ def forward( bsz, q_len, _ = orig_shape query_states, key_states, value_states = self._generate_qkv_mot_opt( hidden_states, - token_types, - start_indices, - end_indices, - probs, - row_id_map, bsz, q_len, + action_starts, + action_ends, + expert_segments, + two_expert_layout, + token_types, ) else: bsz, q_len, _ = hidden_states.size() @@ -907,7 +944,12 @@ def forward( if self.config.mot_opt: output = self._generate_output_mot_opt( - attn_output, token_types, start_indices, end_indices + attn_output, + action_starts, + action_ends, + expert_segments, + two_expert_layout, + token_types, ) else: output = self._generate_output(attn_output, masks) diff --git a/wall_x/model/core/attention/mask.py b/wall_x/model/core/attention/mask.py deleted file mode 100644 index ec446083..00000000 --- a/wall_x/model/core/attention/mask.py +++ /dev/null @@ -1,259 +0,0 @@ -import torch - - -def find_first_last_ones(tensor): - """ - Input: tensor of shape (bs, seq_len) containing 0s and 1s - Output: (first_indices, last_indices), each of shape (bs,) - first_indices[i] is the first index of 1 in batch i, or -1 if none exists - last_indices[i] is the last index of 1 in batch i, or -1 if none exists - """ - bs, seq_len = tensor.shape - masks = tensor == 1 - has_ones = masks.any(dim=1) - - first = torch.full((bs,), -1, dtype=torch.long, device=tensor.device) - last = first.clone() - - # Compute the first index of 1 - first[has_ones] = torch.argmax(masks[has_ones].float(), dim=1) - - # Compute the last index of 1 - flipped_masks = masks.flip(dims=[1]) - last_argmax = torch.argmax(flipped_masks[has_ones].float(), dim=1) - last[has_ones] = seq_len - 1 - last_argmax - - return first, last - - -def update_position_ids(position_ids, moe_token_types, positional_masks): - """Extracted from ActionModelMixMin._update_position_ids (was @staticmethod).""" - if positional_masks is None or "ar_predict_token_positions" not in positional_masks: - return position_ids - - new_position_ids = position_ids.clone() - ar_predict_token_positions = positional_masks["ar_predict_token_positions"] - flow_mask = moe_token_types == 1 - - start_ar_pos, end_ar_pos = find_first_last_ones(ar_predict_token_positions) - start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) - - for bs_i in range(position_ids.shape[1]): - if start_ar_pos[bs_i] != -1 and end_ar_pos[bs_i] != -1: - start_ar_ids = new_position_ids[:, bs_i, start_ar_pos[bs_i]] - start_flow_ids = new_position_ids[:, bs_i, start_flow_pos[bs_i]] - diff = start_flow_ids - start_ar_ids - new_position_ids[:, bs_i, start_flow_pos[bs_i] :] = position_ids[ - :, bs_i, start_flow_pos[bs_i] : - ] - diff.unsqueeze(-1) - - return new_position_ids - - -def update_joint_attention_mask_2d( - attention_mask, - moe_token_types, - positional_masks, - causal_action_attention_mask=False, -): - """Extracted from ActionModelMixMin._update_joint_attention_mask_2d. - The only self attribute used was self.config.causal_action_attention_mask, now passed as parameter. - """ - if attention_mask.dim() == 3: # bs, seq_len, seq_len - return attention_mask - - bs, seq_len = moe_token_types.shape[0], moe_token_types.shape[1] - # Create a lower-triangular causal mask - causal_mask = torch.tril( - torch.ones( - (seq_len, seq_len), dtype=torch.bfloat16, device=moe_token_types.device - ) - ) - # Expand to the batch dimension - attention_mask = causal_mask.unsqueeze(0).expand(bs, -1, -1) - - if positional_masks is not None and "padding_positions" in positional_masks: - padding_positions = positional_masks["padding_positions"] - # Set padding rows to zero - attention_mask = torch.where( - padding_positions[:, None, :], - torch.zeros_like(attention_mask), - attention_mask, - ) - # Set padding columns to zero - attention_mask = torch.where( - padding_positions[:, :, None], - torch.zeros_like(attention_mask), - attention_mask, - ) - - # Set the moe1 region to 1 and mask it from the fast region - moe1_mask = (moe_token_types[:, :, None]) & (moe_token_types[:, None, :]) - - if ( - not causal_action_attention_mask - ): # If causal action attention mask is disabled, set the whole moe1 region to 1 - attention_mask = torch.where( - moe1_mask, torch.ones_like(attention_mask), attention_mask - ) - - if ( - positional_masks is not None - and "ar_predict_token_positions" in positional_masks - ): - ar_predict_token_positions = positional_masks["ar_predict_token_positions"] - moe1_mask = (moe_token_types[:, :, None]) & ( - ar_predict_token_positions[:, None, :] - ) - attention_mask = torch.where( - moe1_mask, torch.zeros_like(attention_mask), attention_mask - ) - - if ( - positional_masks is not None - and "valid_flow_action_positions" in positional_masks - ): - # true in moe_token_types but false in valid_flow_action_positions - nonvalid_flow_action_positions = ( - moe_token_types & ~positional_masks["valid_flow_action_positions"] - ) - attention_mask = torch.where( - nonvalid_flow_action_positions[:, None, :], - torch.zeros_like(attention_mask), - attention_mask, - ) - attention_mask = torch.where( - nonvalid_flow_action_positions[:, :, None], - torch.zeros_like(attention_mask), - attention_mask, - ) - - # AR and flow are bidirectional - if positional_masks is not None and "ar_action_mask" in positional_masks: - ar_action_mask = positional_masks["ar_action_mask"] != 0 - flow_positions = moe_token_types == 1 - if positional_masks.get("ar_visible", True): - flow_ar_position = ar_action_mask | flow_positions - flow_ar_mask = flow_ar_position[:, :, None] & flow_ar_position[:, None, :] - attention_mask = torch.where( - flow_ar_mask, torch.ones_like(attention_mask), attention_mask - ) - else: - flow_flow_mask = flow_positions[:, :, None] & flow_positions[:, None, :] - ar_ar_mask = ar_action_mask[:, :, None] & ar_action_mask[:, None, :] - flow_ar_mask = flow_flow_mask | ar_ar_mask - - affected = ar_action_mask | flow_positions # (B, N) - affected_pair = affected[:, :, None] & affected[:, None, :] - attention_mask = attention_mask.masked_fill(affected_pair, 0) - - attention_mask = torch.where( - flow_ar_mask, torch.ones_like(attention_mask), attention_mask - ) - - return attention_mask - - -def update_joint_attention_flash_mask( - attention_mask, - moe_token_types, - positional_masks, - causal_action_attention_mask=False, - debug=False, -): - """Extracted from ActionModelMixMin._update_joint_attention_flash_mask. - The only self attribute used was self.config.causal_action_attention_mask, now passed as parameter. - """ - device = moe_token_types.device - B, S = moe_token_types.shape - i32 = torch.int32 - - # ---- Return-vector initialization ---- - LTS = torch.ones((B, S), device=device, dtype=i32) * S - UTE = torch.arange(S, device=device, dtype=i32).unsqueeze(0).expand(B, S).clone() - - # Handle padding positions - if positional_masks is not None and "padding_positions" in positional_masks: - padding_positions = positional_masks["padding_positions"] - LTS[padding_positions] = 0 - UTE[padding_positions] = S - - # Handle AR predict tokens - if ( - positional_masks is not None - and "ar_predict_token_positions" in positional_masks - ): - start_ar_pos, end_ar_pos = find_first_last_ones( - positional_masks["ar_predict_token_positions"] - ) - for bs_i in range(B): - if end_ar_pos[bs_i] != -1: - LTS[bs_i, positional_masks["ar_predict_token_positions"][bs_i]] = ( - end_ar_pos[bs_i].to(i32) + 1 - ) - - flow_mask = moe_token_types == 1 - start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) - if positional_masks is None or "ar_action_mask" not in positional_masks: - # Handle bidirectional flow action masks - if not causal_action_attention_mask: - for bs_i in range(B): - if start_flow_pos[bs_i] != -1: - UTE[bs_i, flow_mask[bs_i]] = start_flow_pos[bs_i].to(i32) - else: - # AR and flow are bidirectional - ar_action_mask = positional_masks["ar_action_mask"] != 0 - flow_mask = moe_token_types == 1 - if positional_masks.get("ar_visible", True): - flow_ar_position = ar_action_mask | flow_mask - for bs_i in range(B): - idx = flow_ar_position[bs_i].nonzero(as_tuple=True)[0] - if idx.numel() == 0: - continue - block_start = idx.min() - block_end = idx.max() + 1 - # Set the visible range of every token in the block to [block_start, block_end) - UTE[bs_i, idx] = block_start.to(i32) - LTS[bs_i, idx] = block_end.to(i32) - else: - for bs_i in range(B): - # 1) Flow sub-block: flow attends only to flow - flow_idx = flow_mask[bs_i].nonzero(as_tuple=True)[0] - if flow_idx.numel() > 0: - flow_start = flow_idx.min() - flow_end = flow_idx.max() + 1 - # Only flow-token rows are set to [flow_start, flow_end) - UTE[bs_i, flow_idx] = flow_start.to(i32) - LTS[bs_i, flow_idx] = flow_end.to(i32) - - # 2) AR sub-block: AR attends only to AR - ar_idx = ar_action_mask[bs_i].nonzero(as_tuple=True)[0] - if ar_idx.numel() > 0: - ar_start = ar_idx.min() - ar_end = ar_idx.max() + 1 - # Only AR-token rows are set to [ar_start, ar_end) - UTE[bs_i, ar_idx] = ar_start.to(i32) - LTS[bs_i, ar_idx] = ar_end.to(i32) - - # Handle validation flow - if ( - positional_masks is not None - and "valid_flow_action_positions" in positional_masks - ): - flow_mask = moe_token_types == 1 - nonvalid_flow_action_positions = ( - flow_mask & ~positional_masks["valid_flow_action_positions"] - ) - if nonvalid_flow_action_positions.any(): - LTS[nonvalid_flow_action_positions] = 0 - UTE[nonvalid_flow_action_positions] = S - - LTS = LTS.unsqueeze(-1) - UTE = UTE.unsqueeze(-1) - - startend_row_indices = torch.cat([LTS, UTE], dim=-1) - - # add num_heads dimension - startend_row_indices = startend_row_indices.unsqueeze(1) - - return startend_row_indices diff --git a/wall_x/model/core/attention/selector.py b/wall_x/model/core/attention/selector.py index a288935e..6cda10c4 100644 --- a/wall_x/model/core/attention/selector.py +++ b/wall_x/model/core/attention/selector.py @@ -1,9 +1,11 @@ -from typing import Dict, Optional, Union +from typing import Optional, Union, Dict -import torch from packaging import version + from transformers.modeling_utils import AttentionInterface -from transformers.utils import is_torch_xla_available, logging +from transformers.utils import logging, is_torch_xla_available + +import torch ALL_ATTENTION_FUNCTIONS: AttentionInterface = AttentionInterface() logger = logging.get_logger(__name__) diff --git a/wall_x/model/core/ops/README.md b/wall_x/model/core/ops/README.md new file mode 100644 index 00000000..c82e6fdf --- /dev/null +++ b/wall_x/model/core/ops/README.md @@ -0,0 +1,71 @@ +# `wall_x.model.core.ops` — operator boundary contract + +This package is the **single boundary** between wall-x and the internal +`external_accel` kernel library. It is also a **public (open-source) module**: the +open-source export ships `wall_x/model/core/ops/` but **not** `external_accel`, so the +PyTorch fallback in each operator is the code path open-source users actually +run. That makes the contract below load-bearing, not stylistic. + +## The contract + +Every operator that can be accelerated by an external_accel kernel MUST: + +1. **Touch external_accel only here.** No module outside `core/ops/` may + `import external_accel` (directly or via `from external_accel import ...`). Higher layers + import the operator from `wall_x.model.core.ops`. Two legacy exceptions + (`model/pi0/*`) predate this rule and are tracked for cleanup. + +2. **Import cleanly without external_accel.** Kernel access is wrapped so the module + imports on a machine with no external_accel: + - `OpsProxy` subclasses resolve lazily via `importlib.import_module`. + - Plain-function ops use `try: from external_accel import except Exception: = None`. + The `from external_accel import` line must be indented under `try` (never column-0), + so the open-source export grep guard passes. + +3. **Own a self-contained PyTorch fallback.** The fallback lives **in this + package**, computes the same result as the kernel, and is a **first-class, + independently callable** entry — never buried in a caller module. Rationale: + - open-source runs the fallback (no external_accel), so it must be complete here; + - the fallback must be selectable in isolation for the bitexact test below. + +4. **Select backend by capability — no env switches.** An op uses its kernel + iff the loaded external_accel build provides it (and inputs are on CUDA); otherwise + it runs the fallback. There are **no per-op switches and no global disable + flag** — capability detection ("kernel present ⇒ use it") is the only + mechanism. To exercise the fallback, run without external_accel (the open-source + env) or, in a test, call the fallback directly / null the module's kernel + handle. + +5. **Be covered by `tests/test_ops_fallback.py`** with two assertions: + - **fallback vs reference** — the PyTorch fallback matches an independent + reference (this runs in the open-source overlay, guaranteeing OSS + correctness with no external_accel present); + - **kernel vs fallback** (`test_vs_external_accel`) — when external_accel is present, the + kernel and the fallback agree. Use an explicit tolerance for + native-dtype/fused rounding (fp32 ~2e-6, fp16 ~2e-2, bf16 ~3e-2); a + divergence far outside tolerance is a **convention/layout bug**, not a + rounding artifact — fix the operator, do not widen the tolerance. + +## Two conforming shapes + +- **`OpsProxy` subclass** (rmsnorm, rope, m_rope, silu_and_mul, …): define + `_external_accel_name` and `_pytorch_fallback`; get `__call__` dispatch, + `call_with_backend("pytorch"|"external_accel")`, and `available_backends()` for free. + Preferred for single tensor→tensor ops. + +- **Plain-function module** (gated_delta, qkv_norm_replace, attention_mask): + when the op is not a single callable (multiple entry points, or a fused op + with a structured signature), expose `(...)` (capability-dispatched) and a + directly-callable `_pytorch(...)` fallback in the same module, plus a + `_available(...)` probe. The fallback must be a pure tensor function (it + may take a `module` for parameters but must not embed caller-only control + flow), so the bitexact test can call kernel and fallback side by side. + +## Anti-patterns (do not ship) + +- Fallback logic living in a caller module (e.g. an attention forward) while the + ops module only wraps the kernel — the op then has no OSS-runnable, testable + fallback. +- Any env switch (per-op or global) to force the fallback — capability + detection covers every real need; nulling the kernel handle covers tests. +- Widening a `test_vs_external_accel` tolerance to hide a large divergence. diff --git a/wall_x/model/core/ops/__init__.py b/wall_x/model/core/ops/__init__.py index 1e4a6ad2..6466b4d0 100644 --- a/wall_x/model/core/ops/__init__.py +++ b/wall_x/model/core/ops/__init__.py @@ -1,17 +1,35 @@ """Operator proxy layer with runtime fallback backends.""" +from wall_x.model.core.ops.norm import rmsnorm, layernorm, add_rmsnorm, gated_rmsnorm +from wall_x.model.core.ops.adarmsnorm import adarmsnorm +from wall_x.model.core.ops.activation import silu_and_mul +from wall_x.model.core.ops.position import pos_embed_interp +from wall_x.model.core.ops.rope import rope, m_rope, rot_pos_emb +from wall_x.model.core.ops.moe import ( + permute, + unpermute, + two_expert_pack, + two_expert_unpack, +) from wall_x.model.core.ops.index import get_rope_index, get_window_index -from wall_x.model.core.ops.moe import permute, unpermute -from wall_x.model.core.ops.norm import rmsnorm -from wall_x.model.core.ops.rope import m_rope, rope, rot_pos_emb +from wall_x.model.core.ops.qwen35_index import qwen35_get_rope_index __all__ = [ "rmsnorm", + "layernorm", + "adarmsnorm", + "add_rmsnorm", + "gated_rmsnorm", + "silu_and_mul", + "pos_embed_interp", "rope", "m_rope", "rot_pos_emb", "permute", "unpermute", + "two_expert_pack", + "two_expert_unpack", "get_rope_index", + "qwen35_get_rope_index", "get_window_index", ] diff --git a/wall_x/model/core/ops/_cuda_wrappers.py b/wall_x/model/core/ops/_cuda_wrappers.py index 0fef8a23..269d841f 100644 --- a/wall_x/model/core/ops/_cuda_wrappers.py +++ b/wall_x/model/core/ops/_cuda_wrappers.py @@ -56,6 +56,56 @@ def backward(ctx, grad_q_embed, grad_k_embed): return grad_q, grad_k, None, None, None, None, None +class _RopeFullFunction(Function): + @staticmethod + def forward( + ctx, + q, + k, + cos, + sin, + interleave, + rotary_dim, + fwd_kernel=None, + bwd_kernel=None, + ): + if q.stride(-1) != 1: + q = q.contiguous() + if k.stride(-1) != 1: + k = k.contiguous() + q_embed = q.new_empty(q.shape) + k_embed = k.new_empty(k.shape) + ctx.save_for_backward(cos, sin) + ctx.interleave = interleave + ctx.rotary_dim = rotary_dim + ctx.bwd_kernel = bwd_kernel + fwd_kernel(q, k, q_embed, k_embed, cos, sin, interleave, rotary_dim) + return q_embed, k_embed + + @staticmethod + def backward(ctx, grad_q_embed, grad_k_embed): + if grad_q_embed is None or grad_k_embed is None: + return None, None, None, None, None, None, None, None + if grad_q_embed.stride(-1) != 1: + grad_q_embed = grad_q_embed.contiguous() + if grad_k_embed.stride(-1) != 1: + grad_k_embed = grad_k_embed.contiguous() + cos, sin = ctx.saved_tensors + grad_q = grad_q_embed.new_empty(grad_q_embed.shape) + grad_k = grad_k_embed.new_empty(grad_k_embed.shape) + ctx.bwd_kernel( + grad_q_embed, + grad_k_embed, + grad_q, + grad_k, + cos, + sin, + ctx.interleave, + ctx.rotary_dim, + ) + return grad_q, grad_k, None, None, None, None, None, None + + class _RopePackFunction(Function): @staticmethod def forward( @@ -102,10 +152,12 @@ class Rope: def __init__(self): pass self.fwd_kernel = _m().rope + self.full_kernel = _m().rope_full self.inplace_kernel = _m().rope_inplace self.pack_kernel = _m().rope_inplace_pack self.pack_bwd_kernel = _m().rope_inplace_pack_bwd self.bwd_kernel = _m().rope_bwd + self.full_bwd_kernel = _m().rope_full_bwd def pack( self, @@ -117,11 +169,11 @@ def pack( interleave=False, inference=False, ): - """packed qkvq/krope,qkv. + """原地对packed qkv的q/k部分应用rope，并返回同一个qkv张量。 Split/view/offset calculations are done on the C++ side. Callers consume qkv through the paired flash_attn pack interface without splitting q/k/v in Python. - autograd,backwarddqkvdq/dk. + 训练路径支持autograd，backward会原地更新dqkv中的dq/dk部分。 Args: qkv: [seq_len, q_dim + 2*kv_dim] @@ -129,7 +181,7 @@ def pack( kv_num_heads: number of key/value heads cos, sin: rotary position encoding [1, seq_len, head_dim/2] interleave: whether to use interleaved RoPE - inference: kernel,autograd + inference: 只执行前向原地kernel，不接入autograd """ if inference: self.pack_kernel(qkv, cos, sin, q_num_heads, kv_num_heads, interleave) @@ -149,10 +201,30 @@ def pack( def pack_backward( self, dqkv, q_num_heads, kv_num_heads, cos, sin, interleave=False ): - """packed dqkvdq/dkrope backward,dqkv.""" + """原地对packed dqkv的dq/dk部分应用rope backward，并返回dqkv。""" self.pack_bwd_kernel(dqkv, cos, sin, q_num_heads, kv_num_heads, interleave) return dqkv + def full(self, q, k, cos, sin, interleave=False, rotary_dim=None): + """Apply RoPE to the rotary prefix and return full-head q/k outputs. + + This preserves the existing ``rope(q_rot, k_rot, ...)`` API while + avoiding Python-side ``torch.cat([rotated, pass_through], dim=-1)`` for + partial rotary embeddings. + """ + if rotary_dim is None: + rotary_dim = cos.shape[-1] * 2 + return _RopeFullFunction.apply( + q, + k, + cos, + sin, + interleave, + int(rotary_dim), + self.full_kernel, + self.full_bwd_kernel, + ) + def __call__(self, q, k, cos, sin, interleave=False, inference=False): if inference: self.inplace_kernel(q, k, cos, sin, interleave) @@ -275,6 +347,7 @@ def __init__(self): self.inplace_kernel = _m().m_rope_inplace self.pack_kernel = _m().m_rope_inplace_pack self.pack_bwd_kernel = _m().m_rope_inplace_pack_bwd + self.cfg_cache_write_kernel = _m().m_rope_inplace_cfg_cache_write self.bwd_kernel = _m().m_rope_bwd def pack( @@ -287,7 +360,7 @@ def pack( mrope_section, inference=False, ): - """packed qkvq/kmrope,qkv.GQA. + """原地对packed qkv的q/k部分应用mrope，并返回同一个qkv张量。支持GQA。 Split/view/offset calculations are done on the C++ side. qkv must be 3D: [bz, seq_len, q_dim+2*kv_dim]. @@ -297,7 +370,7 @@ def pack( kv_num_heads: number of key/value heads cos, sin: rotary position encoding [3, bz, seq_len, head_dim/2] mrope_section: (first, second) M-RoPE section parameters - inference: kernel,autograd + inference: 只执行前向原地kernel，不接入autograd """ if inference: self.pack_kernel( @@ -323,7 +396,7 @@ def pack( ) def pack_backward(self, dqkv, q_num_heads, kv_num_heads, cos, sin, mrope_section): - """packed dqkvdq/dkmrope backward,dqkv.""" + """原地对packed dqkv的dq/dk部分应用mrope backward，并返回dqkv。""" self.pack_bwd_kernel( dqkv, cos, @@ -335,6 +408,42 @@ def pack_backward(self, dqkv, q_num_heads, kv_num_heads, cos, sin, mrope_section ) return dqkv + def cfg_cache_write( + self, + q, + k, + v, + cos, + sin, + mrope_section, + kvcache, + valid0, + valid1, + common_len, + cond_tail_len, + ): + """In-place mrope(q, k) fused with CFG packed prefill KV-cache write. + + This is inference-only. `q`, `k`, and `v` are packed single-batch + tensors [1, packed_seq, heads, head_dim]. `kvcache` is the per-layer + static cache view [2, max_batch, max_len, kv_heads, head_dim]. + """ + self.cfg_cache_write_kernel( + q, + k, + v, + cos, + sin, + kvcache, + valid0, + valid1, + int(common_len), + int(cond_tail_len), + mrope_section[0], + mrope_section[1], + ) + return q, k + def __call__(self, q, k, cos, sin, mrope_section, inference=False): if inference: first = mrope_section[0] @@ -359,7 +468,9 @@ class RotPos: def __init__(self): pass - def __call__(self, inv_freq, grid_thw, spatial_merge_size): + def __call__( + self, inv_freq, grid_thw, spatial_merge_size, total_output_tokens=None + ): assert ( inv_freq.dtype == torch.float32 ), f"Expected float32, got {inv_freq.dtype}" @@ -380,8 +491,11 @@ def __call__(self, inv_freq, grid_thw, spatial_merge_size): dim=0, ).to(grid_thw.dtype) + if total_output_tokens is None: + total_output_tokens = cumsum_tokens[-1].item() + output = torch.empty( - (cumsum_tokens[-1].item(), inv_freq.size(0) * 2), + (total_output_tokens, inv_freq.size(0) * 2), dtype=torch.float, device=inv_freq.device, ) @@ -748,3 +862,629 @@ def permute_kernel(input_act, indices, num_out_tokens=None, max_token_num=0): def unpermute_kernel(input_act, row_id_map, probs=None): return UnpermuteMoE_topK.apply(input_act, row_id_map, probs) + + + + + +def _check_index_tensor( + tensor: torch.Tensor, + name: str, + *, + expected_size: int, + device: torch.device, +): + if not isinstance(tensor, torch.Tensor): + raise RuntimeError(f"{name} must be a tensor") + if tensor.dim() != 1: + raise RuntimeError(f"{name} must have shape [{expected_size}]") + if tensor.numel() != expected_size: + raise RuntimeError( + f"{name} size mismatch: got {tensor.numel()}, expected {expected_size}" + ) + if tensor.dtype != torch.int64: + raise RuntimeError(f"{name} must be int64") + if tensor.device != device: + raise RuntimeError(f"{name} must be on the same device as input") + if not tensor.is_contiguous(): + raise RuntimeError(f"{name} must be contiguous") + + +def _check_layout_tensors( + action_starts: torch.Tensor, + action_ends: torch.Tensor, + llm_offsets: torch.Tensor, + action_offsets: torch.Tensor, + *, + batch_size: int, + device: torch.device, +): + _check_index_tensor( + action_starts, + "action_starts", + expected_size=batch_size, + device=device, + ) + _check_index_tensor( + action_ends, + "action_ends", + expected_size=batch_size, + device=device, + ) + _check_index_tensor( + llm_offsets, + "llm_offsets", + expected_size=batch_size + 1, + device=device, + ) + _check_index_tensor( + action_offsets, + "action_offsets", + expected_size=batch_size + 1, + device=device, + ) + + +class _TwoExpertPack(torch.autograd.Function): + @staticmethod + def forward( + ctx, + input_act: torch.Tensor, + action_starts: torch.Tensor, + action_ends: torch.Tensor, + llm_offsets: torch.Tensor, + action_offsets: torch.Tensor, + ): + if input_act.is_cpu: + raise RuntimeError("two_expert_pack expects a CUDA tensor") + if input_act.dim() != 3: + raise RuntimeError("two_expert_pack expects input shape [B, S, D]") + + batch_size, seq_length, hidden_dim = input_act.shape + _check_layout_tensors( + action_starts, + action_ends, + llm_offsets, + action_offsets, + batch_size=batch_size, + device=input_act.device, + ) + + if not input_act.is_contiguous(): + input_act = input_act.contiguous() + + output = torch.empty( + (batch_size * seq_length, hidden_dim), + dtype=input_act.dtype, + device=input_act.device, + ) + _m().two_expert_pack( + output, input_act, action_starts, action_ends, llm_offsets, action_offsets + ) + + ctx.batch_size = batch_size + ctx.seq_length = seq_length + ctx.save_for_backward(action_starts, action_ends, llm_offsets, action_offsets) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + if not ctx.needs_input_grad[0]: + return None, None, None, None, None + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + action_starts, action_ends, llm_offsets, action_offsets = ctx.saved_tensors + grad_input = torch.empty( + (ctx.batch_size, ctx.seq_length, grad_output.size(1)), + dtype=grad_output.dtype, + device=grad_output.device, + ) + _m().two_expert_unpack( + grad_input, + grad_output, + ctx.batch_size, + ctx.seq_length, + action_starts, + action_ends, + llm_offsets, + action_offsets, + ) + return grad_input, None, None, None, None + + +class _TwoExpertUnpack(torch.autograd.Function): + @staticmethod + def forward( + ctx, + input_act: torch.Tensor, + batch_size: int, + seq_length: int, + action_starts: torch.Tensor, + action_ends: torch.Tensor, + llm_offsets: torch.Tensor, + action_offsets: torch.Tensor, + ): + if input_act.is_cpu: + raise RuntimeError("two_expert_unpack expects a CUDA tensor") + if input_act.dim() != 2: + raise RuntimeError("two_expert_unpack expects input shape [B*S, D]") + + batch_size = int(batch_size) + seq_length = int(seq_length) + if input_act.size(0) != batch_size * seq_length: + raise RuntimeError( + f"packed rows mismatch: got {input_act.size(0)}, " + f"expected {batch_size * seq_length}" + ) + _check_layout_tensors( + action_starts, + action_ends, + llm_offsets, + action_offsets, + batch_size=batch_size, + device=input_act.device, + ) + + if not input_act.is_contiguous(): + input_act = input_act.contiguous() + + output = torch.empty( + (batch_size, seq_length, input_act.size(1)), + dtype=input_act.dtype, + device=input_act.device, + ) + _m().two_expert_unpack( + output, + input_act, + batch_size, + seq_length, + action_starts, + action_ends, + llm_offsets, + action_offsets, + ) + + ctx.save_for_backward(action_starts, action_ends, llm_offsets, action_offsets) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + if not ctx.needs_input_grad[0]: + return None, None, None, None, None, None, None + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + batch_size, seq_length, hidden_dim = grad_output.shape + action_starts, action_ends, llm_offsets, action_offsets = ctx.saved_tensors + grad_input = torch.empty( + (batch_size * seq_length, hidden_dim), + dtype=grad_output.dtype, + device=grad_output.device, + ) + _m().two_expert_pack( + grad_input, + grad_output, + action_starts, + action_ends, + llm_offsets, + action_offsets, + ) + return grad_input, None, None, None, None, None, None + + +def two_expert_pack_kernel( + input_act: torch.Tensor, + action_starts: torch.Tensor, + action_ends: torch.Tensor, + llm_offsets: torch.Tensor, + action_offsets: torch.Tensor, +) -> torch.Tensor: + return _TwoExpertPack.apply( + input_act, action_starts, action_ends, llm_offsets, action_offsets + ) + + +def two_expert_unpack_kernel( + input_act: torch.Tensor, + batch_size: int, + seq_length: int, + action_starts: torch.Tensor, + action_ends: torch.Tensor, + llm_offsets: torch.Tensor, + action_offsets: torch.Tensor, +) -> torch.Tensor: + return _TwoExpertUnpack.apply( + input_act, + batch_size, + seq_length, + action_starts, + action_ends, + llm_offsets, + action_offsets, + ) + + + + +def _get_sm_count(): + """Cache SM count for the current GPU.""" + if not hasattr(_get_sm_count, "_count"): + _get_sm_count._count = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _get_sm_count._count + + +def _get_rmsnorm_backward_blocks(num_tokens: int) -> int: + return min(_get_sm_count() * 2, num_tokens) + + +def _is_rmsnorm_kernel_layout(tensor: torch.Tensor) -> bool: + if tensor.dim() < 1 or tensor.stride(-1) != 1: + return False + if tensor.dim() == 1: + return True + + hidden_size = tensor.size(-1) + row_stride = tensor.stride(-2) + if row_stride < hidden_size: + return False + + expected_stride = row_stride + for dim in range(tensor.dim() - 2, -1, -1): + if tensor.stride(dim) != expected_stride: + return False + expected_stride *= tensor.size(dim) + return True + + +def _force_kernel_layout(tensor: torch.Tensor) -> torch.Tensor: + """Return a tensor the RMSNorm kernel accepts (fully-packed row layout). + + ``.contiguous()`` is NOT sufficient: a leading size-1 dim (e.g. a + ``[1, S/sp, H]`` sequence-parallel shard of a ``[1, S, H]`` tensor) keeps a + stale outer stride yet torch still reports ``is_contiguous()==True``, so + ``.contiguous()`` is a no-op while the kernel's strict + ``stride(i) == expected_stride`` check still fails. Flatten the leading dims + and re-expand to force genuine packed strides.""" + if _is_rmsnorm_kernel_layout(tensor): + return tensor + flat = tensor.reshape(-1, tensor.shape[-1]).contiguous() + return flat.reshape(tensor.shape) + + +class _RmsnormFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, weight, epsilon, zero_centered): + input = _force_kernel_layout(input) + output = torch.empty_like(input) + if zero_centered: + _m().rmsnorm_zero_centered(output, input, weight, epsilon) + else: + _m().rmsnorm(output, input, weight, epsilon) + ctx.save_for_backward(input, weight) + ctx.epsilon = epsilon + ctx.zero_centered = zero_centered + return output + + @staticmethod + def backward(ctx, grad_output): + input, weight = ctx.saved_tensors + epsilon = ctx.epsilon + zero_centered = ctx.zero_centered + grad_output = _force_kernel_layout(grad_output) + hidden_size = weight.numel() + num_tokens = input.numel() // hidden_size + + dx = torch.empty_like(input) + dw_float = torch.empty(hidden_size, dtype=torch.float32, device=input.device) + + # Pre-allocate partial_dw buffer: [num_blocks, hidden_size] + num_blocks = _get_rmsnorm_backward_blocks(num_tokens) + partial_dw = torch.empty( + num_blocks, hidden_size, dtype=torch.float32, device=input.device + ) + + if zero_centered: + _m().rmsnorm_zero_centered_backward( + dx, dw_float, partial_dw, grad_output, input, weight, epsilon + ) + else: + _m().rmsnorm_backward( + dx, dw_float, partial_dw, grad_output, input, weight, epsilon + ) + + dw = dw_float.to(weight.dtype) + return dx, dw, None, None + + +def _prepare_add_rmsnorm_inputs(residual, input): + """Shared prologue for the autograd and inference add_rmsnorm paths: + validate shapes, force kernel layout, allocate the two outputs.""" + if residual.shape != input.shape: + raise ValueError( + f"add_rmsnorm residual shape {tuple(residual.shape)} " + f"does not match input shape {tuple(input.shape)}" + ) + residual = _force_kernel_layout(residual) + input = _force_kernel_layout(input) + return residual, input, torch.empty_like(input), torch.empty_like(input) + + +class _AddRmsnormFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, residual, input, weight, epsilon, zero_centered): + residual, input, sum_output, norm_output = _prepare_add_rmsnorm_inputs( + residual, input + ) + hidden_size = weight.numel() + num_tokens = input.numel() // hidden_size + rstd = torch.empty(num_tokens, dtype=torch.float32, device=input.device) + if zero_centered: + _m().add_rmsnorm_zero_centered_with_rstd( + sum_output, norm_output, rstd, residual, input, weight, epsilon + ) + else: + _m().add_rmsnorm_with_rstd( + sum_output, norm_output, rstd, residual, input, weight, epsilon + ) + ctx.save_for_backward(sum_output, weight, rstd) + ctx.epsilon = epsilon + ctx.zero_centered = zero_centered + return sum_output, norm_output + + @staticmethod + def backward(ctx, grad_sum_output, grad_norm_output): + sum_output, weight, rstd = ctx.saved_tensors + epsilon = ctx.epsilon + zero_centered = ctx.zero_centered + hidden_size = weight.numel() + + if grad_norm_output is None: + d_rms = None + dw = None + else: + grad_norm_output = _force_kernel_layout(grad_norm_output) + + num_tokens = sum_output.numel() // hidden_size + dw_float = torch.empty( + hidden_size, dtype=torch.float32, device=sum_output.device + ) + + use_rstd_backward = grad_sum_output is not None and hidden_size <= 4096 + num_blocks = _get_rmsnorm_backward_blocks(num_tokens) + partial_dw = torch.empty( + num_blocks, hidden_size, dtype=torch.float32, device=sum_output.device + ) + + if use_rstd_backward: + grad_sum_output = _force_kernel_layout(grad_sum_output) + d_add = torch.empty_like(sum_output) + if zero_centered: + _m().add_rmsnorm_zero_centered_backward_with_rstd( + d_add, + dw_float, + partial_dw, + grad_sum_output, + grad_norm_output, + sum_output, + weight, + rstd, + ) + else: + _m().add_rmsnorm_backward_with_rstd( + d_add, + dw_float, + partial_dw, + grad_sum_output, + grad_norm_output, + sum_output, + weight, + rstd, + ) + return d_add, d_add, dw_float.to(weight.dtype), None, None + + d_rms = torch.empty_like(sum_output) + if zero_centered: + _m().rmsnorm_zero_centered_backward( + d_rms, + dw_float, + partial_dw, + grad_norm_output, + sum_output, + weight, + epsilon, + ) + else: + _m().rmsnorm_backward( + d_rms, + dw_float, + partial_dw, + grad_norm_output, + sum_output, + weight, + epsilon, + ) + dw = dw_float.to(weight.dtype) + + if grad_sum_output is None: + d_add = d_rms + elif d_rms is None: + d_add = grad_sum_output + else: + d_add = grad_sum_output + d_rms + + return d_add, d_add, dw, None, None + + +class Rmsnorm: + def __init__(self): + pass + + def __call__(self, input, weight, epsilon, inference=False, zero_centered=False): + if inference: + input = _force_kernel_layout(input) + output = torch.empty_like(input) + if zero_centered: + _m().rmsnorm_zero_centered(output, input, weight, epsilon) + else: + _m().rmsnorm(output, input, weight, epsilon) + return output + return _RmsnormFunction.apply(input, weight, epsilon, zero_centered) + + +class AddRmsnorm: + def __init__(self): + pass + + def __call__( + self, residual, input, weight, epsilon, inference=False, zero_centered=False + ): + if inference: + residual, input, sum_output, norm_output = _prepare_add_rmsnorm_inputs( + residual, input + ) + if zero_centered: + _m().add_rmsnorm_zero_centered( + sum_output, norm_output, residual, input, weight, epsilon + ) + else: + _m().add_rmsnorm( + sum_output, norm_output, residual, input, weight, epsilon + ) + return sum_output, norm_output + + return _AddRmsnormFunction.apply( + residual, input, weight, epsilon, zero_centered + ) + + + + +class _SiluAndMulSplitFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input1: torch.Tensor) -> torch.Tensor: + intermediate_size = input1.size(-1) // 2 + output = torch.empty( + (*input1.shape[:-1], intermediate_size), + dtype=input1.dtype, + device=input1.device, + ) + _m().silu_and_mul_split(output, input1) + ctx.save_for_backward(input1) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + if not ctx.needs_input_grad[0]: + return None + + (gate_up,) = ctx.saved_tensors + if grad_output.stride(-1) != 1: + grad_output = grad_output.contiguous() + + grad_gate_up = torch.empty_like(gate_up) + _m().silu_and_mul_split_inplace_backward( + grad_gate_up, grad_output, gate_up + ) + return grad_gate_up + + +class _SiluAndMulSplitInplaceFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input1: torch.Tensor) -> torch.Tensor: + _m().silu_and_mul_split_inplace(input1) + intermediate_size = input1.size(-1) // 2 + return input1[..., :intermediate_size] + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + raise RuntimeError( + "silu_and_mul split inplace mutates its input and is only used " + "for inference/no-grad execution" + ) + + +def silu_and_mul_split_inplace(input1: torch.Tensor) -> torch.Tensor: + if torch.is_grad_enabled() and input1.requires_grad: + return _SiluAndMulSplitFunction.apply(input1) + return _SiluAndMulSplitInplaceFunction.apply(input1) + + +class _SiluAndMulFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: + if not input1.is_contiguous(): + input1 = input1.contiguous() + if not input2.is_contiguous(): + input2 = input2.contiguous() + output = torch.empty_like(input1) + _m().silu_and_mul(output, input1, input2) + ctx.save_for_backward(input1, input2) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + needs_input1, needs_input2 = ctx.needs_input_grad + if not needs_input1 and not needs_input2: + return None, None + + input1, input2 = ctx.saved_tensors + if grad_output.stride(-1) != 1 or not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + grad_input1 = torch.empty_like(input1) + grad_input2 = torch.empty_like(input2) + _m().silu_and_mul_backward( + grad_input1, grad_input2, grad_output, input1, input2 + ) + if not needs_input1: + grad_input1 = None + if not needs_input2: + grad_input2 = None + return grad_input1, grad_input2 + + +def silu_and_mul_two_input(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: + return _SiluAndMulFunction.apply(input1, input2) + + +class SiluAndMul: + def __init__(self): + pass + + def __call__( + self, input1: torch.Tensor, input2: torch.Tensor = None, inplace: bool = False + ) -> torch.Tensor: + """Fused SiLU + elementwise multiply, with optional fused split. + + Two calling conventions: + + 1. Split mode (input2 is None, inplace must be True): + out = op(gate_up, inplace=True) + gate_up has shape [..., 2*N]. The kernel reads gate = gate_up[..., :N] + and up = gate_up[..., N:], computes silu(gate)*up, and writes the result + back into gate_up[..., :N] in-place. Returns gate_up[..., :N]. + + 2. Two-tensor mode (input2 is provided): + out = op(gate, up) # non-inplace, allocates output + out = op(gate, up, inplace=True) # inplace, returns gate + gate and up must have the same shape. + + Args: + input1: gate tensor (two-tensor mode) or combined gate_up tensor (split mode). + input2: up tensor; omit to use split mode. + inplace: if True, writes result back into input1 (saves one allocation). + + Returns: + Tensor: silu(gate) * up, shape [..., N]. + """ + if input2 is None: + assert inplace, "split mode requires inplace=True" + assert input1.size(-1) % 2 == 0, "last dim must be even for split mode" + return silu_and_mul_split_inplace(input1) + if inplace: + _m().silu_and_mul_inplace(input1, input2) + return input1 + return silu_and_mul_two_input(input1, input2) diff --git a/wall_x/model/core/ops/activation.py b/wall_x/model/core/ops/activation.py new file mode 100644 index 00000000..0e0aebee --- /dev/null +++ b/wall_x/model/core/ops/activation.py @@ -0,0 +1,58 @@ +"""Activation fusion operators.""" + +import logging + +import torch +import torch.nn.functional as F + +from wall_x.model.core.ops.base import OpsProxy + +logger = logging.getLogger(__name__) + + +class SiluAndMulOp(OpsProxy): + """Fused SiLU + elementwise multiply. + + Signatures: + - silu_and_mul(gate, up) -> silu(gate) * up + - silu_and_mul(gate_up, inplace=True) -> silu(gate) * up + + In split mode, ``inplace=True`` is best-effort for inference. During + autograd-enabled execution the PyTorch fallback returns a fresh tensor to + avoid writing into a view. + """ + + @property + def _external_accel_name(self): + return "silu_and_mul" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import SiluAndMul + + return SiluAndMul() + except ImportError: + return None + except Exception as e: + logger.warning("SiluAndMulOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback(self, input1, input2=None, inplace=False): + if input2 is None: + if input1.size(-1) % 2 != 0: + raise AssertionError("last dim must be even for split mode") + gate, up = input1.chunk(2, dim=-1) + output = F.silu(gate) * up + if inplace and not torch.is_grad_enabled(): + gate.copy_(output) + return gate + return output + + output = F.silu(input1) * input2 + if inplace: + input1.copy_(output) + return input1 + return output + + +silu_and_mul = SiluAndMulOp() diff --git a/wall_x/model/core/ops/adarmsnorm.py b/wall_x/model/core/ops/adarmsnorm.py new file mode 100644 index 00000000..9f7575cc --- /dev/null +++ b/wall_x/model/core/ops/adarmsnorm.py @@ -0,0 +1,47 @@ +"""Adaptive RMS normalization operator (fused variance+scale+shift+gate). + +Used by pi05's expert post_attention_layernorm where modulation comes from +``self.dense(cond)`` and gets chunk'd into (scale, shift, gate). +""" + +import logging + +import torch +from wall_x.model.core.ops.base import OpsProxy + +logger = logging.getLogger(__name__) + + +class AdaRMSNormOp(OpsProxy): + """Adaptive RMSNorm: adarmsnorm(x, modulation, eps) -> (out, gate). + + Args: + x: input [..., H] + modulation: [..., 3H] or [B, 3H] when x is [B, S, H]. Will be + chunked into (scale, shift, gate) along the last dim. + eps: epsilon for variance. + + Returns: + out: x normalized then ``(1+scale)*x + shift`` matching x shape. + gate: extracted from modulation, broadcast-shape to align with x + (so subsequent gated residual ops work directly). + """ + + @property + def _external_accel_name(self): + return "adarmsnorm" + + def _pytorch_fallback(self, x, modulation, eps=1e-6): + """Pure PyTorch reference matching gemma adaptive RMSNorm math.""" + dtype = x.dtype + # Variance + rsqrt in fp32 (mirrors GemmaRMSNorm._norm). + var = torch.mean(torch.square(x.float()), dim=-1, keepdim=True) + normed = x * torch.rsqrt(var + eps) + if x.dim() == 3 and modulation.dim() == 2: + modulation = modulation.unsqueeze(1) + scale, shift, gate = torch.chunk(modulation, 3, dim=-1) + out = normed * (1.0 + scale.to(torch.float32)) + shift.to(torch.float32) + return out.to(dtype), gate.to(dtype) + + +adarmsnorm = AdaRMSNormOp() diff --git a/wall_x/model/core/ops/attention_mask.py b/wall_x/model/core/ops/attention_mask.py new file mode 100644 index 00000000..ce0661e6 --- /dev/null +++ b/wall_x/model/core/ops/attention_mask.py @@ -0,0 +1,504 @@ +import logging + +import torch + +logger = logging.getLogger(__name__) + +# ops is the sole boundary that touches external_accel: import the kernels directly +# here (capability-based — present in the loaded build => used) and expose the +# operators (kernel + PyTorch fallback) so higher layers import only from ops. +try: + from external_accel import ( + qwen35_update_joint_attention_flash_mask as _external_accel_update_flash_mask, + qwen35_update_position_ids as _external_accel_update_position_ids, + ) +except Exception: + _external_accel_update_flash_mask = None + _external_accel_update_position_ids = None + + +def _handle_external_accel_attention_mask_error(exc): + # Never swallow silently: an intermittently failing kernel degrading to the + # slow path must be visible in logs. + logger.warning( + "external_accel attention-mask kernel failed (%s); using PyTorch fallback", exc + ) + + +def find_first_last_ones(tensor): + """ + Input: tensor of shape (bs, seq_len) containing 0s and 1s + Output: (first_indices, last_indices), each of shape (bs,) + first_indices[i] is the first index of 1 in batch i, or -1 if none exists + last_indices[i] is the last index of 1 in batch i, or -1 if none exists + """ + bs, seq_len = tensor.shape + masks = tensor == 1 + has_ones = masks.any(dim=1) + + first = torch.full((bs,), -1, dtype=torch.long, device=tensor.device) + last = first.clone() + + # Compute the first index of 1 + first[has_ones] = torch.argmax(masks[has_ones].float(), dim=1) + + # Compute the last index of 1 + flipped_masks = masks.flip(dims=[1]) + last_argmax = torch.argmax(flipped_masks[has_ones].float(), dim=1) + last[has_ones] = seq_len - 1 - last_argmax + + return first, last + + +def update_position_ids(position_ids, moe_token_types, positional_masks): + """Extracted from ActionModelMixMin._update_position_ids (was @staticmethod).""" + if positional_masks is None: + return position_ids + + if ( + _external_accel_update_position_ids is not None + and torch.is_tensor(position_ids) + and position_ids.is_cuda + and torch.is_tensor(moe_token_types) + and moe_token_types.is_cuda + ): + try: + return _external_accel_update_position_ids( + position_ids, + moe_token_types, + ar_predict_token_positions=positional_masks.get( + "ar_predict_token_positions" + ), + flow_block_id=positional_masks.get("flow_block_id"), + rtc_pinned=positional_masks.get("rtc_pinned"), + rtc_rope_gap=positional_masks.get("rtc_rope_gap", 0), + ) + except Exception as exc: + _handle_external_accel_attention_mask_error(exc) + + new_position_ids = position_ids.clone() + flow_mask = moe_token_types == 1 + + if "ar_predict_token_positions" in positional_masks: + ar_predict_token_positions = positional_masks["ar_predict_token_positions"] + + start_ar_pos, end_ar_pos = find_first_last_ones(ar_predict_token_positions) + start_flow_pos, _ = find_first_last_ones(flow_mask) + + for bs_i in range(position_ids.shape[1]): + if start_ar_pos[bs_i] != -1 and end_ar_pos[bs_i] != -1: + start_ar_ids = new_position_ids[:, bs_i, start_ar_pos[bs_i]] + start_flow_ids = new_position_ids[:, bs_i, start_flow_pos[bs_i]] + diff = start_flow_ids - start_ar_ids + new_position_ids[:, bs_i, start_flow_pos[bs_i] :] = position_ids[ + :, bs_i, start_flow_pos[bs_i] : + ] - diff.unsqueeze(-1) + + # 修复: K 个 flow block 共享 baseline 位置(数据驱动多样本 A′)。dataset 把 input_ids + # tile 成 K 个 flow block 但不产出 position_ids; compute_3d 在【扩展后】序列上算出递增 + # 位置 → K 个并行去噪样本拿到各不相同的 RoPE 位置, 破坏"同一段轨迹的并行样本共享 + # baseline"。用 dataset 提供的 flow_block_id 把每个 block reset 成 block 0 的位置; 在 + # rope_gap 之前(gap 是各 block 基于自身 rtc_pinned 的后续修饰)。 + # 注: 仅在会运行 _update_position_ids 的注意力实现下生效(flash_mask / sdpa / + # flash_attention_2_ki / flash_attention_2_triton); 普通 flash_attention_2 本就不支持 + # 多样本(flow block 无隔离), 见 docs/multisample_rtc_attention_support.md。 + flow_block_id = positional_masks.get("flow_block_id") + if flow_block_id is not None: + for bs_i in range(new_position_ids.shape[1]): + fb = flow_block_id[bs_i] + if not bool((fb >= 0).any()): + continue + n_blocks = int(fb[fb >= 0].max().item()) + 1 + block0 = (fb == 0).nonzero(as_tuple=True)[0] + for blk in range(1, n_blocks): + blk_pos = (fb == blk).nonzero(as_tuple=True)[0] + if blk_pos.numel() != block0.numel(): + raise ValueError( + "flow_block_id requires equal-sized flow blocks for " + "position id reset: " + f"batch={bs_i}, block=0 has {block0.numel()} tokens, " + f"block={blk} has {blk_pos.numel()} tokens" + ) + new_position_ids[:, bs_i, blk_pos] = new_position_ids[:, bs_i, block0] + + rtc_rope_gap = positional_masks.get("rtc_rope_gap", 0) + if isinstance(rtc_rope_gap, torch.Tensor): + rtc_rope_gap = rtc_rope_gap.item() + rtc_rope_gap = int(rtc_rope_gap or 0) + if rtc_rope_gap > 0 and "rtc_pinned" in positional_masks: + rtc_pinned = positional_masks["rtc_pinned"].bool() + rtc_suffix_flow = flow_mask & ~rtc_pinned + if rtc_suffix_flow.any(): + new_position_ids = new_position_ids + ( + rtc_suffix_flow.unsqueeze(0).to(new_position_ids.dtype) * rtc_rope_gap + ) + + return new_position_ids + + +def _apply_multisample_flow_blocks_flash( + LTS, UTE, moe_token_types, flow_block_id, rtc_pinned=None +): + """Restrict FlashMask visibility by dataset-provided flow block ids.""" + i32 = LTS.dtype + B, S = LTS.shape + flow_seq = moe_token_types == 1 + if not bool(flow_seq.any()): + return + + valid_block = flow_seq & (flow_block_id >= 0) + if not bool(valid_block.any()): + return + + n_blocks = int(flow_block_id[valid_block].max().item()) + 1 + safe_block_id = torch.where(valid_block, flow_block_id, 0).long() + pos = torch.arange(S, device=flow_block_id.device).expand(B, S) + + block_start = torch.full((B, n_blocks), S, dtype=torch.long, device=pos.device) + block_end = torch.full((B, n_blocks), -1, dtype=torch.long, device=pos.device) + block_start.scatter_reduce_( + 1, + safe_block_id, + torch.where(valid_block, pos, S), + reduce="amin", + include_self=True, + ) + block_end.scatter_reduce_( + 1, + safe_block_id, + torch.where(valid_block, pos, -1), + reduce="amax", + include_self=True, + ) + + UTE[valid_block] = block_start.gather(1, safe_block_id)[valid_block].to(i32) + LTS[valid_block] = (block_end.gather(1, safe_block_id)[valid_block] + 1).to(i32) + + if rtc_pinned is not None: + pinned = rtc_pinned.bool() & valid_block + suffix = valid_block & ~rtc_pinned.bool() + suffix_start = torch.full((B, n_blocks), S, dtype=torch.long, device=pos.device) + suffix_start.scatter_reduce_( + 1, + safe_block_id, + torch.where(suffix, pos, S), + reduce="amin", + include_self=True, + ) + + ones = torch.ones_like(safe_block_id, dtype=torch.long) + pinned_counts = torch.zeros((B, n_blocks), dtype=torch.long, device=pos.device) + suffix_counts = torch.zeros_like(pinned_counts) + pinned_counts.scatter_add_(1, safe_block_id, torch.where(pinned, ones, 0)) + suffix_counts.scatter_add_(1, safe_block_id, torch.where(suffix, ones, 0)) + mixed_block = (pinned_counts > 0) & (suffix_counts > 0) + suffix_in_mixed_block = suffix & mixed_block.gather(1, safe_block_id) + UTE[suffix_in_mixed_block] = suffix_start.gather(1, safe_block_id)[ + suffix_in_mixed_block + ].to(i32) + + +def _apply_multisample_flow_blocks_2d(attention_mask, flow_block_id, rtc_pinned=None): + """Apply per-block flow attention and RTC asymmetric masking to dense masks.""" + flow_seq = flow_block_id >= 0 + if not bool(flow_seq.any()): + return attention_mask + + zero = torch.zeros((), dtype=attention_mask.dtype, device=attention_mask.device) + one = torch.ones((), dtype=attention_mask.dtype, device=attention_mask.device) + + flow_pairs = flow_seq[:, :, None] & flow_seq[:, None, :] + attention_mask = torch.where(flow_pairs, zero, attention_mask) + + same_block_pairs = flow_pairs & ( + flow_block_id[:, :, None] == flow_block_id[:, None, :] + ) + attention_mask = torch.where(same_block_pairs, one, attention_mask) + + if rtc_pinned is not None: + pinned_query = rtc_pinned.bool()[:, :, None] + suffix_key = ~rtc_pinned.bool()[:, None, :] + pinned_to_suffix = same_block_pairs & pinned_query & suffix_key + attention_mask = torch.where(pinned_to_suffix, zero, attention_mask) + + return attention_mask + + +def update_joint_attention_mask_2d( + attention_mask, + moe_token_types, + positional_masks, + causal_action_attention_mask=False, +): + """Extracted from ActionModelMixMin._update_joint_attention_mask_2d. + The only self attribute used was self.config.causal_action_attention_mask, now passed as parameter. + """ + if attention_mask.dim() == 3: # bs, seq_len, seq_len + return attention_mask + + bs, seq_len = moe_token_types.shape[0], moe_token_types.shape[1] + # Create a lower-triangular causal mask + causal_mask = torch.tril( + torch.ones( + (seq_len, seq_len), dtype=torch.bfloat16, device=moe_token_types.device + ) + ) + # Expand to the batch dimension + attention_mask = causal_mask.unsqueeze(0).expand(bs, -1, -1) + + if positional_masks is not None and "padding_positions" in positional_masks: + padding_positions = positional_masks["padding_positions"] + # Set padding rows to zero + attention_mask = torch.where( + padding_positions[:, None, :], + torch.zeros_like(attention_mask), + attention_mask, + ) + # Set padding columns to zero + attention_mask = torch.where( + padding_positions[:, :, None], + torch.zeros_like(attention_mask), + attention_mask, + ) + + # Set the moe1 region to 1 and mask it from the fast region + moe1_mask = (moe_token_types[:, :, None]) & (moe_token_types[:, None, :]) + + if ( + not causal_action_attention_mask + ): # If causal action attention mask is disabled, set the whole moe1 region to 1 + attention_mask = torch.where( + moe1_mask, torch.ones_like(attention_mask), attention_mask + ) + + if ( + positional_masks is not None + and "ar_predict_token_positions" in positional_masks + ): + ar_predict_token_positions = positional_masks["ar_predict_token_positions"] + moe1_mask = (moe_token_types[:, :, None]) & ( + ar_predict_token_positions[:, None, :] + ) + attention_mask = torch.where( + moe1_mask, torch.zeros_like(attention_mask), attention_mask + ) + + # AR and flow are bidirectional + if positional_masks is not None and "ar_action_mask" in positional_masks: + ar_action_mask = positional_masks["ar_action_mask"] != 0 + flow_positions = moe_token_types == 1 + if positional_masks.get("ar_visible", True): + flow_ar_position = ar_action_mask | flow_positions + flow_ar_mask = flow_ar_position[:, :, None] & flow_ar_position[:, None, :] + attention_mask = torch.where( + flow_ar_mask, torch.ones_like(attention_mask), attention_mask + ) + else: + flow_flow_mask = flow_positions[:, :, None] & flow_positions[:, None, :] + ar_ar_mask = ar_action_mask[:, :, None] & ar_action_mask[:, None, :] + flow_ar_mask = flow_flow_mask | ar_ar_mask + + affected = ar_action_mask | flow_positions # (B, N) + affected_pair = affected[:, :, None] & affected[:, None, :] + attention_mask = attention_mask.masked_fill(affected_pair, 0) + + attention_mask = torch.where( + flow_ar_mask, torch.ones_like(attention_mask), attention_mask + ) + + if positional_masks is not None and "flow_block_id" in positional_masks: + rtc_pinned = positional_masks.get("rtc_pinned") + attention_mask = _apply_multisample_flow_blocks_2d( + attention_mask, positional_masks["flow_block_id"], rtc_pinned + ) + + if ( + positional_masks is not None + and "valid_flow_action_positions" in positional_masks + ): + # true in moe_token_types but false in valid_flow_action_positions + nonvalid_flow_action_positions = ( + moe_token_types & ~positional_masks["valid_flow_action_positions"] + ) + attention_mask = torch.where( + nonvalid_flow_action_positions[:, None, :], + torch.zeros_like(attention_mask), + attention_mask, + ) + attention_mask = torch.where( + nonvalid_flow_action_positions[:, :, None], + torch.zeros_like(attention_mask), + attention_mask, + ) + + return attention_mask + + +def update_joint_attention_flash_mask( + attention_mask, + moe_token_types, + positional_masks, + causal_action_attention_mask=False, + debug=False, + batch_size=None, + seq_len=None, + device=None, +): + """Extracted from ActionModelMixMin._update_joint_attention_flash_mask. + The only self attribute used was self.config.causal_action_attention_mask, now passed as parameter. + """ + if moe_token_types is None: + if batch_size is None or seq_len is None: + if attention_mask is None: + raise ValueError( + "`batch_size` and `seq_len` are required when " + "`moe_token_types` and `attention_mask` are both None." + ) + batch_size, seq_len = attention_mask.shape[:2] + if device is None: + if attention_mask is not None: + device = attention_mask.device + elif positional_masks: + for value in positional_masks.values(): + if hasattr(value, "device"): + device = value.device + break + if device is None: + raise ValueError( + "`device` is required when `positional_masks` has no tensors." + ) + else: + raise ValueError("`device` is required when `moe_token_types` is None.") + else: + device = moe_token_types.device + batch_size, seq_len = moe_token_types.shape + + B, S = batch_size, seq_len + i32 = torch.int32 + + if ( + _external_accel_update_flash_mask is not None + and moe_token_types is not None + and torch.is_tensor(moe_token_types) + and moe_token_types.is_cuda + ): + masks = positional_masks or {} + try: + return _external_accel_update_flash_mask( + moe_token_types, + padding_positions=masks.get("padding_positions"), + ar_predict_token_positions=masks.get("ar_predict_token_positions"), + valid_flow_action_positions=masks.get("valid_flow_action_positions"), + ar_action_mask=masks.get("ar_action_mask"), + flow_block_id=masks.get("flow_block_id"), + rtc_pinned=masks.get("rtc_pinned"), + causal_action_attention_mask=causal_action_attention_mask, + ar_visible=masks.get("ar_visible", True), + ) + except Exception as exc: + _handle_external_accel_attention_mask_error(exc) + + # ---- Return-vector initialization ---- + LTS = torch.ones((B, S), device=device, dtype=i32) * S + UTE = torch.arange(S, device=device, dtype=i32).unsqueeze(0).expand(B, S).clone() + + # Handle padding positions + if positional_masks is not None and "padding_positions" in positional_masks: + padding_positions = positional_masks["padding_positions"] + LTS[padding_positions] = 0 + UTE[padding_positions] = S + + # Handle AR predict tokens + if ( + positional_masks is not None + and "ar_predict_token_positions" in positional_masks + ): + start_ar_pos, end_ar_pos = find_first_last_ones( + positional_masks["ar_predict_token_positions"] + ) + for bs_i in range(B): + if end_ar_pos[bs_i] != -1: + LTS[bs_i, positional_masks["ar_predict_token_positions"][bs_i]] = ( + end_ar_pos[bs_i].to(i32) + 1 + ) + + flow_mask = moe_token_types == 1 if moe_token_types is not None else None + if positional_masks is None or "ar_action_mask" not in positional_masks: + # Handle bidirectional flow action masks + if flow_mask is not None and not causal_action_attention_mask: + start_flow_pos, _ = find_first_last_ones(flow_mask) + for bs_i in range(B): + if start_flow_pos[bs_i] != -1: + UTE[bs_i, flow_mask[bs_i]] = start_flow_pos[bs_i].to(i32) + else: + # AR and flow are bidirectional + ar_action_mask = positional_masks["ar_action_mask"] != 0 + if positional_masks.get("ar_visible", True): + flow_ar_position = ( + ar_action_mask if flow_mask is None else ar_action_mask | flow_mask + ) + for bs_i in range(B): + idx = flow_ar_position[bs_i].nonzero(as_tuple=True)[0] + if idx.numel() == 0: + continue + block_start = idx.min() + block_end = idx.max() + 1 + # Set the visible range of every token in the block to [block_start, block_end) + UTE[bs_i, idx] = block_start.to(i32) + LTS[bs_i, idx] = block_end.to(i32) + else: + for bs_i in range(B): + # 1) Flow sub-block: flow attends only to flow + if flow_mask is not None: + flow_idx = flow_mask[bs_i].nonzero(as_tuple=True)[0] + if flow_idx.numel() > 0: + flow_start = flow_idx.min() + flow_end = flow_idx.max() + 1 + # Only flow-token rows are set to [flow_start, flow_end) + UTE[bs_i, flow_idx] = flow_start.to(i32) + LTS[bs_i, flow_idx] = flow_end.to(i32) + + # 2) AR sub-block: AR attends only to AR + ar_idx = ar_action_mask[bs_i].nonzero(as_tuple=True)[0] + if ar_idx.numel() > 0: + ar_start = ar_idx.min() + ar_end = ar_idx.max() + 1 + # Only AR-token rows are set to [ar_start, ar_end) + UTE[bs_i, ar_idx] = ar_start.to(i32) + LTS[bs_i, ar_idx] = ar_end.to(i32) + + if ( + positional_masks is not None + and "flow_block_id" in positional_masks + and flow_mask is not None + ): + _apply_multisample_flow_blocks_flash( + LTS, + UTE, + moe_token_types, + positional_masks["flow_block_id"], + positional_masks.get("rtc_pinned"), + ) + + # Handle validation flow + if ( + positional_masks is not None + and "valid_flow_action_positions" in positional_masks + and flow_mask is not None + ): + nonvalid_flow_action_positions = ( + flow_mask & ~positional_masks["valid_flow_action_positions"] + ) + if nonvalid_flow_action_positions.any(): + LTS[nonvalid_flow_action_positions] = 0 + UTE[nonvalid_flow_action_positions] = S + + LTS = LTS.unsqueeze(-1) + UTE = UTE.unsqueeze(-1) + + startend_row_indices = torch.cat([LTS, UTE], dim=-1) + + # add num_heads dimension + startend_row_indices = startend_row_indices.unsqueeze(1) + + return startend_row_indices diff --git a/wall_x/model/core/ops/base.py b/wall_x/model/core/ops/base.py index 1a70e178..b0a78625 100644 --- a/wall_x/model/core/ops/base.py +++ b/wall_x/model/core/ops/base.py @@ -1,6 +1,5 @@ """Base class for operator proxies with lazy backend resolution.""" -import importlib import logging import threading @@ -10,15 +9,9 @@ class OpsProxy: """Callable proxy that lazily resolves to the best available backend. - Three-level fallback: - Level 1: external_accel (fastest, requires external_accel package) - Level 2: CUDA inline kernel (JIT compiled, requires CUDA) - Level 3: Pure PyTorch (always available) - - Subclasses define _external_accel_name, _pytorch_fallback, and optionally _cuda_kernel. - - Set _external_accel_name to None to skip external_accel resolution (e.g., when the - external_accel API is incompatible). + Public fallback order: + Level 1: CUDA inline kernel, when available + Level 2: Pure PyTorch fallback """ def __init__(self): @@ -27,49 +20,42 @@ def __init__(self): self._resolve_lock = threading.Lock() def _resolve(self): - """Resolve the best available backend. Called once on first use (thread-safe).""" + """Resolve the best available backend. Called once on first use.""" if self._resolved_fn is not None: return with self._resolve_lock: - # Double-check after acquiring lock if self._resolved_fn is not None: return - # Level 1: try external_accel (skip if _external_accel_name is None) - fn = self._probe_external_accel() - if fn is not None: - self._backend = "external_accel" - self._resolved_fn = fn - return - - # Level 2: try CUDA inline kernel (subclass override) cuda_fn = self._get_cuda_kernel() if cuda_fn is not None: self._backend = "cuda_inline" self._resolved_fn = cuda_fn return - # Level 3: PyTorch fallback + pytorch_fn = self._get_pytorch_fallback() + if pytorch_fn is None: + raise RuntimeError( + f"{self.__class__.__name__}: no available operator backend. " + f"Available: {self.available_backends()}" + ) self._backend = "pytorch" - self._resolved_fn = self._pytorch_fallback - logger.info(f"{self.__class__.__name__}: using PyTorch fallback") + self._resolved_fn = pytorch_fn + logger.info("%s: using PyTorch fallback", self.__class__.__name__) def _get_cuda_kernel(self): - """Override in subclass to provide Level 2 CUDA kernel. Returns None if unavailable.""" - return None - - @property - def _external_accel_name(self): - """Return external_accel attribute name, or None to skip external_accel resolution. - - Base returns None (skip external_accel). Subclasses override to return the - attribute name when external_accel integration is desired. - """ + """Override in subclasses to provide a CUDA inline kernel.""" return None def _pytorch_fallback(self, *args, **kwargs): raise NotImplementedError(f"{self.__class__.__name__} has no PyTorch fallback") + def _get_pytorch_fallback(self): + """Return the PyTorch fallback only when a subclass implements one.""" + if type(self)._pytorch_fallback is OpsProxy._pytorch_fallback: + return None + return self._pytorch_fallback + def __call__(self, *args, **kwargs): if self._resolved_fn is None: self._resolve() @@ -81,55 +67,25 @@ def backend(self) -> str: self._resolve() return self._backend - # ------------------------------------------------------------------ - # Multi-backend API (for testing / benchmarking) - # ------------------------------------------------------------------ - - def _probe_external_accel(self): - """Return the external_accel function for this op, or None if unavailable.""" - if self._external_accel_name is None: - return None - try: - mod = importlib.import_module("external_accel") - return getattr(mod, self._external_accel_name) - except (ImportError, AttributeError): - return None - def available_backends(self): - """Return list of available backend names for this operator.""" - backends = ["pytorch"] - if self._probe_external_accel() is not None: - backends.append("external_accel") + """Return available backend names for this operator.""" + backends = [] if self._get_cuda_kernel() is not None: backends.append("cuda_inline") + if self._get_pytorch_fallback() is not None: + backends.append("pytorch") return backends def _get_backend_fn(self, backend): """Return the callable for a specific backend, or None if unavailable.""" if backend == "pytorch": - return self._pytorch_fallback - elif backend == "external_accel": - return self._probe_external_accel() - elif backend == "cuda_inline": + return self._get_pytorch_fallback() + if backend == "cuda_inline": return self._get_cuda_kernel() - else: - raise ValueError(f"Unknown backend: {backend}") + raise ValueError(f"Unknown backend: {backend}") def call_with_backend(self, backend, *args, **kwargs): - """Call this operator using a specific backend. - - This does NOT change the default backend used by ``__call__``. - The ``backend`` property still reflects the auto-resolved default. - - Args: - backend: One of "external_accel", "cuda_inline", "pytorch". - - Returns: - Operator output from the specified backend. - - Raises: - RuntimeError: If the requested backend is not available. - """ + """Call this operator using a specific backend.""" fn = self._get_backend_fn(backend) if fn is None: raise RuntimeError( diff --git a/wall_x/model/core/ops/csrc/binding.cu b/wall_x/model/core/ops/csrc/binding.cu index 7b860616..d2450d32 100644 --- a/wall_x/model/core/ops/csrc/binding.cu +++ b/wall_x/model/core/ops/csrc/binding.cu @@ -7,6 +7,27 @@ #include #include +namespace tvm_ffi_rmsnorm { +void RMSNorm(const at::Tensor& output, const at::Tensor& input, const at::Tensor& weight, double epsilon); +void RMSNormZeroCentered(const at::Tensor& output, const at::Tensor& input, const at::Tensor& weight, double epsilon); +void AddRMSNorm(const at::Tensor& sum_output, const at::Tensor& norm_output, const at::Tensor& residual, const at::Tensor& input, const at::Tensor& weight, double epsilon); +void AddRMSNormZeroCentered(const at::Tensor& sum_output, const at::Tensor& norm_output, const at::Tensor& residual, const at::Tensor& input, const at::Tensor& weight, double epsilon); +void AddRMSNormWithRstd(const at::Tensor& sum_output, const at::Tensor& norm_output, const at::Tensor& rstd_output, const at::Tensor& residual, const at::Tensor& input, const at::Tensor& weight, double epsilon); +void AddRMSNormZeroCenteredWithRstd(const at::Tensor& sum_output, const at::Tensor& norm_output, const at::Tensor& rstd_output, const at::Tensor& residual, const at::Tensor& input, const at::Tensor& weight, double epsilon); +void RMSNormBackward( const at::Tensor& dx, const at::Tensor& dw, const at::Tensor& partial_dw_view, const at::Tensor& dy, const at::Tensor& input, const at::Tensor& weight, double epsilon); +void RMSNormZeroCenteredBackward( const at::Tensor& dx, const at::Tensor& dw, const at::Tensor& partial_dw_view, const at::Tensor& dy, const at::Tensor& input, const at::Tensor& weight, double epsilon); +void AddRMSNormBackward( const at::Tensor& d_add, const at::Tensor& dw, const at::Tensor& partial_dw_view, const at::Tensor& grad_sum, const at::Tensor& grad_norm, const at::Tensor& sum_input, const at::Tensor& weight, double epsilon); +void AddRMSNormZeroCenteredBackward( const at::Tensor& d_add, const at::Tensor& dw, const at::Tensor& partial_dw_view, const at::Tensor& grad_sum, const at::Tensor& grad_norm, const at::Tensor& sum_input, const at::Tensor& weight, double epsilon); +void AddRMSNormBackwardWithRstd( const at::Tensor& d_add, const at::Tensor& dw, const at::Tensor& partial_dw_view, const at::Tensor& grad_sum, const at::Tensor& grad_norm, const at::Tensor& sum_input, const at::Tensor& weight, const at::Tensor& rstd); +void AddRMSNormZeroCenteredBackwardWithRstd( const at::Tensor& d_add, const at::Tensor& dw, const at::Tensor& partial_dw_view, const at::Tensor& grad_sum, const at::Tensor& grad_norm, const at::Tensor& sum_input, const at::Tensor& weight, const at::Tensor& rstd); +void GatedRMSNorm(const at::Tensor& output, const at::Tensor& input, const at::Tensor& gate, const at::Tensor& weight, double epsilon); +void GatedRMSNormBackward( const at::Tensor& dx, const at::Tensor& dgate, const at::Tensor& dw, const at::Tensor& partial_dw_view, const at::Tensor& dy, const at::Tensor& input, const at::Tensor& gate, const at::Tensor& weight, double epsilon); +void GatedRMSNormWithRstd(const at::Tensor& output, const at::Tensor& rstd, const at::Tensor& input, const at::Tensor& gate, const at::Tensor& weight, double epsilon, int64_t warps_per_block); +void GatedRMSNormBackwardWithRstd( const at::Tensor& dx, const at::Tensor& dgate, const at::Tensor& dw, const at::Tensor& partial_dw_view, const at::Tensor& dy, const at::Tensor& input, const at::Tensor& gate, const at::Tensor& rstd, const at::Tensor& weight, double epsilon, int64_t warps_per_block); +void AdaRMSNorm(const at::Tensor& output, const at::Tensor& gate, const at::Tensor& input, const at::Tensor& modulation, double epsilon); +void AdaRMSNormBackward( const at::Tensor& dx, const at::Tensor& dmod, const at::Tensor& dy, const at::Tensor& dgate, const at::Tensor& input, const at::Tensor& modulation, double epsilon); +} // namespace tvm_ffi_rmsnorm + namespace wallx_cuda_get_rope_index { void GetRopeIndex(const c10::optional& input_ids, const c10::optional& image_grid_thw, const c10::optional& video_grid_thw, const c10::optional& second_per_grid_ts, const c10::optional& attention_mask, const at::Tensor& position_ids, const at::Tensor& mrope_deltas, const at::Tensor& workspace, int spatial_merge_size, int image_token_id, int video_token_id, int vision_start_token_id, float tokens_per_second); int64_t GetRopeIndexGetWorkSpace(const c10::optional& input_ids, const c10::optional& image_grid_thw, const c10::optional& video_grid_thw); @@ -17,6 +38,7 @@ void MRope(const at::Tensor& q, const at::Tensor& k, const at::Tensor& q_embed, void MRopeInplace(const at::Tensor& q, const at::Tensor& k, const at::Tensor& cos, const at::Tensor& sin, const int first, const int second ); void MRopeInplacePack(const at::Tensor& qkv, const at::Tensor& cos, const at::Tensor& sin, int64_t q_num_heads, int64_t kv_num_heads, const int first, const int second); void MRopeInplacePackBackward(const at::Tensor& dqkv, const at::Tensor& cos, const at::Tensor& sin, int64_t q_num_heads, int64_t kv_num_heads, const int first, const int second); +void MRopeInplaceCfgCacheWrite(const at::Tensor& q, const at::Tensor& k, const at::Tensor& v, const at::Tensor& cos, const at::Tensor& sin, const at::Tensor& kvcache, const at::Tensor& valid0, const at::Tensor& valid1, int64_t common_len, int64_t cond_tail_len, const int first, const int second); void MRopeBackward( const at::Tensor& grad_q_embed, const at::Tensor& grad_k_embed, const at::Tensor& grad_q, const at::Tensor& grad_k, const at::Tensor& cos, const at::Tensor& sin, const int first, const int second ); } // namespace wallx_cuda_m_rope @@ -29,10 +51,12 @@ size_t CubSortPairGetStorageBytes(int64_t num_items); namespace wallx_cuda_rope { void Rope(const at::Tensor& q, const at::Tensor& k, const at::Tensor& q_embed, const at::Tensor& k_embed, const at::Tensor& cos, const at::Tensor& sin, bool interleave ); +void RopeFull(const at::Tensor& q, const at::Tensor& k, const at::Tensor& q_embed, const at::Tensor& k_embed, const at::Tensor& cos, const at::Tensor& sin, bool interleave, int64_t rotary_dim ); void RopeInplace(const at::Tensor& q, const at::Tensor& k, const at::Tensor& cos, const at::Tensor& sin, bool interleave ); void RopeInplacePack(const at::Tensor& qkv, const at::Tensor& cos, const at::Tensor& sin, int64_t q_num_heads, int64_t kv_num_heads, bool interleave); void RopeInplacePackBackward(const at::Tensor& dqkv, const at::Tensor& cos, const at::Tensor& sin, int64_t q_num_heads, int64_t kv_num_heads, bool interleave); void RopeBackward( const at::Tensor& grad_q_embed, const at::Tensor& grad_k_embed, const at::Tensor& grad_q, const at::Tensor& grad_k, const at::Tensor& cos, const at::Tensor& sin, bool interleave ); +void RopeFullBackward( const at::Tensor& grad_q_embed, const at::Tensor& grad_k_embed, const at::Tensor& grad_q, const at::Tensor& grad_k, const at::Tensor& cos, const at::Tensor& sin, bool interleave, int64_t rotary_dim ); } // namespace wallx_cuda_rope namespace wallx_cuda_rot_pos { @@ -40,6 +64,20 @@ void GetTokenCounts( const at::Tensor& grid_thw, const at::Tensor& token_counts, void RotPosEmb( const at::Tensor& inv_freq, const at::Tensor& grid_thw, const at::Tensor& output, const at::Tensor& cumsum_tokens, int spatial_merge_size ); } // namespace wallx_cuda_rot_pos +namespace wallx_cuda_silu_and_mul { +void SiluAndMul( const at::Tensor& output, const at::Tensor& input1, const at::Tensor& input2); +void SiluAndMulBackward( const at::Tensor& grad_input1, const at::Tensor& grad_input2, const at::Tensor& grad_output, const at::Tensor& input1, const at::Tensor& input2); +void SiluAndMulInplace( const at::Tensor& input1, const at::Tensor& input2); +void SiluAndMulSplit( const at::Tensor& output, const at::Tensor& gate_up); +void SiluAndMulSplitInplace(const at::Tensor& gate_up); +void SiluAndMulSplitInplaceBackward( const at::Tensor& grad_gate_up, const at::Tensor& grad_output, const at::Tensor& gate_up); +} // namespace wallx_cuda_silu_and_mul + +namespace wallx_cuda_two_expert_scatter { +void TwoExpertPack( const at::Tensor& output, const at::Tensor& input, const at::Tensor& action_starts, const at::Tensor& action_ends, const at::Tensor& llm_offsets, const at::Tensor& action_offsets); +void TwoExpertUnpack( const at::Tensor& output, const at::Tensor& input, int64_t batch_size, int64_t seq_length, const at::Tensor& action_starts, const at::Tensor& action_ends, const at::Tensor& llm_offsets, const at::Tensor& action_offsets); +} // namespace wallx_cuda_two_expert_scatter + namespace wallx_cuda_window_index { void GetWindowIndex( const at::Tensor& grid_thw, const at::Tensor& grid_info_tensor, const at::Tensor& window_indices, const at::Tensor& cu_window_seqlens, const at::Tensor& window_counts_tensor, int max_grid_t, int spatial_merge_size, int vit_merger_window_size, int patch_size, int spatial_merge_unit); void GetTotals( const at::Tensor& grid_thw, const at::Tensor& grid_info_tensor, const at::Tensor& global_totals_tensor, int spatial_merge_size, int vit_merger_window_size ); @@ -47,14 +85,17 @@ void GetTotals( const at::Tensor& grid_thw, const at::Tensor& grid_info_tensor, PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("rope", &wallx_cuda_rope::Rope); + m.def("rope_full", &wallx_cuda_rope::RopeFull); m.def("rope_inplace", &wallx_cuda_rope::RopeInplace); m.def("rope_inplace_pack", &wallx_cuda_rope::RopeInplacePack); m.def("rope_inplace_pack_bwd", &wallx_cuda_rope::RopeInplacePackBackward); m.def("rope_bwd", &wallx_cuda_rope::RopeBackward); + m.def("rope_full_bwd", &wallx_cuda_rope::RopeFullBackward); m.def("m_rope", &wallx_cuda_m_rope::MRope); m.def("m_rope_inplace", &wallx_cuda_m_rope::MRopeInplace); m.def("m_rope_inplace_pack", &wallx_cuda_m_rope::MRopeInplacePack); m.def("m_rope_inplace_pack_bwd", &wallx_cuda_m_rope::MRopeInplacePackBackward); + m.def("m_rope_inplace_cfg_cache_write", &wallx_cuda_m_rope::MRopeInplaceCfgCacheWrite); m.def("m_rope_bwd", &wallx_cuda_m_rope::MRopeBackward); m.def("get_token_counts", &wallx_cuda_rot_pos::GetTokenCounts); m.def("rot_pos", &wallx_cuda_rot_pos::RotPosEmb); @@ -66,4 +107,30 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("unpermute", &wallx_cuda_permute::MoeRecoverTopKOp); m.def("unpermute_bwd", &wallx_cuda_permute::MoeRecoverTopKBwdOp); m.def("cub_sort_pair_get_storage_bytes", &wallx_cuda_permute::CubSortPairGetStorageBytes); + m.def("two_expert_pack", &wallx_cuda_two_expert_scatter::TwoExpertPack); + m.def("two_expert_unpack", &wallx_cuda_two_expert_scatter::TwoExpertUnpack); + m.def("rmsnorm", &tvm_ffi_rmsnorm::RMSNorm); + m.def("rmsnorm_zero_centered", &tvm_ffi_rmsnorm::RMSNormZeroCentered); + m.def("add_rmsnorm", &tvm_ffi_rmsnorm::AddRMSNorm); + m.def("add_rmsnorm_zero_centered", &tvm_ffi_rmsnorm::AddRMSNormZeroCentered); + m.def("add_rmsnorm_with_rstd", &tvm_ffi_rmsnorm::AddRMSNormWithRstd); + m.def("add_rmsnorm_zero_centered_with_rstd", &tvm_ffi_rmsnorm::AddRMSNormZeroCenteredWithRstd); + m.def("rmsnorm_backward", &tvm_ffi_rmsnorm::RMSNormBackward); + m.def("rmsnorm_zero_centered_backward", &tvm_ffi_rmsnorm::RMSNormZeroCenteredBackward); + m.def("add_rmsnorm_backward", &tvm_ffi_rmsnorm::AddRMSNormBackward); + m.def("add_rmsnorm_zero_centered_backward", &tvm_ffi_rmsnorm::AddRMSNormZeroCenteredBackward); + m.def("add_rmsnorm_backward_with_rstd", &tvm_ffi_rmsnorm::AddRMSNormBackwardWithRstd); + m.def("add_rmsnorm_zero_centered_backward_with_rstd", &tvm_ffi_rmsnorm::AddRMSNormZeroCenteredBackwardWithRstd); + m.def("gated_rmsnorm", &tvm_ffi_rmsnorm::GatedRMSNorm); + m.def("gated_rmsnorm_backward", &tvm_ffi_rmsnorm::GatedRMSNormBackward); + m.def("gated_rmsnorm_with_rstd", &tvm_ffi_rmsnorm::GatedRMSNormWithRstd); + m.def("gated_rmsnorm_backward_with_rstd", &tvm_ffi_rmsnorm::GatedRMSNormBackwardWithRstd); + m.def("adarmsnorm", &tvm_ffi_rmsnorm::AdaRMSNorm); + m.def("adarmsnorm_backward", &tvm_ffi_rmsnorm::AdaRMSNormBackward); + m.def("silu_and_mul", &wallx_cuda_silu_and_mul::SiluAndMul); + m.def("silu_and_mul_backward", &wallx_cuda_silu_and_mul::SiluAndMulBackward); + m.def("silu_and_mul_inplace", &wallx_cuda_silu_and_mul::SiluAndMulInplace); + m.def("silu_and_mul_split", &wallx_cuda_silu_and_mul::SiluAndMulSplit); + m.def("silu_and_mul_split_inplace", &wallx_cuda_silu_and_mul::SiluAndMulSplitInplace); + m.def("silu_and_mul_split_inplace_backward", &wallx_cuda_silu_and_mul::SiluAndMulSplitInplaceBackward); } diff --git a/wall_x/model/core/ops/csrc/common/activations.cuh b/wall_x/model/core/ops/csrc/common/activations.cuh index 4560eba4..9c0a6de3 100644 --- a/wall_x/model/core/ops/csrc/common/activations.cuh +++ b/wall_x/model/core/ops/csrc/common/activations.cuh @@ -27,4 +27,51 @@ __forceinline__ __device__ T applyActivation(const T &x) { } } +// Forward activation and its derivative computed together so the costly +// transcendental (tanhf for GELU, __expf for SILU) is evaluated once per +// element and shared. Returns both as fp32 — caller does the chain-rule +// multiply in high precision before casting back to T. +// +// Used by activ_and_mul backward kernel: +// dL/d(gate) = grad * up * dact(gate) +// dL/d(up) = grad * act(gate) +// — both consume act / dact evaluated at the same gate element. +// +// Numerics: GELU forward formula matches applyActivation byte-for-byte +// (same x³-first ordering for the tanh argument). GELU' uses the standard +// chain-rule expansion of the tanh approximation: +// dGELU/dx = 0.5·(1+t) + 0.5·x·(1-t²)·(√(2/π) + √(2/π)·3·0.044715·x²) +// where 0.1070322243 = √(2/π) · 3 · 0.044715 (pre-multiplied). +struct ActAndDeriv { + float act; + float dact; +}; + +template +__forceinline__ __device__ ActAndDeriv applyActivationAndDerivative(const T &x) { + const float x_f = (float)x; + ActAndDeriv r; + if constexpr (activation_type == ActivationType::RELU) { + r.act = x_f > 0.0f ? x_f : 0.0f; + r.dact = x_f > 0.0f ? 1.0f : 0.0f; + } + else if constexpr (activation_type == ActivationType::SILU) { + const float s = 1.0f / (1.0f + __expf(-x_f)); + r.act = x_f * s; + r.dact = s * (1.0f + x_f * (1.0f - s)); + } + else if constexpr (activation_type == ActivationType::GELU) { + const float x3 = x_f * x_f * x_f; + const float t = tanhf(0.79788456f * (x_f + 0.044715f * x3)); + const float x2 = x_f * x_f; + r.act = 0.5f * x_f * (1.0f + t); + r.dact = 0.5f * x_f * ((1.0f - t * t) * (0.79788456f + 0.1070322243f * x2)) + + 0.5f * (1.0f + t); + } + else { + assert(false); + } + return r; +} + } // namespace wallx_cuda diff --git a/wall_x/model/core/ops/csrc/get_rope_index/get_rope_index.cu b/wall_x/model/core/ops/csrc/get_rope_index/get_rope_index.cu index bc531901..060cfc8d 100644 --- a/wall_x/model/core/ops/csrc/get_rope_index/get_rope_index.cu +++ b/wall_x/model/core/ops/csrc/get_rope_index/get_rope_index.cu @@ -129,7 +129,7 @@ __global__ void compute_vision_counts( } -// Stage 2: Preprocessing kernel - parse and analyze all vision content +// Stage 2: Preprocessing kernel – parse and analyze all vision content __global__ void preprocess_vision_tokens( const int64_t *input_ids, // (batch_size, seq_len) const int64_t *attention_mask, // (batch_size, seq_len) @@ -276,7 +276,7 @@ __global__ void preprocess_vision_tokens( } -// Stage 3: Main kernel - compute 3D position IDs for all tokens in parallel +// Stage 3: Main kernel – compute 3D position IDs for all tokens in parallel __global__ void compute_3d_positions( const int64_t *input_ids, // (batch_size, seq_len) const int64_t *attention_mask, // (batch_size, seq_len) @@ -284,8 +284,8 @@ __global__ void compute_3d_positions( const int64_t *vision_counts, // (batch_size,) const int64_t *text_lengths, // (batch_size, MAX_VISION_TOKENS+1) const int64_t *position_offsets, // (batch_size, MAX_VISION_TOKENS+1) - int64_t *position_ids, // (3, batch_size, seq_len) - output 3D position IDs - int64_t *mrope_deltas, // (batch_size,) - RoPE length adjustment + int64_t *position_ids, // (3, batch_size, seq_len) — output 3D position IDs + int64_t *mrope_deltas, // (batch_size,) — RoPE length adjustment const int64_t batch_size, const int64_t seq_len, const float tokens_per_second) @@ -443,7 +443,7 @@ __global__ void compute_3d_positions( // Fallback kernel when no vision tokens exist: assign sequential positions based on attention mask __global__ void compute_3d_positions_mask_text( const int64_t* __restrict__ attention_mask, // (batch_size, seq_len) - int64_t* __restrict__ position_ids, // (3, batch_size, seq_len) - flattened row-major + int64_t* __restrict__ position_ids, // (3, batch_size, seq_len) — flattened row-major int64_t* __restrict__ mrope_deltas, // (batch_size,) const int64_t batch_size, const int64_t seq_len @@ -466,7 +466,7 @@ __global__ void compute_3d_positions_mask_text( int64_t total_valid = 0; // Compute prefix sum (cumulative count of valid tokens) - // Simple O(seq_len^2) per-thread loop - acceptable for moderate seq_len + // Simple O(seq_len^2) per-thread loop — acceptable for moderate seq_len for (int idx = tid; idx < seq_len; idx += blockDim.x) { int64_t sum = 0; for (int i = 0; i <= idx; ++i) { @@ -511,7 +511,7 @@ __global__ void arange_3d(int64_t seq_len, return; } - // Flatten index: (dim, batch, token) -> linear + // Flatten index: (dim, batch, token) → linear int64_t idx = i * gridDim.y * seq_len + j * seq_len + k; out[idx] = k; } @@ -605,7 +605,7 @@ void GetRopeIndex(const c10::optional& input_ids, cudaStream_t stream = static_cast( at::cuda::getCurrentCUDAStream().stream()); - // Case: no vision tokens -> use simple positional encoding + // Case: no vision tokens → use simple positional encoding if (!image_grid_thw.has_value() && !video_grid_thw.has_value()) { if (attention_mask.has_value()) { // Use attention mask to compute valid positions @@ -619,7 +619,7 @@ void GetRopeIndex(const c10::optional& input_ids, ); return; } else { - // No mask, no vision -> use simple arange + // No mask, no vision → use simple arange dim3 thread_num(256); dim3 grid(3, batch_size, (seq_len + 255) / 256); arange_3d<<>>(seq_len, static_cast(position_ids.data_ptr())); diff --git a/wall_x/model/core/ops/csrc/m_rope/m_rope_kernel.cu b/wall_x/model/core/ops/csrc/m_rope/m_rope_kernel.cu index 12422a98..8d055602 100644 --- a/wall_x/model/core/ops/csrc/m_rope/m_rope_kernel.cu +++ b/wall_x/model/core/ops/csrc/m_rope/m_rope_kernel.cu @@ -12,10 +12,30 @@ #include #include #include +#include #include "../common/cuda_utils.h" namespace wallx_cuda_m_rope { +constexpr int64_t kMaxMRopeGridDimYZ = 65535; + +inline dim3 MakeMRopeGrid(int64_t batch, int64_t seq_len) { + ASSERT_CHECK(batch > 0); + ASSERT_CHECK(seq_len > 0); + ASSERT_CHECK(batch <= 2147483647LL); + + const int64_t grid_z = (seq_len + kMaxMRopeGridDimYZ - 1) / kMaxMRopeGridDimYZ; + ASSERT_CHECK(grid_z <= kMaxMRopeGridDimYZ); + const int64_t grid_y = (seq_len + grid_z - 1) / grid_z; + return dim3(static_cast(batch), + static_cast(grid_y), + static_cast(grid_z)); +} + +__device__ __forceinline__ int64_t MRopeSeqBlockIdx() { + return static_cast(blockIdx.y) + static_cast(blockIdx.z) * gridDim.y; +} + // Helper: convert vector of T to float4 template __device__ __forceinline__ float4 to_float4(const T* ptr) { @@ -66,6 +86,21 @@ __device__ __forceinline__ void from_float4(T* ptr, const float4& f4) { } } +template +__device__ __forceinline__ void copy4(T* dst, const T* src) { + if constexpr (std::is_same_v) { + reinterpret_cast(dst)[0] = reinterpret_cast(src)[0]; + reinterpret_cast(dst)[1] = reinterpret_cast(src)[1]; + } else if constexpr (std::is_same_v) { + reinterpret_cast<__nv_bfloat162*>(dst)[0] = reinterpret_cast(src)[0]; + reinterpret_cast<__nv_bfloat162*>(dst)[1] = reinterpret_cast(src)[1]; + } else if constexpr (std::is_same_v) { + reinterpret_cast(dst)[0] = reinterpret_cast(src)[0]; + } else { + static_assert(sizeof(T) == 0, "Unsupported type"); + } +} + template __global__ void MRopeKernel(const float* cos, const float* sin, @@ -78,27 +113,31 @@ __global__ void MRopeKernel(const float* cos, const int first, const int second, const int d, - const int qb_stride, - const int qs_stride, - const int qh_stride, - const int kb_stride, - const int ks_stride, - const int kh_stride, - const int qeb_stride, - const int qes_stride, - const int qeh_stride, - const int keb_stride, - const int kes_stride, - const int keh_stride) { + const int64_t seq_len, + // int64 strides: (batch)*(b_stride) overflows int32 for + // tensors past 2^31 elements (e.g. b>=128 at s=4096,h=32,d=128). + const int64_t qb_stride, + const int64_t qs_stride, + const int64_t qh_stride, + const int64_t kb_stride, + const int64_t ks_stride, + const int64_t kh_stride, + const int64_t qeb_stride, + const int64_t qes_stride, + const int64_t qeh_stride, + const int64_t keb_stride, + const int64_t kes_stride, + const int64_t keh_stride) { extern __shared__ char cos_sin[]; const int half_dim = d / 2; float* cos_smem = reinterpret_cast(cos_sin); float* sin_smem = cos_smem + half_dim; int b = blockIdx.x; - int s = blockIdx.y; + int64_t s = MRopeSeqBlockIdx(); + if (s >= seq_len) return; - int64_t offset = gridDim.x * gridDim.y * half_dim; - int64_t cos_sin_b_stride = gridDim.y * half_dim; + int64_t offset = static_cast(gridDim.x) * seq_len * half_dim; + int64_t cos_sin_b_stride = seq_len * half_dim; int64_t cos_sin_s_stride = half_dim; #define SIN_GMEM(a, b, c, d) sin[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] #define COS_GMEM(a, b, c, d) cos[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] @@ -185,17 +224,19 @@ template __global__ void MRopeInplaceKernel(float* cos, float* sin, T* q, const int q_h, T* k, const int k_h, const int first, const int second, - const int d, const int qb_stride, const int qs_stride, const int qh_stride, - const int kb_stride, const int ks_stride, const int kh_stride) { + const int d, const int64_t seq_len, + const int64_t qb_stride, const int64_t qs_stride, const int64_t qh_stride, + const int64_t kb_stride, const int64_t ks_stride, const int64_t kh_stride) { extern __shared__ char cos_sin[]; const int half_dim = d / 2; float* cos_smem = reinterpret_cast(cos_sin); float* sin_smem = cos_smem + half_dim; int b = blockIdx.x; - int s = blockIdx.y; + int64_t s = MRopeSeqBlockIdx(); + if (s >= seq_len) return; - int64_t offset = gridDim.x * gridDim.y * half_dim; - int64_t cos_sin_b_stride = gridDim.y * half_dim; + int64_t offset = static_cast(gridDim.x) * seq_len * half_dim; + int64_t cos_sin_b_stride = seq_len * half_dim; int64_t cos_sin_s_stride = half_dim; #define SIN_GMEM(a, b, c, d) sin[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] #define COS_GMEM(a, b, c, d) cos[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] @@ -276,6 +317,153 @@ __global__ void MRopeInplaceKernel(float* cos, float* sin, } } +template +__global__ void MRopeInplaceCfgCacheWriteKernel( + const float* cos, + const float* sin, + T* q, + const int q_h, + T* k, + const int k_h, + const T* v, + T* kvcache, + const int64_t* valid0, + const int64_t* valid1, + const int common_len, + const int cond_tail_len, + const int first, + const int second, + const int d, + const int64_t qb_stride, + const int64_t qs_stride, + const int64_t qh_stride, + const int64_t kb_stride, + const int64_t ks_stride, + const int64_t kh_stride, + const int64_t vb_stride, + const int64_t vs_stride, + const int64_t vh_stride, + const int64_t cache_kv_stride, + const int64_t cache_batch_stride, + const int64_t cache_seq_stride, + const int64_t cache_head_stride) { + extern __shared__ char cos_sin[]; + const int half_dim = d / 2; + float* cos_smem = reinterpret_cast(cos_sin); + float* sin_smem = cos_smem + half_dim; + const int b = blockIdx.x; + const int s = blockIdx.y; + + const int64_t offset = static_cast(gridDim.x) * gridDim.y * half_dim; + const int64_t cos_sin_b_stride = static_cast(gridDim.y) * half_dim; + const int64_t cos_sin_s_stride = half_dim; + + for (int i = threadIdx.x; i < half_dim; i += blockDim.x) { + const int rope_axis = (i < first) ? 0 : ((i < first + second) ? 1 : 2); + const int64_t rope_offset = static_cast(rope_axis) * offset + + static_cast(b) * cos_sin_b_stride + + static_cast(s) * cos_sin_s_stride + i; + cos_smem[i] = cos[rope_offset]; + sin_smem[i] = sin[rope_offset]; + } + __syncthreads(); + + const int cond_end = common_len + cond_tail_len; + const bool is_common = s < common_len; + const bool is_cond_tail = (s >= common_len) && (s < cond_end); + int64_t cache_pos0 = -1; + int64_t cache_pos1 = -1; + if (is_common) { + cache_pos0 = valid0[s]; + cache_pos1 = valid1[s]; + } else if (is_cond_tail) { + cache_pos0 = valid0[s]; + } else { + const int valid1_idx = common_len + (s - cond_end); + cache_pos1 = valid1[valid1_idx]; + } + + for (int j = threadIdx.x * 4; j < q_h * half_dim; j += blockDim.x * 4) { + const int h_idx = j / half_dim; + const int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + T* q0_ptr = q + static_cast(b) * qb_stride + static_cast(s) * qs_stride + + static_cast(h_idx) * qh_stride + base_d; + T* q1_ptr = q0_ptr + half_dim; + const float4 q_x0 = to_float4(q0_ptr); + const float4 q_x1 = to_float4(q1_ptr); + const float4 cos_vec = to_float4(cos_smem + base_d); + const float4 sin_vec = to_float4(sin_smem + base_d); + + const float4 q0_rot = make_float4( + q_x0.x * cos_vec.x - q_x1.x * sin_vec.x, + q_x0.y * cos_vec.y - q_x1.y * sin_vec.y, + q_x0.z * cos_vec.z - q_x1.z * sin_vec.z, + q_x0.w * cos_vec.w - q_x1.w * sin_vec.w); + const float4 q1_rot = make_float4( + q_x1.x * cos_vec.x + q_x0.x * sin_vec.x, + q_x1.y * cos_vec.y + q_x0.y * sin_vec.y, + q_x1.z * cos_vec.z + q_x0.z * sin_vec.z, + q_x1.w * cos_vec.w + q_x0.w * sin_vec.w); + + from_float4(q0_ptr, q0_rot); + from_float4(q1_ptr, q1_rot); + } + + T* key_cache = kvcache; + T* value_cache = kvcache + cache_kv_stride; + for (int j = threadIdx.x * 4; j < k_h * half_dim; j += blockDim.x * 4) { + const int h_idx = j / half_dim; + const int base_d = j % half_dim; + if (base_d + 3 >= half_dim) continue; + + T* k0_ptr = k + static_cast(b) * kb_stride + static_cast(s) * ks_stride + + static_cast(h_idx) * kh_stride + base_d; + T* k1_ptr = k0_ptr + half_dim; + const T* v0_ptr = v + static_cast(b) * vb_stride + static_cast(s) * vs_stride + + static_cast(h_idx) * vh_stride + base_d; + const T* v1_ptr = v0_ptr + half_dim; + + const float4 k_x0 = to_float4(k0_ptr); + const float4 k_x1 = to_float4(k1_ptr); + const float4 cos_vec = to_float4(cos_smem + base_d); + const float4 sin_vec = to_float4(sin_smem + base_d); + + const float4 k0_rot = make_float4( + k_x0.x * cos_vec.x - k_x1.x * sin_vec.x, + k_x0.y * cos_vec.y - k_x1.y * sin_vec.y, + k_x0.z * cos_vec.z - k_x1.z * sin_vec.z, + k_x0.w * cos_vec.w - k_x1.w * sin_vec.w); + const float4 k1_rot = make_float4( + k_x1.x * cos_vec.x + k_x0.x * sin_vec.x, + k_x1.y * cos_vec.y + k_x0.y * sin_vec.y, + k_x1.z * cos_vec.z + k_x0.z * sin_vec.z, + k_x1.w * cos_vec.w + k_x0.w * sin_vec.w); + + from_float4(k0_ptr, k0_rot); + from_float4(k1_ptr, k1_rot); + + const int64_t cache_head_offset = static_cast(h_idx) * cache_head_stride + base_d; + if (is_common || is_cond_tail) { + T* kc0 = key_cache + 0 * cache_batch_stride + cache_pos0 * cache_seq_stride + cache_head_offset; + T* vc0 = value_cache + 0 * cache_batch_stride + cache_pos0 * cache_seq_stride + cache_head_offset; + from_float4(kc0, k0_rot); + from_float4(kc0 + half_dim, k1_rot); + copy4(vc0, v0_ptr); + copy4(vc0 + half_dim, v1_ptr); + } + if (is_common || !is_cond_tail) { + T* kc1 = key_cache + 1 * cache_batch_stride + cache_pos1 * cache_seq_stride + cache_head_offset; + T* vc1 = value_cache + 1 * cache_batch_stride + cache_pos1 * cache_seq_stride + cache_head_offset; + from_float4(kc1, k0_rot); + from_float4(kc1 + half_dim, k1_rot); + copy4(vc1, v0_ptr); + copy4(vc1 + half_dim, v1_ptr); + } + } +} + template __global__ void MRopeKernelBackward( const float* cos, @@ -288,8 +476,9 @@ __global__ void MRopeKernelBackward( const int q_h, const int k_h, const int d, - const int qb_stride, const int qs_stride, const int qh_stride, - const int kb_stride, const int ks_stride, const int kh_stride) { + const int64_t seq_len, + const int64_t qb_stride, const int64_t qs_stride, const int64_t qh_stride, + const int64_t kb_stride, const int64_t ks_stride, const int64_t kh_stride) { extern __shared__ char cos_sin[]; const int half_dim = d / 2; @@ -297,10 +486,11 @@ __global__ void MRopeKernelBackward( float* sin_smem = cos_smem + half_dim; int b = blockIdx.x; - int s = blockIdx.y; + int64_t s = MRopeSeqBlockIdx(); + if (s >= seq_len) return; - int64_t offset = gridDim.x * gridDim.y * half_dim; - int64_t cos_sin_b_stride = gridDim.y * half_dim; + int64_t offset = static_cast(gridDim.x) * seq_len * half_dim; + int64_t cos_sin_b_stride = seq_len * half_dim; int64_t cos_sin_s_stride = half_dim; #define SIN_GMEM(a, b, c, d) sin[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] @@ -444,7 +634,7 @@ void MRope(const at::Tensor& q, // [b, s, h, d] int64_t qeh_stride = q_embed.stride(2); int64_t keh_stride = k_embed.stride(2); - dim3 grid(batch, seq_len); + dim3 grid = MakeMRopeGrid(batch, seq_len); if (q.scalar_type() == at::kHalf) { const half* q_data = static_cast(q.data_ptr()); @@ -458,7 +648,7 @@ void MRope(const at::Tensor& q, // [b, s, h, d] q_data, q_head_num, k_data, k_head_num, q_embed_data, k_embed_data, - first, second, dim, + first, second, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride); } else if (q.scalar_type() == at::kBFloat16) { @@ -473,7 +663,7 @@ void MRope(const at::Tensor& q, // [b, s, h, d] q_data, q_head_num, k_data, k_head_num, q_embed_data, k_embed_data, - first, second, dim, + first, second, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride); } else if (q.scalar_type() == at::kFloat) { @@ -488,7 +678,7 @@ void MRope(const at::Tensor& q, // [b, s, h, d] q_data, q_head_num, k_data, k_head_num, q_embed_data, k_embed_data, - first, second, dim, + first, second, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride ); @@ -537,7 +727,7 @@ void MRopeInplace(const at::Tensor& q, // [b, s, h, d] int64_t qh_stride = q.stride(2); int64_t kh_stride = k.stride(2); - dim3 grid(batch, seq_len); + dim3 grid = MakeMRopeGrid(batch, seq_len); if (q.scalar_type() == at::kHalf) { half* q_data = static_cast(q.data_ptr()); @@ -548,7 +738,7 @@ void MRopeInplace(const at::Tensor& q, // [b, s, h, d] cos_data, sin_data, q_data, q_head_num, k_data, k_head_num, - first, second, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); + first, second, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); } else if (q.scalar_type() == at::kBFloat16) { __nv_bfloat16* q_data = static_cast<__nv_bfloat16*>(q.data_ptr()); __nv_bfloat16* k_data = static_cast<__nv_bfloat16*>(k.data_ptr()); @@ -558,7 +748,7 @@ void MRopeInplace(const at::Tensor& q, // [b, s, h, d] cos_data, sin_data, q_data, q_head_num, k_data, k_head_num, - first, second, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); + first, second, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); } else if (q.scalar_type() == at::kFloat) { float* q_data = static_cast(q.data_ptr()); float* k_data = static_cast(k.data_ptr()); @@ -568,7 +758,7 @@ void MRopeInplace(const at::Tensor& q, // [b, s, h, d] cos_data, sin_data, q_data, q_head_num, k_data, k_head_num, - first, second, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + first, second, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride ); } else { throw std::runtime_error("Unsupported data type for m_rope"); @@ -577,6 +767,130 @@ void MRopeInplace(const at::Tensor& q, // [b, s, h, d] sync_check_cuda_error(); } +void MRopeInplaceCfgCacheWrite(const at::Tensor& q, // [1, packed_seq, q_h, d] + const at::Tensor& k, // [1, packed_seq, kv_h, d] + const at::Tensor& v, // [1, packed_seq, kv_h, d] + const at::Tensor& cos, // [3, 1, packed_seq, d / 2] + const at::Tensor& sin, // [3, 1, packed_seq, d / 2] + const at::Tensor& kvcache, // [2, max_batch, max_len, kv_h, d] + const at::Tensor& valid0, // int64 original prefix indices for cond row + const at::Tensor& valid1, // int64 original prefix indices for uncond row + int64_t common_len, + int64_t cond_tail_len, + const int first, + const int second) { + int Nthreads = 256; + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + ASSERT_CHECK(q.dim() == 4); + ASSERT_CHECK(k.dim() == 4); + ASSERT_CHECK(v.dim() == 4); + ASSERT_CHECK(kvcache.dim() == 5); + ASSERT_CHECK(q.scalar_type() == k.scalar_type()); + ASSERT_CHECK(q.scalar_type() == v.scalar_type()); + ASSERT_CHECK(q.scalar_type() == kvcache.scalar_type()); + ASSERT_CHECK(q.size(0) == 1); + ASSERT_CHECK(k.size(0) == q.size(0)); + ASSERT_CHECK(v.size(0) == q.size(0)); + ASSERT_CHECK(k.size(1) == q.size(1)); + ASSERT_CHECK(v.size(1) == q.size(1)); + ASSERT_CHECK(v.size(2) == k.size(2)); + ASSERT_CHECK(k.size(3) == q.size(3)); + ASSERT_CHECK(v.size(3) == q.size(3)); + ASSERT_CHECK(q.size(3) % 8 == 0); + ASSERT_CHECK(q.stride(3) == 1); + ASSERT_CHECK(k.stride(3) == 1); + ASSERT_CHECK(v.stride(3) == 1); + ASSERT_CHECK(kvcache.stride(4) == 1); + + ASSERT_CHECK(cos.size(0) == 3); + ASSERT_CHECK(sin.size(0) == 3); + ASSERT_CHECK(cos.size(1) == q.size(0)); + ASSERT_CHECK(sin.size(1) == q.size(0)); + ASSERT_CHECK(cos.size(2) == q.size(1)); + ASSERT_CHECK(sin.size(2) == q.size(1)); + ASSERT_CHECK(cos.size(3) == q.size(3) / 2); + ASSERT_CHECK(sin.size(3) == q.size(3) / 2); + ASSERT_CHECK(cos.scalar_type() == at::kFloat); + ASSERT_CHECK(sin.scalar_type() == at::kFloat); + ASSERT_CHECK(cos.is_contiguous()); + ASSERT_CHECK(sin.is_contiguous()); + + ASSERT_CHECK(kvcache.size(0) == 2); + ASSERT_CHECK(kvcache.size(1) >= 2); + ASSERT_CHECK(kvcache.size(3) == k.size(2)); + ASSERT_CHECK(kvcache.size(4) == k.size(3)); + ASSERT_CHECK(valid0.dim() == 1); + ASSERT_CHECK(valid1.dim() == 1); + ASSERT_CHECK(valid0.scalar_type() == at::kLong); + ASSERT_CHECK(valid1.scalar_type() == at::kLong); + + int64_t seq_len = q.size(1); + int64_t uncond_tail_len = seq_len - common_len - cond_tail_len; + ASSERT_CHECK(common_len >= 0); + ASSERT_CHECK(cond_tail_len >= 0); + ASSERT_CHECK(uncond_tail_len >= 0); + ASSERT_CHECK(valid0.size(0) >= common_len + cond_tail_len); + ASSERT_CHECK(valid1.size(0) >= common_len + uncond_tail_len); + + int64_t q_head_num = q.size(2); + int64_t k_head_num = k.size(2); + int64_t dim = q.size(3); + dim3 grid(q.size(0), seq_len); + + if (q.scalar_type() == at::kHalf) { + MRopeInplaceCfgCacheWriteKernel<<>>( + static_cast(cos.data_ptr()), + static_cast(sin.data_ptr()), + static_cast(q.data_ptr()), q_head_num, + static_cast(k.data_ptr()), k_head_num, + static_cast(v.data_ptr()), + static_cast(kvcache.data_ptr()), + static_cast(valid0.data_ptr()), + static_cast(valid1.data_ptr()), + common_len, cond_tail_len, first, second, dim, + q.stride(0), q.stride(1), q.stride(2), + k.stride(0), k.stride(1), k.stride(2), + v.stride(0), v.stride(1), v.stride(2), + kvcache.stride(0), kvcache.stride(1), kvcache.stride(2), kvcache.stride(3)); + } else if (q.scalar_type() == at::kBFloat16) { + MRopeInplaceCfgCacheWriteKernel<__nv_bfloat16><<>>( + static_cast(cos.data_ptr()), + static_cast(sin.data_ptr()), + static_cast<__nv_bfloat16*>(q.data_ptr()), q_head_num, + static_cast<__nv_bfloat16*>(k.data_ptr()), k_head_num, + static_cast(v.data_ptr()), + static_cast<__nv_bfloat16*>(kvcache.data_ptr()), + static_cast(valid0.data_ptr()), + static_cast(valid1.data_ptr()), + common_len, cond_tail_len, first, second, dim, + q.stride(0), q.stride(1), q.stride(2), + k.stride(0), k.stride(1), k.stride(2), + v.stride(0), v.stride(1), v.stride(2), + kvcache.stride(0), kvcache.stride(1), kvcache.stride(2), kvcache.stride(3)); + } else if (q.scalar_type() == at::kFloat) { + MRopeInplaceCfgCacheWriteKernel<<>>( + static_cast(cos.data_ptr()), + static_cast(sin.data_ptr()), + static_cast(q.data_ptr()), q_head_num, + static_cast(k.data_ptr()), k_head_num, + static_cast(v.data_ptr()), + static_cast(kvcache.data_ptr()), + static_cast(valid0.data_ptr()), + static_cast(valid1.data_ptr()), + common_len, cond_tail_len, first, second, dim, + q.stride(0), q.stride(1), q.stride(2), + k.stride(0), k.stride(1), k.stride(2), + v.stride(0), v.stride(1), v.stride(2), + kvcache.stride(0), kvcache.stride(1), kvcache.stride(2), kvcache.stride(3)); + } else { + throw std::runtime_error("Unsupported data type for m_rope_inplace_cfg_cache_write"); + } + + sync_check_cuda_error(); +} + void MRopeBackward( const at::Tensor& grad_q_embed, const at::Tensor& grad_k_embed, @@ -620,7 +934,7 @@ void MRopeBackward( int64_t qh_stride = grad_q_embed.stride(2); int64_t kh_stride = grad_k_embed.stride(2); - dim3 grid(batch, seq_len); + dim3 grid = MakeMRopeGrid(batch, seq_len); if (grad_q_embed.scalar_type() == at::kHalf) { const half* gq_embed = static_cast(grad_q_embed.data_ptr()); @@ -634,7 +948,7 @@ void MRopeBackward( gq_embed, gk_embed, gq, gk, first, second, - q_head_num, k_head_num, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + q_head_num, k_head_num, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride ); } else if (grad_q_embed.scalar_type() == at::kBFloat16) { const __nv_bfloat16* gq_embed = static_cast(grad_q_embed.data_ptr()); @@ -648,7 +962,7 @@ void MRopeBackward( gq_embed, gk_embed, gq, gk, first, second, - q_head_num, k_head_num, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + q_head_num, k_head_num, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride ); } else if (grad_q_embed.scalar_type() == at::kFloat) { const float* gq_embed = static_cast(grad_q_embed.data_ptr()); @@ -662,7 +976,7 @@ void MRopeBackward( gq_embed, gk_embed, gq, gk, first, second, - q_head_num, k_head_num, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + q_head_num, k_head_num, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride ); } else { throw std::runtime_error("Unsupported data type for m_rope backward"); @@ -707,7 +1021,7 @@ void MRopeInplacePack(const at::Tensor& qkv, // [bz, seq_len, q_dim + 2*kv_dim] int64_t qkv_b_stride = qkv.stride(0); int64_t qkv_s_stride = qkv.stride(1); - dim3 grid(bz, seq_len); + dim3 grid = MakeMRopeGrid(bz, seq_len); if (qkv.scalar_type() == at::kHalf) { half* q_data = static_cast(qkv.data_ptr()); @@ -717,7 +1031,7 @@ void MRopeInplacePack(const at::Tensor& qkv, // [bz, seq_len, q_dim + 2*kv_dim] MRopeInplaceKernel<<>>( cos_data, sin_data, q_data, q_num_heads, k_data, kv_num_heads, - first, second, head_dim, + first, second, head_dim, seq_len, qkv_b_stride, qkv_s_stride, head_dim, qkv_b_stride, qkv_s_stride, head_dim); } else if (qkv.scalar_type() == at::kBFloat16) { @@ -728,7 +1042,7 @@ void MRopeInplacePack(const at::Tensor& qkv, // [bz, seq_len, q_dim + 2*kv_dim] MRopeInplaceKernel<__nv_bfloat16><<>>( cos_data, sin_data, q_data, q_num_heads, k_data, kv_num_heads, - first, second, head_dim, + first, second, head_dim, seq_len, qkv_b_stride, qkv_s_stride, head_dim, qkv_b_stride, qkv_s_stride, head_dim); } else if (qkv.scalar_type() == at::kFloat) { @@ -739,7 +1053,7 @@ void MRopeInplacePack(const at::Tensor& qkv, // [bz, seq_len, q_dim + 2*kv_dim] MRopeInplaceKernel<<>>( cos_data, sin_data, q_data, q_num_heads, k_data, kv_num_heads, - first, second, head_dim, + first, second, head_dim, seq_len, qkv_b_stride, qkv_s_stride, head_dim, qkv_b_stride, qkv_s_stride, head_dim); } else { @@ -783,7 +1097,7 @@ void MRopeInplacePackBackward(const at::Tensor& dqkv, // [bz, seq_len, q_dim + 2 int64_t dqkv_b_stride = dqkv.stride(0); int64_t dqkv_s_stride = dqkv.stride(1); - dim3 grid(bz, seq_len); + dim3 grid = MakeMRopeGrid(bz, seq_len); if (dqkv.scalar_type() == at::kHalf) { half* dq_data = static_cast(dqkv.data_ptr()); @@ -795,7 +1109,7 @@ void MRopeInplacePackBackward(const at::Tensor& dqkv, // [bz, seq_len, q_dim + 2 dq_data, dk_data, dq_data, dk_data, first, second, - q_num_heads, kv_num_heads, head_dim, + q_num_heads, kv_num_heads, head_dim, seq_len, dqkv_b_stride, dqkv_s_stride, head_dim, dqkv_b_stride, dqkv_s_stride, head_dim); } else if (dqkv.scalar_type() == at::kBFloat16) { @@ -808,7 +1122,7 @@ void MRopeInplacePackBackward(const at::Tensor& dqkv, // [bz, seq_len, q_dim + 2 dq_data, dk_data, dq_data, dk_data, first, second, - q_num_heads, kv_num_heads, head_dim, + q_num_heads, kv_num_heads, head_dim, seq_len, dqkv_b_stride, dqkv_s_stride, head_dim, dqkv_b_stride, dqkv_s_stride, head_dim); } else if (dqkv.scalar_type() == at::kFloat) { @@ -821,7 +1135,7 @@ void MRopeInplacePackBackward(const at::Tensor& dqkv, // [bz, seq_len, q_dim + 2 dq_data, dk_data, dq_data, dk_data, first, second, - q_num_heads, kv_num_heads, head_dim, + q_num_heads, kv_num_heads, head_dim, seq_len, dqkv_b_stride, dqkv_s_stride, head_dim, dqkv_b_stride, dqkv_s_stride, head_dim); } else { @@ -836,5 +1150,6 @@ void MRopeInplacePackBackward(const at::Tensor& dqkv, // [bz, seq_len, q_dim + 2 // Removed TVM FFI export (see pybind11 registration) // Removed TVM FFI export (see pybind11 registration) // Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) } // wallx_cuda_m_rope diff --git a/wall_x/model/core/ops/csrc/permute_unpermute/permute.cu b/wall_x/model/core/ops/csrc/permute_unpermute/permute.cu index 498ebe97..d245b1d5 100644 --- a/wall_x/model/core/ops/csrc/permute_unpermute/permute.cu +++ b/wall_x/model/core/ops/csrc/permute_unpermute/permute.cu @@ -15,7 +15,7 @@ #include "device_launch_parameters.h" // Global FP8 guard: use CUDA_VERSION (defined in both host and device passes), -// not __CUDA_ARCH__ (device-pass only - leaves host dispatchers unable to see +// not __CUDA_ARCH__ (device-pass only — leaves host dispatchers unable to see // FP8 typedefs/branches). FP8 storage types arrived in CUDA 11.8. #if defined(CUDA_VERSION) && (CUDA_VERSION >= 11800) #include diff --git a/wall_x/model/core/ops/csrc/rmsnorm/rmsnorm_kernel.cu b/wall_x/model/core/ops/csrc/rmsnorm/rmsnorm_kernel.cu new file mode 100644 index 00000000..504a51d1 --- /dev/null +++ b/wall_x/model/core/ops/csrc/rmsnorm/rmsnorm_kernel.cu @@ -0,0 +1,3947 @@ +// Re-enable CUDA half operators (PyTorch disables them) +#undef __CUDA_NO_HALF_OPERATORS__ +#undef __CUDA_NO_HALF_CONVERSIONS__ +#undef __CUDA_NO_HALF2_OPERATORS__ +#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include +#include + +#include +#include +#include + +#include + +#include "../common/cuda_utils.h" + +namespace tvm_ffi_rmsnorm { + +struct RMSNormShape { + int64_t num_tokens; + int64_t hidden_size; + int64_t row_stride; +}; + +struct TensorRowLayout { + int ndim; + int64_t sizes[8]; + int64_t strides[8]; + __host__ __device__ int64_t stride(int i) const { return strides[i]; } +}; + +RMSNormShape GetRMSNormShape(const at::Tensor& tensor) { + ASSERT_CHECK(tensor.dim() >= 1); + const int ndim = tensor.dim(); + const int64_t hidden_size = tensor.size(ndim - 1); + ASSERT_CHECK(hidden_size > 0); + ASSERT_CHECK(tensor.stride(ndim - 1) == 1); + + int64_t num_tokens = 1; + for (int i = 0; i < ndim - 1; ++i) { + num_tokens *= tensor.size(i); + } + ASSERT_CHECK(num_tokens > 0); + + int64_t row_stride = hidden_size; + if (ndim > 1) { + row_stride = tensor.stride(ndim - 2); + ASSERT_CHECK(row_stride >= hidden_size); + int64_t expected_stride = row_stride; + for (int i = ndim - 2; i >= 0; --i) { + ASSERT_CHECK(tensor.stride(i) == expected_stride); + expected_stride *= tensor.size(i); + } + } + return {num_tokens, hidden_size, row_stride}; +} + +void CheckSameShape(const at::Tensor& lhs, + const at::Tensor& rhs) { + ASSERT_CHECK(lhs.dim() == rhs.dim()); + for (int i = 0; i < lhs.dim(); ++i) { + ASSERT_CHECK(lhs.size(i) == rhs.size(i)); + } +} + +TensorRowLayout GetTensorRowLayout(const at::Tensor& tensor) { + ASSERT_CHECK(tensor.dim() >= 1); + ASSERT_CHECK(tensor.dim() <= 8); + const int ndim = tensor.dim(); + ASSERT_CHECK(tensor.size(ndim - 1) > 0); + ASSERT_CHECK(tensor.stride(ndim - 1) == 1); + + TensorRowLayout layout; + layout.ndim = ndim; + for (int i = 0; i < 8; ++i) { + layout.sizes[i] = 1; + layout.strides[i] = 0; + } + for (int i = 0; i < ndim; ++i) { + layout.sizes[i] = tensor.size(i); + layout.strides[i] = tensor.stride(i); + } + return layout; +} + +int64_t GetTensorNumTokens(const at::Tensor& tensor) { + ASSERT_CHECK(tensor.dim() >= 1); + int64_t num_tokens = 1; + for (int i = 0; i < tensor.dim() - 1; ++i) { + num_tokens *= tensor.size(i); + } + return num_tokens; +} + +void CheckTensorRowLayoutVectorized(const TensorRowLayout& layout, int64_t vec_size) { + for (int i = 0; i < layout.ndim - 1; ++i) { + ASSERT_CHECK(layout.strides[i] % vec_size == 0); + } +} + +bool IsFloat32(at::ScalarType dtype) { + return dtype == at::kFloat; +} + +bool IsFloat16(at::ScalarType dtype) { + return dtype == at::kHalf; +} + +bool IsBFloat16(at::ScalarType dtype) { + return dtype == at::kBFloat16; +} + +int64_t GetVecSize(at::ScalarType dtype) { + if (IsFloat32(dtype)) return VecTraits::vec_size; + if (IsFloat16(dtype)) return VecTraits<__half>::vec_size; + if (IsBFloat16(dtype)) return VecTraits<__nv_bfloat16>::vec_size; + throw std::runtime_error("Unsupported data type for RMSNorm"); +} + +template +struct dependent_false : std::false_type {}; + +template +__device__ __forceinline__ float ScalarToFloat(T value) { + if constexpr (std::is_same_v) { + return value; + } else if constexpr (std::is_same_v) { + return __half2float(value); + } else if constexpr (std::is_same_v) { + return __bfloat162float(value); + } else { + static_assert(dependent_false::value, "Unsupported RMSNorm scalar type"); + } +} + +template +__device__ __forceinline__ float RoundFloatToScalarType(float value) { + if constexpr (std::is_same_v) { + return value; + } else if constexpr (std::is_same_v) { + return __half2float(__float2half(value)); + } else if constexpr (std::is_same_v) { + return __bfloat162float(__float2bfloat16(value)); + } else { + static_assert(dependent_false::value, "Unsupported RMSNorm scalar type"); + } +} + +template +__device__ __forceinline__ void VecToFloats( + typename VecTraits::Type vec, float* values) { + if constexpr (std::is_same_v) { + values[0] = vec.x; + values[1] = vec.y; + values[2] = vec.z; + values[3] = vec.w; + } else if constexpr (std::is_same_v) { + values[0] = __half2float(vec.x); + values[1] = __half2float(vec.y); + } else if constexpr (std::is_same_v) { + values[0] = __bfloat162float(vec.x); + values[1] = __bfloat162float(vec.y); + } else { + static_assert(dependent_false::value, "Unsupported RMSNorm vector type"); + } +} + +template +__device__ __forceinline__ typename VecTraits::Type FloatsToVec(const float* values) { + typename VecTraits::Type vec; + if constexpr (std::is_same_v) { + vec.x = values[0]; + vec.y = values[1]; + vec.z = values[2]; + vec.w = values[3]; + } else if constexpr (std::is_same_v) { + vec.x = __float2half(values[0]); + vec.y = __float2half(values[1]); + } else if constexpr (std::is_same_v) { + vec.x = __float2bfloat16(values[0]); + vec.y = __float2bfloat16(values[1]); + } else { + static_assert(dependent_false::value, "Unsupported RMSNorm vector type"); + } + return vec; +} + +template +__device__ __forceinline__ void LoadWeightFloats( + const WeightDType* weight, int64_t offset, float* values) { + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + values[j] = ScalarToFloat(weight[offset + j]); + } +} + +template +__device__ __forceinline__ void Load4Floats(const T* ptr, float* values) { + if constexpr (std::is_same_v) { + VecToFloats(VecTraits::load(ptr), values); + } else { + VecToFloats(VecTraits::load(ptr), values); + VecToFloats(VecTraits::load(ptr + 2), values + 2); + } +} + +template +__device__ __forceinline__ void Store4Scalars(T* ptr, const float* values) { + if constexpr (std::is_same_v) { + VecTraits::store(ptr, FloatsToVec(values)); + } else { + VecTraits::store(ptr, FloatsToVec(values)); + VecTraits::store(ptr + 2, FloatsToVec(values + 2)); + } +} + +__device__ __forceinline__ void CpAsync8(void* smem_dst, const void* gmem_src) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + const unsigned smem_addr = + static_cast(__cvta_generic_to_shared(smem_dst)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 8, 8;\n" :: + "r"(smem_addr), "l"(gmem_src)); +#else + *reinterpret_cast(smem_dst) = + *reinterpret_cast(gmem_src); +#endif +} + +__device__ __forceinline__ void CpAsyncCommitGroup() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.commit_group;\n" ::); +#endif +} + +template +__device__ __forceinline__ void CpAsyncWaitGroup() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.wait_group %0;\n" :: "n"(PendingGroups)); +#endif +} + +template +__device__ __forceinline__ uint2 Load4Raw(const T* ptr) { + static_assert(!std::is_same_v, "Load4Raw is for 16-bit RMSNorm data"); + return *reinterpret_cast(ptr); +} + +template +__device__ __forceinline__ uint32_t PackVec2Bits(typename VecTraits::Type vec) { + static_assert(!std::is_same_v, "PackVec2Bits is for 16-bit RMSNorm data"); + union Bits { + typename VecTraits::Type vec; + uint32_t raw; + }; + Bits bits; + bits.vec = vec; + return bits.raw; +} + +template +__device__ __forceinline__ typename VecTraits::Type UnpackVec2Bits(uint32_t raw) { + static_assert(!std::is_same_v, "UnpackVec2Bits is for 16-bit RMSNorm data"); + union Bits { + uint32_t raw; + typename VecTraits::Type vec; + }; + Bits bits; + bits.raw = raw; + return bits.vec; +} + +template +__device__ __forceinline__ void Raw4ToFloats(uint2 raw, float* values) { + VecToFloats(UnpackVec2Bits(raw.x), values); + VecToFloats(UnpackVec2Bits(raw.y), values + 2); +} + +template +__device__ __forceinline__ void Store4RawFromFloats(T* ptr, const float* values) { + uint2 raw; + raw.x = PackVec2Bits(FloatsToVec(values)); + raw.y = PackVec2Bits(FloatsToVec(values + 2)); + *reinterpret_cast(ptr) = raw; +} + +template +__device__ __forceinline__ void LoadWeight4Floats( + const WeightDType* weight, int64_t offset, float* values) { + if constexpr (std::is_same_v) { + LoadWeightFloats(weight, offset, values); + } else { + Raw4ToFloats(*reinterpret_cast(weight + offset), + values); + } +} + +__device__ __forceinline__ float warpReduceSum(float val) { + for (int offset = 16; offset > 0; offset /= 2) + val += __shfl_down_sync(0xffffffff, val, offset); + return val; +} + +template +__device__ __forceinline__ float subwarpReduceSum(float val) { + for (int offset = WIDTH / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset, WIDTH); + } + return val; +} + +// Mask-aware variant: for kernels where whole sub-warp row-groups may have exited +// (early-returned) before the reduction, the shuffle mask must name only the active +// lanes of this row-group — 0xffffffff would name exited lanes, which is UB. +template +__device__ __forceinline__ float subwarpReduceSum(float val, unsigned mask) { + for (int offset = WIDTH / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(mask, val, offset, WIDTH); + } + return val; +} + +__device__ __forceinline__ float blockReduceSum(float val) { + static __shared__ float shared[32]; + int lane = threadIdx.x % 32; + int wid = threadIdx.x / 32; + val = warpReduceSum(val); + if (lane == 0) shared[wid] = val; + __syncthreads(); + int num_warps = (blockDim.x + 31) / 32; + val = (threadIdx.x < num_warps) ? shared[lane] : 0.0f; + if (wid == 0) val = warpReduceSum(val); + return val; +} + +// ======================== Forward ======================== + +template +__global__ void rmsnormKernel( + T* output, const T* input, const WeightDType* weight, + const float epsilon, const int64_t hidden_size, + const int64_t input_row_stride, const int64_t output_row_stride) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + float square_sum = 0.0f; + __shared__ float inv_rms; + const int64_t tid = threadIdx.x; + const int64_t bid = blockIdx.x; + const int64_t input_offset = bid * input_row_stride; + const int64_t output_offset = bid * output_row_stride; + const int64_t vec_hidden = hidden_size / VEC_SIZE; + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec x_vec = VecTraits::load(input + input_offset + i * VEC_SIZE); + float x_vals[VEC_SIZE]; + VecToFloats(x_vec, x_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + square_sum += x_vals[j] * x_vals[j]; + } + } + square_sum = blockReduceSum(square_sum); + if (tid == 0) inv_rms = rsqrtf(square_sum / static_cast(hidden_size) + epsilon); + __syncthreads(); + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec in_vec = VecTraits::load(input + input_offset + i * VEC_SIZE); + float x_vals[VEC_SIZE], w_vals[VEC_SIZE], out_vals[VEC_SIZE]; + VecToFloats(in_vec, x_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + if constexpr (ZeroCentered) w_vals[j] += 1.0f; + out_vals[j] = x_vals[j] * w_vals[j] * inv_rms; + } + VecTraits::store(output + output_offset + i * VEC_SIZE, FloatsToVec(out_vals)); + } +} + +template +void LaunchRMSNormKernel( + const at::Tensor& output, + const at::Tensor& input, + const at::Tensor& weight, + float epsilon, + int64_t hidden_size, + int64_t num_tokens, + int64_t input_row_stride, + int64_t output_row_stride, + int64_t block_size, + cudaStream_t stream) { + rmsnormKernel<<>>( + static_cast(output.data_ptr()), + static_cast(input.data_ptr()), + static_cast(weight.data_ptr()), + epsilon, hidden_size, input_row_stride, output_row_stride); +} + +template +__global__ void addRMSNormKernel( + T* sum_output, + T* norm_output, + const T* residual, + const T* input, + const WeightDType* weight, + float* rstd_output, + const float epsilon, + const int64_t hidden_size, + const int64_t residual_row_stride, + const int64_t input_row_stride, + const int64_t sum_row_stride, + const int64_t norm_row_stride) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + float square_sum = 0.0f; + __shared__ float inv_rms; + const int64_t tid = threadIdx.x; + const int64_t bid = blockIdx.x; + const int64_t residual_offset = bid * residual_row_stride; + const int64_t input_offset = bid * input_row_stride; + const int64_t sum_offset = bid * sum_row_stride; + const int64_t norm_offset = bid * norm_row_stride; + const int64_t vec_hidden = hidden_size / VEC_SIZE; + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec residual_vec = VecTraits::load(residual + residual_offset + i * VEC_SIZE); + Vec input_vec = VecTraits::load(input + input_offset + i * VEC_SIZE); + float residual_vals[VEC_SIZE], input_vals[VEC_SIZE], sum_vals[VEC_SIZE]; + VecToFloats(residual_vec, residual_vals); + VecToFloats(input_vec, input_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + // Match the old path: bf16/fp16 add materializes a rounded sum tensor, + // and RMSNorm consumes that rounded tensor. + sum_vals[j] = RoundFloatToScalarType(residual_vals[j] + input_vals[j]); + square_sum += sum_vals[j] * sum_vals[j]; + } + VecTraits::store(sum_output + sum_offset + i * VEC_SIZE, + FloatsToVec(sum_vals)); + } + square_sum = blockReduceSum(square_sum); + if (tid == 0) { + inv_rms = rsqrtf(square_sum / static_cast(hidden_size) + epsilon); + if (rstd_output != nullptr) { + rstd_output[bid] = inv_rms; + } + } + __syncthreads(); + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec sum_vec = VecTraits::load(sum_output + sum_offset + i * VEC_SIZE); + float sum_vals[VEC_SIZE]; + float w_vals[VEC_SIZE], out_vals[VEC_SIZE]; + VecToFloats(sum_vec, sum_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + if constexpr (ZeroCentered) w_vals[j] += 1.0f; + out_vals[j] = sum_vals[j] * w_vals[j] * inv_rms; + } + VecTraits::store(norm_output + norm_offset + i * VEC_SIZE, + FloatsToVec(out_vals)); + } +} + +template +void LaunchAddRMSNormKernel( + const at::Tensor& sum_output, + const at::Tensor& norm_output, + const at::Tensor& residual, + const at::Tensor& input, + const at::Tensor& weight, + float* rstd_output, + float epsilon, + int64_t hidden_size, + int64_t num_tokens, + int64_t residual_row_stride, + int64_t input_row_stride, + int64_t sum_row_stride, + int64_t norm_row_stride, + int64_t block_size, + cudaStream_t stream) { + addRMSNormKernel<<>>( + static_cast(sum_output.data_ptr()), + static_cast(norm_output.data_ptr()), + static_cast(residual.data_ptr()), + static_cast(input.data_ptr()), + static_cast(weight.data_ptr()), + rstd_output, + epsilon, hidden_size, residual_row_stride, input_row_stride, + sum_row_stride, norm_row_stride); +} + +template +void RMSNormImpl(const at::Tensor& output, const at::Tensor& input, + const at::Tensor& weight, double epsilon) { + ASSERT_CHECK(weight.dim() == 1); + ASSERT_CHECK(weight.is_contiguous()); + CheckSameShape(output, input); + const RMSNormShape shape = GetRMSNormShape(input); + const RMSNormShape output_shape = GetRMSNormShape(output); + const int64_t hidden_size = shape.hidden_size; + const int64_t num_tokens = shape.num_tokens; + ASSERT_CHECK(output_shape.hidden_size == hidden_size); + ASSERT_CHECK(output_shape.num_tokens == num_tokens); + ASSERT_CHECK(hidden_size == weight.size(0)); + ASSERT_CHECK(hidden_size % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(shape.row_stride % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(output_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + const int64_t block_size = std::min(hidden_size, 1024L); + if (IsFloat32(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchRMSNormKernel( + output, input, weight, static_cast(epsilon), hidden_size, num_tokens, + shape.row_stride, output_shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat16(weight.scalar_type())) + LaunchRMSNormKernel( + output, input, weight, static_cast(epsilon), hidden_size, num_tokens, + shape.row_stride, output_shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchRMSNormKernel( + output, input, weight, static_cast(epsilon), hidden_size, num_tokens, + shape.row_stride, output_shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchRMSNormKernel( + output, input, weight, static_cast(epsilon), hidden_size, num_tokens, + shape.row_stride, output_shape.row_stride, block_size, stream); + else if (IsBFloat16(input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchRMSNormKernel<__nv_bfloat16, __nv_bfloat16, ZeroCentered>( + output, input, weight, static_cast(epsilon), hidden_size, num_tokens, + shape.row_stride, output_shape.row_stride, block_size, stream); + else if (IsBFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchRMSNormKernel<__nv_bfloat16, float, ZeroCentered>( + output, input, weight, static_cast(epsilon), hidden_size, num_tokens, + shape.row_stride, output_shape.row_stride, block_size, stream); + else throw std::runtime_error("Unsupported data type for RMSNorm"); + sync_check_cuda_error(); +} + +template +void AddRMSNormImpl(const at::Tensor& sum_output, + const at::Tensor& norm_output, + const at::Tensor& residual, + const at::Tensor& input, + const at::Tensor& weight, + double epsilon) { + ASSERT_CHECK(weight.dim() == 1); + ASSERT_CHECK(weight.is_contiguous()); + CheckSameShape(sum_output, input); + CheckSameShape(norm_output, input); + CheckSameShape(residual, input); + const RMSNormShape input_shape = GetRMSNormShape(input); + const RMSNormShape residual_shape = GetRMSNormShape(residual); + const RMSNormShape sum_shape = GetRMSNormShape(sum_output); + const RMSNormShape norm_shape = GetRMSNormShape(norm_output); + const int64_t hidden_size = input_shape.hidden_size; + const int64_t num_tokens = input_shape.num_tokens; + ASSERT_CHECK(residual_shape.hidden_size == hidden_size); + ASSERT_CHECK(residual_shape.num_tokens == num_tokens); + ASSERT_CHECK(sum_shape.hidden_size == hidden_size); + ASSERT_CHECK(sum_shape.num_tokens == num_tokens); + ASSERT_CHECK(norm_shape.hidden_size == hidden_size); + ASSERT_CHECK(norm_shape.num_tokens == num_tokens); + ASSERT_CHECK(hidden_size == weight.size(0)); + ASSERT_CHECK(input.scalar_type() == residual.scalar_type()); + ASSERT_CHECK(input.scalar_type() == sum_output.scalar_type()); + ASSERT_CHECK(input.scalar_type() == norm_output.scalar_type()); + ASSERT_CHECK(hidden_size % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(input_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(residual_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(sum_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(norm_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + const int64_t block_size = std::min(hidden_size, 1024L); + if (IsFloat32(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchAddRMSNormKernel( + sum_output, norm_output, residual, input, weight, nullptr, static_cast(epsilon), + hidden_size, num_tokens, residual_shape.row_stride, input_shape.row_stride, + sum_shape.row_stride, norm_shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat16(weight.scalar_type())) + LaunchAddRMSNormKernel( + sum_output, norm_output, residual, input, weight, nullptr, static_cast(epsilon), + hidden_size, num_tokens, residual_shape.row_stride, input_shape.row_stride, + sum_shape.row_stride, norm_shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchAddRMSNormKernel( + sum_output, norm_output, residual, input, weight, nullptr, static_cast(epsilon), + hidden_size, num_tokens, residual_shape.row_stride, input_shape.row_stride, + sum_shape.row_stride, norm_shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchAddRMSNormKernel( + sum_output, norm_output, residual, input, weight, nullptr, static_cast(epsilon), + hidden_size, num_tokens, residual_shape.row_stride, input_shape.row_stride, + sum_shape.row_stride, norm_shape.row_stride, block_size, stream); + else if (IsBFloat16(input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchAddRMSNormKernel<__nv_bfloat16, __nv_bfloat16, ZeroCentered>( + sum_output, norm_output, residual, input, weight, nullptr, static_cast(epsilon), + hidden_size, num_tokens, residual_shape.row_stride, input_shape.row_stride, + sum_shape.row_stride, norm_shape.row_stride, block_size, stream); + else if (IsBFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchAddRMSNormKernel<__nv_bfloat16, float, ZeroCentered>( + sum_output, norm_output, residual, input, weight, nullptr, static_cast(epsilon), + hidden_size, num_tokens, residual_shape.row_stride, input_shape.row_stride, + sum_shape.row_stride, norm_shape.row_stride, block_size, stream); + else throw std::runtime_error("Unsupported data type for AddRMSNorm"); + sync_check_cuda_error(); +} + +void RMSNorm(const at::Tensor& output, const at::Tensor& input, + const at::Tensor& weight, double epsilon) { + RMSNormImpl(output, input, weight, epsilon); +} + +void RMSNormZeroCentered(const at::Tensor& output, const at::Tensor& input, + const at::Tensor& weight, double epsilon) { + RMSNormImpl(output, input, weight, epsilon); +} + +void AddRMSNorm(const at::Tensor& sum_output, + const at::Tensor& norm_output, + const at::Tensor& residual, + const at::Tensor& input, + const at::Tensor& weight, + double epsilon) { + AddRMSNormImpl(sum_output, norm_output, residual, input, weight, epsilon); +} + +void AddRMSNormZeroCentered(const at::Tensor& sum_output, + const at::Tensor& norm_output, + const at::Tensor& residual, + const at::Tensor& input, + const at::Tensor& weight, + double epsilon) { + AddRMSNormImpl(sum_output, norm_output, residual, input, weight, epsilon); +} + +template +void AddRMSNormWithRstdImpl(const at::Tensor& sum_output, + const at::Tensor& norm_output, + const at::Tensor& rstd_output, + const at::Tensor& residual, + const at::Tensor& input, + const at::Tensor& weight, + double epsilon) { + ASSERT_CHECK(rstd_output.scalar_type() == at::kFloat); + ASSERT_CHECK(rstd_output.is_contiguous()); + ASSERT_CHECK(weight.dim() == 1); + ASSERT_CHECK(weight.is_contiguous()); + CheckSameShape(sum_output, input); + CheckSameShape(norm_output, input); + CheckSameShape(residual, input); + const RMSNormShape input_shape = GetRMSNormShape(input); + const RMSNormShape residual_shape = GetRMSNormShape(residual); + const RMSNormShape sum_shape = GetRMSNormShape(sum_output); + const RMSNormShape norm_shape = GetRMSNormShape(norm_output); + const int64_t hidden_size = input_shape.hidden_size; + const int64_t num_tokens = input_shape.num_tokens; + ASSERT_CHECK(rstd_output.dim() == 1); + ASSERT_CHECK(rstd_output.size(0) == num_tokens); + ASSERT_CHECK(residual_shape.hidden_size == hidden_size); + ASSERT_CHECK(residual_shape.num_tokens == num_tokens); + ASSERT_CHECK(sum_shape.hidden_size == hidden_size); + ASSERT_CHECK(sum_shape.num_tokens == num_tokens); + ASSERT_CHECK(norm_shape.hidden_size == hidden_size); + ASSERT_CHECK(norm_shape.num_tokens == num_tokens); + ASSERT_CHECK(hidden_size == weight.size(0)); + ASSERT_CHECK(input.scalar_type() == residual.scalar_type()); + ASSERT_CHECK(input.scalar_type() == sum_output.scalar_type()); + ASSERT_CHECK(input.scalar_type() == norm_output.scalar_type()); + ASSERT_CHECK(hidden_size % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(input_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(residual_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(sum_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(norm_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + const int64_t block_size = std::min(hidden_size, 1024L); + float* rstd_ptr = static_cast(rstd_output.data_ptr()); + if (IsFloat32(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchAddRMSNormKernel( + sum_output, norm_output, residual, input, weight, rstd_ptr, + static_cast(epsilon), hidden_size, num_tokens, + residual_shape.row_stride, input_shape.row_stride, + sum_shape.row_stride, norm_shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat16(weight.scalar_type())) + LaunchAddRMSNormKernel( + sum_output, norm_output, residual, input, weight, rstd_ptr, + static_cast(epsilon), hidden_size, num_tokens, + residual_shape.row_stride, input_shape.row_stride, + sum_shape.row_stride, norm_shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchAddRMSNormKernel( + sum_output, norm_output, residual, input, weight, rstd_ptr, + static_cast(epsilon), hidden_size, num_tokens, + residual_shape.row_stride, input_shape.row_stride, + sum_shape.row_stride, norm_shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchAddRMSNormKernel( + sum_output, norm_output, residual, input, weight, rstd_ptr, + static_cast(epsilon), hidden_size, num_tokens, + residual_shape.row_stride, input_shape.row_stride, + sum_shape.row_stride, norm_shape.row_stride, block_size, stream); + else if (IsBFloat16(input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchAddRMSNormKernel<__nv_bfloat16, __nv_bfloat16, ZeroCentered>( + sum_output, norm_output, residual, input, weight, rstd_ptr, + static_cast(epsilon), hidden_size, num_tokens, + residual_shape.row_stride, input_shape.row_stride, + sum_shape.row_stride, norm_shape.row_stride, block_size, stream); + else if (IsBFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchAddRMSNormKernel<__nv_bfloat16, float, ZeroCentered>( + sum_output, norm_output, residual, input, weight, rstd_ptr, + static_cast(epsilon), hidden_size, num_tokens, + residual_shape.row_stride, input_shape.row_stride, + sum_shape.row_stride, norm_shape.row_stride, block_size, stream); + else throw std::runtime_error("Unsupported data type for AddRMSNorm with rstd"); + sync_check_cuda_error(); +} + +void AddRMSNormWithRstd(const at::Tensor& sum_output, + const at::Tensor& norm_output, + const at::Tensor& rstd_output, + const at::Tensor& residual, + const at::Tensor& input, + const at::Tensor& weight, + double epsilon) { + AddRMSNormWithRstdImpl( + sum_output, norm_output, rstd_output, residual, input, weight, epsilon); +} + +void AddRMSNormZeroCenteredWithRstd(const at::Tensor& sum_output, + const at::Tensor& norm_output, + const at::Tensor& rstd_output, + const at::Tensor& residual, + const at::Tensor& input, + const at::Tensor& weight, + double epsilon) { + AddRMSNormWithRstdImpl( + sum_output, norm_output, rstd_output, residual, input, weight, epsilon); +} + +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) + +// ======================== Backward: Persistent Kernel + Partial Reduction ======================== +// +// Launch num_blocks = min(num_tokens, SM_count) persistent blocks. +// Each block loops over assigned rows computing dx, accumulating dw in registers. +// After the loop, each block writes its partial dw to partial_dw[block_id, hidden_size]. +// A lightweight reduction kernel sums partial_dw into final dw. +// +// Advantages: +// - No atomicAdd contention (lock-free) +// - Weight W loaded once into L1 cache, reused across all rows +// - X and dY read exactly once per row, producing both dX and partial dW +// + +template +__global__ void rmsnormBackwardPersistentKernel( + T* __restrict__ dx, + float* __restrict__ partial_dw, + const T* __restrict__ dy, + const T* __restrict__ input, + const WeightDType* __restrict__ weight, + const float epsilon, + const int64_t hidden_size, + const int64_t num_tokens, + const int64_t dx_row_stride, + const int64_t dy_row_stride, + const int64_t input_row_stride +) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + + const int64_t tid = threadIdx.x; + const int64_t num_blocks = gridDim.x; + const int64_t vec_hidden = hidden_size / VEC_SIZE; + const float inv_hidden = 1.0f / static_cast(hidden_size); + + __shared__ float s_inv_rms; + __shared__ float s_c; + + // Register accumulator for partial dw + float local_dw[MAX_VEC_PER_THREAD * VEC_SIZE]; + #pragma unroll + for (int j = 0; j < MAX_VEC_PER_THREAD * VEC_SIZE; j++) local_dw[j] = 0.0f; + + // Persistent loop: each block processes multiple rows + for (int64_t row = blockIdx.x; row < num_tokens; row += num_blocks) { + const int64_t dx_off = row * dx_row_stride; + const int64_t dy_off = row * dy_row_stride; + const int64_t input_off = row * input_row_stride; + float sq = 0.0f, dp = 0.0f; + + // Pass 1: reductions + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec xv = VecTraits::load(input + input_off + i * VEC_SIZE); + Vec dv = VecTraits::load(dy + dy_off + i * VEC_SIZE); + float x_vals[VEC_SIZE], d_vals[VEC_SIZE], w_vals[VEC_SIZE]; + VecToFloats(xv, x_vals); + VecToFloats(dv, d_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + if constexpr (ZeroCentered) w_vals[j] += 1.0f; + sq += x_vals[j] * x_vals[j]; + dp += d_vals[j] * w_vals[j] * x_vals[j]; + } + } + sq = blockReduceSum(sq); __syncthreads(); + dp = blockReduceSum(dp); + if (tid == 0) { + float ve = sq * inv_hidden + epsilon; + s_inv_rms = rsqrtf(ve); + s_c = dp / (static_cast(hidden_size) * ve); + } + __syncthreads(); + const float irms = s_inv_rms, cv = s_c; + + // Pass 2: compute dx + accumulate dw in registers + int li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x, li++) { + Vec xv = VecTraits::load(input + input_off + i * VEC_SIZE); + Vec dv = VecTraits::load(dy + dy_off + i * VEC_SIZE); + float x_vals[VEC_SIZE], d_vals[VEC_SIZE], w_vals[VEC_SIZE], out_vals[VEC_SIZE]; + VecToFloats(xv, x_vals); + VecToFloats(dv, d_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + if constexpr (ZeroCentered) w_vals[j] += 1.0f; + out_vals[j] = irms * (d_vals[j] * w_vals[j] - x_vals[j] * cv); + local_dw[li * VEC_SIZE + j] += d_vals[j] * x_vals[j] * irms; + } + VecTraits::store(dx + dx_off + i * VEC_SIZE, FloatsToVec(out_vals)); + } + } + + // Write partial dw to global memory + const int64_t poff = blockIdx.x * hidden_size; + int li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x, li++) { + if constexpr (VEC_SIZE == 4) { + float4 dw_out; + dw_out.x = local_dw[li * VEC_SIZE + 0]; + dw_out.y = local_dw[li * VEC_SIZE + 1]; + dw_out.z = local_dw[li * VEC_SIZE + 2]; + dw_out.w = local_dw[li * VEC_SIZE + 3]; + *reinterpret_cast(partial_dw + poff + i * VEC_SIZE) = dw_out; + } else { + float2 dw_out; + dw_out.x = local_dw[li * VEC_SIZE + 0]; + dw_out.y = local_dw[li * VEC_SIZE + 1]; + *reinterpret_cast(partial_dw + poff + i * VEC_SIZE) = dw_out; + } + } +} + +template +__device__ __forceinline__ float AddGradSumLikeMaterializedAdd( + float grad_sum, float rms_dx) { + const float rounded_rms_dx = RoundFloatToScalarType(rms_dx); + return RoundFloatToScalarType(grad_sum + rounded_rms_dx); +} + +template +__global__ void addRMSNormBackwardPersistentKernel( + T* __restrict__ d_add, + float* __restrict__ partial_dw, + const T* __restrict__ grad_sum, + const T* __restrict__ grad_norm, + const T* __restrict__ sum_input, + const WeightDType* __restrict__ weight, + const float epsilon, + const int64_t hidden_size, + const int64_t num_tokens, + const int64_t d_add_row_stride, + const int64_t grad_sum_row_stride, + const int64_t grad_norm_row_stride, + const int64_t sum_row_stride +) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + + const int64_t tid = threadIdx.x; + const int64_t num_blocks = gridDim.x; + const int64_t vec_hidden = hidden_size / VEC_SIZE; + const float inv_hidden = 1.0f / static_cast(hidden_size); + + __shared__ float s_inv_rms; + __shared__ float s_c; + + float local_dw[MAX_VEC_PER_THREAD * VEC_SIZE]; + #pragma unroll + for (int j = 0; j < MAX_VEC_PER_THREAD * VEC_SIZE; ++j) local_dw[j] = 0.0f; + + for (int64_t row = blockIdx.x; row < num_tokens; row += num_blocks) { + const int64_t d_add_off = row * d_add_row_stride; + const int64_t grad_sum_off = row * grad_sum_row_stride; + const int64_t grad_norm_off = row * grad_norm_row_stride; + const int64_t sum_off = row * sum_row_stride; + float sq = 0.0f, dp = 0.0f; + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec xv = VecTraits::load(sum_input + sum_off + i * VEC_SIZE); + Vec dv = VecTraits::load(grad_norm + grad_norm_off + i * VEC_SIZE); + float x_vals[VEC_SIZE], d_vals[VEC_SIZE], w_vals[VEC_SIZE]; + VecToFloats(xv, x_vals); + VecToFloats(dv, d_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + if constexpr (ZeroCentered) w_vals[j] += 1.0f; + sq += x_vals[j] * x_vals[j]; + dp += d_vals[j] * w_vals[j] * x_vals[j]; + } + } + sq = blockReduceSum(sq); __syncthreads(); + dp = blockReduceSum(dp); + if (tid == 0) { + float ve = sq * inv_hidden + epsilon; + s_inv_rms = rsqrtf(ve); + s_c = dp / (static_cast(hidden_size) * ve); + } + __syncthreads(); + const float irms = s_inv_rms, cv = s_c; + + int li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x, ++li) { + Vec xv = VecTraits::load(sum_input + sum_off + i * VEC_SIZE); + Vec dv = VecTraits::load(grad_norm + grad_norm_off + i * VEC_SIZE); + Vec gs = VecTraits::load(grad_sum + grad_sum_off + i * VEC_SIZE); + float x_vals[VEC_SIZE], d_vals[VEC_SIZE], gs_vals[VEC_SIZE]; + float w_vals[VEC_SIZE], out_vals[VEC_SIZE]; + VecToFloats(xv, x_vals); + VecToFloats(dv, d_vals); + VecToFloats(gs, gs_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + if constexpr (ZeroCentered) w_vals[j] += 1.0f; + const float rms_dx = + irms * (d_vals[j] * w_vals[j] - x_vals[j] * cv); + out_vals[j] = + AddGradSumLikeMaterializedAdd(gs_vals[j], rms_dx); + local_dw[li * VEC_SIZE + j] += d_vals[j] * x_vals[j] * irms; + } + VecTraits::store(d_add + d_add_off + i * VEC_SIZE, + FloatsToVec(out_vals)); + } + } + + const int64_t poff = blockIdx.x * hidden_size; + int li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x, ++li) { + if constexpr (VEC_SIZE == 4) { + float4 dw_out; + dw_out.x = local_dw[li * VEC_SIZE + 0]; + dw_out.y = local_dw[li * VEC_SIZE + 1]; + dw_out.z = local_dw[li * VEC_SIZE + 2]; + dw_out.w = local_dw[li * VEC_SIZE + 3]; + *reinterpret_cast(partial_dw + poff + i * VEC_SIZE) = dw_out; + } else { + float2 dw_out; + dw_out.x = local_dw[li * VEC_SIZE + 0]; + dw_out.y = local_dw[li * VEC_SIZE + 1]; + *reinterpret_cast(partial_dw + poff + i * VEC_SIZE) = dw_out; + } + } +} + +template< + typename T, + typename WeightDType, + int MAX_VEC_PER_THREAD, + bool ZeroCentered, + bool FuseGradSum> +__global__ void rmsnormBackwardCachedKernel( + T* __restrict__ dx, + float* __restrict__ partial_dw, + const T* __restrict__ dy, + const T* __restrict__ input, + const WeightDType* __restrict__ weight, + const T* __restrict__ grad_sum, + const float epsilon, + const int64_t hidden_size, + const int64_t num_tokens, + const int64_t dx_row_stride, + const int64_t dy_row_stride, + const int64_t input_row_stride, + const int64_t grad_sum_row_stride +) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + + const int64_t tid = threadIdx.x; + const int64_t num_blocks = gridDim.x; + const int64_t vec_hidden = hidden_size / VEC_SIZE; + const float inv_hidden = 1.0f / static_cast(hidden_size); + + __shared__ float s_inv_rms; + __shared__ float s_c; + + float local_dw[MAX_VEC_PER_THREAD * VEC_SIZE]; + #pragma unroll + for (int j = 0; j < MAX_VEC_PER_THREAD * VEC_SIZE; ++j) { + local_dw[j] = 0.0f; + } + + for (int64_t row = blockIdx.x; row < num_tokens; row += num_blocks) { + const int64_t dx_off = row * dx_row_stride; + const int64_t dy_off = row * dy_row_stride; + const int64_t input_off = row * input_row_stride; + const int64_t grad_sum_off = row * grad_sum_row_stride; + float sq = 0.0f, dp = 0.0f; + + float x_cache[MAX_VEC_PER_THREAD * VEC_SIZE]; + float d_cache[MAX_VEC_PER_THREAD * VEC_SIZE]; + float w_cache[MAX_VEC_PER_THREAD * VEC_SIZE]; + int cache_count = 0; + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + float x_vals[VEC_SIZE], d_vals[VEC_SIZE], w_vals[VEC_SIZE]; + Vec xv = VecTraits::load(input + input_off + i * VEC_SIZE); + Vec dv = VecTraits::load(dy + dy_off + i * VEC_SIZE); + VecToFloats(xv, x_vals); + VecToFloats(dv, d_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + if constexpr (ZeroCentered) w_vals[j] += 1.0f; + x_cache[cache_count * VEC_SIZE + j] = x_vals[j]; + d_cache[cache_count * VEC_SIZE + j] = d_vals[j]; + w_cache[cache_count * VEC_SIZE + j] = w_vals[j]; + sq += x_vals[j] * x_vals[j]; + dp += d_vals[j] * w_vals[j] * x_vals[j]; + } + ++cache_count; + } + + sq = blockReduceSum(sq); __syncthreads(); + dp = blockReduceSum(dp); + if (tid == 0) { + float ve = sq * inv_hidden + epsilon; + s_inv_rms = rsqrtf(ve); + s_c = dp / (static_cast(hidden_size) * ve); + } + __syncthreads(); + const float irms = s_inv_rms, cv = s_c; + + int li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x, ++li) { + float out_vals[VEC_SIZE]; + float gs_vals[VEC_SIZE]; + if constexpr (FuseGradSum) { + Vec gs_vec = + VecTraits::load(grad_sum + grad_sum_off + i * VEC_SIZE); + VecToFloats(gs_vec, gs_vals); + } + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + const int cache_idx = li * VEC_SIZE + j; + const float x = x_cache[cache_idx]; + const float d = d_cache[cache_idx]; + const float rms_dx = irms * (d * w_cache[cache_idx] - x * cv); + if constexpr (FuseGradSum) { + out_vals[j] = + AddGradSumLikeMaterializedAdd(gs_vals[j], rms_dx); + } else { + out_vals[j] = rms_dx; + } + local_dw[cache_idx] += d * x * irms; + } + VecTraits::store(dx + dx_off + i * VEC_SIZE, + FloatsToVec(out_vals)); + } + __syncthreads(); + } + + const int64_t poff = blockIdx.x * hidden_size; + int li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x, ++li) { + if constexpr (VEC_SIZE == 4) { + float4 dw_out; + dw_out.x = local_dw[li * VEC_SIZE + 0]; + dw_out.y = local_dw[li * VEC_SIZE + 1]; + dw_out.z = local_dw[li * VEC_SIZE + 2]; + dw_out.w = local_dw[li * VEC_SIZE + 3]; + *reinterpret_cast(partial_dw + poff + i * VEC_SIZE) = dw_out; + } else { + float2 dw_out; + dw_out.x = local_dw[li * VEC_SIZE + 0]; + dw_out.y = local_dw[li * VEC_SIZE + 1]; + *reinterpret_cast(partial_dw + poff + i * VEC_SIZE) = dw_out; + } + } +} + +template +void LaunchRMSNormBackwardKernel( + const at::Tensor& dx, + float* partial_dw, + const at::Tensor& dy, + const at::Tensor& input, + const at::Tensor& weight, + float epsilon, + int64_t hidden_size, + int64_t num_tokens, + int64_t num_blocks, + int64_t dx_row_stride, + int64_t dy_row_stride, + int64_t input_row_stride, + int64_t block_size, + cudaStream_t stream) { + const int64_t vec_hidden = hidden_size / VecTraits::vec_size; + rmsnormBackwardPersistentKernel<<>>( + static_cast(dx.data_ptr()), partial_dw, + static_cast(dy.data_ptr()), static_cast(input.data_ptr()), + static_cast(weight.data_ptr()), epsilon, hidden_size, num_tokens, + dx_row_stride, dy_row_stride, input_row_stride); +} + +template +void LaunchAddRMSNormBackwardKernel( + const at::Tensor& d_add, + float* partial_dw, + const at::Tensor& grad_sum, + const at::Tensor& grad_norm, + const at::Tensor& sum_input, + const at::Tensor& weight, + float epsilon, + int64_t hidden_size, + int64_t num_tokens, + int64_t num_blocks, + int64_t d_add_row_stride, + int64_t grad_sum_row_stride, + int64_t grad_norm_row_stride, + int64_t sum_row_stride, + int64_t block_size, + cudaStream_t stream) { + addRMSNormBackwardPersistentKernel + <<>>( + static_cast(d_add.data_ptr()), partial_dw, + static_cast(grad_sum.data_ptr()), + static_cast(grad_norm.data_ptr()), + static_cast(sum_input.data_ptr()), + static_cast(weight.data_ptr()), + epsilon, hidden_size, num_tokens, d_add_row_stride, + grad_sum_row_stride, grad_norm_row_stride, sum_row_stride); +} + +template +__global__ void addRMSNormBackwardRstdCachedKernel( + T* __restrict__ d_add, + float* __restrict__ partial_dw, + const T* __restrict__ grad_sum, + const T* __restrict__ grad_norm, + const T* __restrict__ sum_input, + const WeightDType* __restrict__ weight, + const float* __restrict__ rstd, + const int64_t hidden_size, + const int64_t num_tokens, + const int64_t d_add_row_stride, + const int64_t grad_sum_row_stride, + const int64_t grad_norm_row_stride, + const int64_t sum_row_stride +) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + + const int64_t tid = threadIdx.x; + const int64_t num_blocks = gridDim.x; + const int64_t vec_hidden = hidden_size / VEC_SIZE; + const float inv_hidden = 1.0f / static_cast(hidden_size); + + __shared__ float s_norm_scale; + + float local_dw[MAX_VEC_PER_THREAD * VEC_SIZE]; + #pragma unroll + for (int j = 0; j < MAX_VEC_PER_THREAD * VEC_SIZE; ++j) local_dw[j] = 0.0f; + + for (int64_t row = blockIdx.x; row < num_tokens; row += num_blocks) { + const int64_t d_add_off = row * d_add_row_stride; + const int64_t grad_sum_off = row * grad_sum_row_stride; + const int64_t grad_norm_off = row * grad_norm_row_stride; + const int64_t sum_off = row * sum_row_stride; + const float row_rstd = rstd[row]; + float dp = 0.0f; + float x_cache[MAX_VEC_PER_THREAD * VEC_SIZE]; + float d_cache[MAX_VEC_PER_THREAD * VEC_SIZE]; + float w_cache[MAX_VEC_PER_THREAD * VEC_SIZE]; + + int li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec xv = VecTraits::load(sum_input + sum_off + i * VEC_SIZE); + Vec dv = VecTraits::load(grad_norm + grad_norm_off + i * VEC_SIZE); + float x_vals[VEC_SIZE], d_vals[VEC_SIZE], w_vals[VEC_SIZE]; + VecToFloats(xv, x_vals); + VecToFloats(dv, d_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + if constexpr (ZeroCentered) w_vals[j] += 1.0f; + const int cache_idx = li * VEC_SIZE + j; + x_cache[cache_idx] = x_vals[j]; + d_cache[cache_idx] = d_vals[j]; + w_cache[cache_idx] = w_vals[j]; + dp += d_vals[j] * w_vals[j] * x_vals[j]; + local_dw[cache_idx] += d_vals[j] * x_vals[j] * row_rstd; + } + ++li; + } + dp = blockReduceSum(dp); + if (tid == 0) { + s_norm_scale = dp * inv_hidden * row_rstd * row_rstd; + } + __syncthreads(); + const float norm_scale = s_norm_scale; + + li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x, ++li) { + Vec gs = VecTraits::load(grad_sum + grad_sum_off + i * VEC_SIZE); + float x_vals[VEC_SIZE], d_vals[VEC_SIZE], gs_vals[VEC_SIZE]; + float w_vals[VEC_SIZE], out_vals[VEC_SIZE]; + VecToFloats(gs, gs_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + const int cache_idx = li * VEC_SIZE + j; + x_vals[j] = x_cache[cache_idx]; + d_vals[j] = d_cache[cache_idx]; + w_vals[j] = w_cache[cache_idx]; + const float rms_dx = + row_rstd * (d_vals[j] * w_vals[j] - x_vals[j] * norm_scale); + out_vals[j] = + AddGradSumLikeMaterializedAdd(gs_vals[j], rms_dx); + } + VecTraits::store(d_add + d_add_off + i * VEC_SIZE, + FloatsToVec(out_vals)); + } + } + + const int64_t poff = blockIdx.x * hidden_size; + int li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x, ++li) { + if constexpr (VEC_SIZE == 4) { + float4 dw_out; + dw_out.x = local_dw[li * VEC_SIZE + 0]; + dw_out.y = local_dw[li * VEC_SIZE + 1]; + dw_out.z = local_dw[li * VEC_SIZE + 2]; + dw_out.w = local_dw[li * VEC_SIZE + 3]; + *reinterpret_cast(partial_dw + poff + i * VEC_SIZE) = dw_out; + } else { + float2 dw_out; + dw_out.x = local_dw[li * VEC_SIZE + 0]; + dw_out.y = local_dw[li * VEC_SIZE + 1]; + *reinterpret_cast(partial_dw + poff + i * VEC_SIZE) = dw_out; + } + } +} + +template +__global__ void addRMSNormBackwardRstdDirectVec4Kernel( + T* __restrict__ d_add, + float* __restrict__ partial_dw, + const T* __restrict__ grad_sum, + const T* __restrict__ grad_norm, + const T* __restrict__ sum_input, + const WeightDType* __restrict__ weight, + const float* __restrict__ rstd, + const int64_t hidden_size, + const int64_t num_tokens, + const int64_t d_add_row_stride, + const int64_t grad_sum_row_stride, + const int64_t grad_norm_row_stride, + const int64_t sum_row_stride +) { + constexpr int VEC4_SIZE = 4; + static_assert(!std::is_same_v, + "DirectVec4 fast path is intended for fp16/bf16 tensors"); + + const int64_t tid = threadIdx.x; + const int64_t num_blocks = gridDim.x; + const int64_t vec4_hidden = hidden_size / VEC4_SIZE; + const float inv_hidden = 1.0f / static_cast(hidden_size); + + __shared__ float s_norm_scale; + + float w_cache[MAX_VEC4_PER_THREAD * VEC4_SIZE]; + float local_dw[MAX_VEC4_PER_THREAD * VEC4_SIZE]; + #pragma unroll + for (int j = 0; j < MAX_VEC4_PER_THREAD * VEC4_SIZE; ++j) { + w_cache[j] = 0.0f; + local_dw[j] = 0.0f; + } + + int w_li = 0; + for (int64_t i = tid; i < vec4_hidden; i += blockDim.x, ++w_li) { + const int64_t col = i * VEC4_SIZE; + LoadWeightFloats( + weight, col, w_cache + w_li * VEC4_SIZE); + #pragma unroll + for (int j = 0; j < VEC4_SIZE; ++j) { + if constexpr (ZeroCentered) { + w_cache[w_li * VEC4_SIZE + j] += 1.0f; + } + } + } + + for (int64_t row = blockIdx.x; row < num_tokens; row += num_blocks) { + const int64_t d_add_off = row * d_add_row_stride; + const int64_t grad_sum_off = row * grad_sum_row_stride; + const int64_t grad_norm_off = row * grad_norm_row_stride; + const int64_t sum_off = row * sum_row_stride; + const float row_rstd = rstd[row]; + float dp = 0.0f; + + int li = 0; + for (int64_t i = tid; i < vec4_hidden; i += blockDim.x, ++li) { + const int64_t col = i * VEC4_SIZE; + float x_vals[VEC4_SIZE], d_vals[VEC4_SIZE]; + Load4Floats(sum_input + sum_off + col, x_vals); + Load4Floats(grad_norm + grad_norm_off + col, d_vals); + #pragma unroll + for (int j = 0; j < VEC4_SIZE; ++j) { + const float w = w_cache[li * VEC4_SIZE + j]; + dp += d_vals[j] * w * x_vals[j]; + local_dw[li * VEC4_SIZE + j] += + d_vals[j] * x_vals[j] * row_rstd; + } + } + dp = blockReduceSum(dp); + if (tid == 0) { + s_norm_scale = dp * inv_hidden * row_rstd * row_rstd; + } + __syncthreads(); + const float norm_scale = s_norm_scale; + + li = 0; + for (int64_t i = tid; i < vec4_hidden; i += blockDim.x, ++li) { + const int64_t col = i * VEC4_SIZE; + float x_vals[VEC4_SIZE], d_vals[VEC4_SIZE], gs_vals[VEC4_SIZE]; + float out_vals[VEC4_SIZE]; + Load4Floats(sum_input + sum_off + col, x_vals); + Load4Floats(grad_norm + grad_norm_off + col, d_vals); + Load4Floats(grad_sum + grad_sum_off + col, gs_vals); + #pragma unroll + for (int j = 0; j < VEC4_SIZE; ++j) { + const float w = w_cache[li * VEC4_SIZE + j]; + const float rms_dx = + row_rstd * (d_vals[j] * w - x_vals[j] * norm_scale); + out_vals[j] = + AddGradSumLikeMaterializedAdd(gs_vals[j], rms_dx); + } + Store4Scalars(d_add + d_add_off + col, out_vals); + } + } + + const int64_t poff = blockIdx.x * hidden_size; + int li = 0; + for (int64_t i = tid; i < vec4_hidden; i += blockDim.x, ++li) { + const int64_t col = i * VEC4_SIZE; + *reinterpret_cast(partial_dw + poff + col) = + make_float4(local_dw[li * VEC4_SIZE + 0], + local_dw[li * VEC4_SIZE + 1], + local_dw[li * VEC4_SIZE + 2], + local_dw[li * VEC4_SIZE + 3]); + } +} + +template +__global__ void addRMSNormBackwardRstdD2560Vec4CpAsyncKernel( + T* __restrict__ d_add, + float* __restrict__ partial_dw, + const T* __restrict__ grad_sum, + const T* __restrict__ grad_norm, + const T* __restrict__ sum_input, + const WeightDType* __restrict__ weight, + const float* __restrict__ rstd, + const int64_t num_tokens, + const int64_t d_add_row_stride, + const int64_t grad_sum_row_stride, + const int64_t grad_norm_row_stride, + const int64_t sum_row_stride +) { + static_assert(!std::is_same_v, + "D2560 vec4 cp.async fast path is for fp16/bf16 tensors"); + constexpr int HIDDEN_SIZE = 2560; + constexpr int THREADS = 128; + constexpr int VEC_ELEMS = 4; + constexpr int VEC_STRIDE = THREADS * VEC_ELEMS; + constexpr int VEC_PER_THREAD = HIDDEN_SIZE / VEC_STRIDE; + constexpr int VALUES_PER_THREAD = VEC_PER_THREAD * VEC_ELEMS; + constexpr float INV_HIDDEN = 1.0f / static_cast(HIDDEN_SIZE); + static_assert(VEC_PER_THREAD == 5); + + extern __shared__ __align__(16) unsigned char smem_bytes[]; + T* smem = reinterpret_cast(smem_bytes); + T* sx0 = smem; + T* sd0 = sx0 + HIDDEN_SIZE; + T* sx1 = sd0 + HIDDEN_SIZE; + T* sd1 = sx1 + HIDDEN_SIZE; + + const int tid = threadIdx.x; + const int64_t num_blocks = gridDim.x; + __shared__ float s_norm_scale; + + float w_cache[VALUES_PER_THREAD]; + float local_dw[VALUES_PER_THREAD]; + #pragma unroll + for (int j = 0; j < VALUES_PER_THREAD; ++j) { + w_cache[j] = 0.0f; + local_dw[j] = 0.0f; + } + + #pragma unroll + for (int k = 0; k < VEC_PER_THREAD; ++k) { + const int col = tid * VEC_ELEMS + k * VEC_STRIDE; + LoadWeight4Floats( + weight, col, w_cache + k * VEC_ELEMS); + #pragma unroll + for (int j = 0; j < VEC_ELEMS; ++j) { + if constexpr (ZeroCentered) { + w_cache[k * VEC_ELEMS + j] += 1.0f; + } + } + } + + int stage = 0; + const int64_t first_row = blockIdx.x; + if (first_row < num_tokens) { + T* sx = sx0; + T* sd = sd0; + const int64_t sum_off = first_row * sum_row_stride; + const int64_t grad_norm_off = first_row * grad_norm_row_stride; + #pragma unroll + for (int k = 0; k < VEC_PER_THREAD; ++k) { + const int col = tid * VEC_ELEMS + k * VEC_STRIDE; + CpAsync8(sx + col, sum_input + sum_off + col); + CpAsync8(sd + col, grad_norm + grad_norm_off + col); + } + CpAsyncCommitGroup(); + } + + for (int64_t row = first_row; row < num_tokens; row += num_blocks) { + const int next_stage = stage ^ 1; + const int64_t next_row = row + num_blocks; + if (next_row < num_tokens) { + T* sx_next = next_stage == 0 ? sx0 : sx1; + T* sd_next = next_stage == 0 ? sd0 : sd1; + const int64_t next_sum_off = next_row * sum_row_stride; + const int64_t next_grad_norm_off = next_row * grad_norm_row_stride; + #pragma unroll + for (int k = 0; k < VEC_PER_THREAD; ++k) { + const int col = tid * VEC_ELEMS + k * VEC_STRIDE; + CpAsync8(sx_next + col, sum_input + next_sum_off + col); + CpAsync8(sd_next + col, grad_norm + next_grad_norm_off + col); + } + CpAsyncCommitGroup(); + CpAsyncWaitGroup<1>(); + } else { + CpAsyncWaitGroup<0>(); + } + + T* sx = stage == 0 ? sx0 : sx1; + T* sd = stage == 0 ? sd0 : sd1; + const float row_rstd = rstd[row]; + float dp = 0.0f; + uint2 x_raw[VEC_PER_THREAD]; + uint2 d_raw[VEC_PER_THREAD]; + + #pragma unroll + for (int k = 0; k < VEC_PER_THREAD; ++k) { + const int col = tid * VEC_ELEMS + k * VEC_STRIDE; + float x_vals[VEC_ELEMS], d_vals[VEC_ELEMS]; + const uint2 xr = Load4Raw(sx + col); + const uint2 dr = Load4Raw(sd + col); + x_raw[k] = xr; + d_raw[k] = dr; + Raw4ToFloats(xr, x_vals); + Raw4ToFloats(dr, d_vals); + #pragma unroll + for (int j = 0; j < VEC_ELEMS; ++j) { + const int idx = k * VEC_ELEMS + j; + const float x = x_vals[j]; + const float d = d_vals[j]; + dp += d * w_cache[idx] * x; + local_dw[idx] += d * x * row_rstd; + } + } + + dp = blockReduceSum(dp); + if (tid == 0) { + s_norm_scale = dp * INV_HIDDEN * row_rstd * row_rstd; + } + __syncthreads(); + const float norm_scale = s_norm_scale; + + const int64_t d_add_off = row * d_add_row_stride; + const int64_t grad_sum_off = row * grad_sum_row_stride; + #pragma unroll + for (int k = 0; k < VEC_PER_THREAD; ++k) { + const int col = tid * VEC_ELEMS + k * VEC_STRIDE; + float x_vals[VEC_ELEMS], d_vals[VEC_ELEMS], gs_vals[VEC_ELEMS]; + float out_vals[VEC_ELEMS]; + Raw4ToFloats(x_raw[k], x_vals); + Raw4ToFloats(d_raw[k], d_vals); + Raw4ToFloats(Load4Raw(grad_sum + grad_sum_off + col), gs_vals); + #pragma unroll + for (int j = 0; j < VEC_ELEMS; ++j) { + const int idx = k * VEC_ELEMS + j; + const float rms_dx = + row_rstd * (d_vals[j] * w_cache[idx] - x_vals[j] * norm_scale); + out_vals[j] = + AddGradSumLikeMaterializedAdd(gs_vals[j], rms_dx); + } + Store4RawFromFloats(d_add + d_add_off + col, out_vals); + } + stage = next_stage; + } + + const int64_t poff = blockIdx.x * HIDDEN_SIZE; + #pragma unroll + for (int k = 0; k < VEC_PER_THREAD; ++k) { + const int col = tid * VEC_ELEMS + k * VEC_STRIDE; + *reinterpret_cast(partial_dw + poff + col) = + make_float4(local_dw[k * VEC_ELEMS + 0], + local_dw[k * VEC_ELEMS + 1], + local_dw[k * VEC_ELEMS + 2], + local_dw[k * VEC_ELEMS + 3]); + } +} + +template +__global__ void addRMSNormBackwardRstdDirectKernel( + T* __restrict__ d_add, + float* __restrict__ partial_dw, + const T* __restrict__ grad_sum, + const T* __restrict__ grad_norm, + const T* __restrict__ sum_input, + const WeightDType* __restrict__ weight, + const float* __restrict__ rstd, + const int64_t hidden_size, + const int64_t num_tokens, + const int64_t d_add_row_stride, + const int64_t grad_sum_row_stride, + const int64_t grad_norm_row_stride, + const int64_t sum_row_stride +) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + + const int64_t tid = threadIdx.x; + const int64_t num_blocks = gridDim.x; + const int64_t vec_hidden = hidden_size / VEC_SIZE; + const float inv_hidden = 1.0f / static_cast(hidden_size); + + __shared__ float s_norm_scale; + + float local_dw[MAX_VEC_PER_THREAD * VEC_SIZE]; + #pragma unroll + for (int j = 0; j < MAX_VEC_PER_THREAD * VEC_SIZE; ++j) local_dw[j] = 0.0f; + + for (int64_t row = blockIdx.x; row < num_tokens; row += num_blocks) { + const int64_t d_add_off = row * d_add_row_stride; + const int64_t grad_sum_off = row * grad_sum_row_stride; + const int64_t grad_norm_off = row * grad_norm_row_stride; + const int64_t sum_off = row * sum_row_stride; + const float row_rstd = rstd[row]; + float dp = 0.0f; + + int li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x, ++li) { + Vec xv = VecTraits::load(sum_input + sum_off + i * VEC_SIZE); + Vec dv = VecTraits::load(grad_norm + grad_norm_off + i * VEC_SIZE); + float x_vals[VEC_SIZE], d_vals[VEC_SIZE], w_vals[VEC_SIZE]; + VecToFloats(xv, x_vals); + VecToFloats(dv, d_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + if constexpr (ZeroCentered) w_vals[j] += 1.0f; + dp += d_vals[j] * w_vals[j] * x_vals[j]; + local_dw[li * VEC_SIZE + j] += d_vals[j] * x_vals[j] * row_rstd; + } + } + dp = blockReduceSum(dp); + if (tid == 0) { + s_norm_scale = dp * inv_hidden * row_rstd * row_rstd; + } + __syncthreads(); + const float norm_scale = s_norm_scale; + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec xv = VecTraits::load(sum_input + sum_off + i * VEC_SIZE); + Vec dv = VecTraits::load(grad_norm + grad_norm_off + i * VEC_SIZE); + Vec gs = VecTraits::load(grad_sum + grad_sum_off + i * VEC_SIZE); + float x_vals[VEC_SIZE], d_vals[VEC_SIZE], gs_vals[VEC_SIZE]; + float w_vals[VEC_SIZE], out_vals[VEC_SIZE]; + VecToFloats(xv, x_vals); + VecToFloats(dv, d_vals); + VecToFloats(gs, gs_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + if constexpr (ZeroCentered) w_vals[j] += 1.0f; + const float rms_dx = + row_rstd * (d_vals[j] * w_vals[j] - x_vals[j] * norm_scale); + // Match the other backward paths (and the unfused reference): round + // rms_dx to the storage dtype before the add, like a materialized add. + out_vals[j] = + AddGradSumLikeMaterializedAdd(gs_vals[j], rms_dx); + } + VecTraits::store(d_add + d_add_off + i * VEC_SIZE, + FloatsToVec(out_vals)); + } + } + + const int64_t poff = blockIdx.x * hidden_size; + int li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x, ++li) { + if constexpr (VEC_SIZE == 4) { + float4 dw_out; + dw_out.x = local_dw[li * VEC_SIZE + 0]; + dw_out.y = local_dw[li * VEC_SIZE + 1]; + dw_out.z = local_dw[li * VEC_SIZE + 2]; + dw_out.w = local_dw[li * VEC_SIZE + 3]; + *reinterpret_cast(partial_dw + poff + i * VEC_SIZE) = dw_out; + } else { + float2 dw_out; + dw_out.x = local_dw[li * VEC_SIZE + 0]; + dw_out.y = local_dw[li * VEC_SIZE + 1]; + *reinterpret_cast(partial_dw + poff + i * VEC_SIZE) = dw_out; + } + } +} + +template +void LaunchAddRMSNormBackwardWithRstdKernel( + const at::Tensor& d_add, + float* partial_dw, + const at::Tensor& grad_sum, + const at::Tensor& grad_norm, + const at::Tensor& sum_input, + const at::Tensor& weight, + const at::Tensor& rstd, + int64_t hidden_size, + int64_t num_tokens, + int64_t num_blocks, + int64_t d_add_row_stride, + int64_t grad_sum_row_stride, + int64_t grad_norm_row_stride, + int64_t sum_row_stride, + int64_t block_size, + cudaStream_t stream) { + const int64_t vec_hidden = hidden_size / VecTraits::vec_size; + if constexpr (!std::is_same_v) { + constexpr int D2560_VEC4_CP_ASYNC_BLOCK_SIZE = 128; + constexpr int D2560_VEC4_CP_ASYNC_VEC_ELEMS = 4; + if (hidden_size == 2560 && + d_add_row_stride % D2560_VEC4_CP_ASYNC_VEC_ELEMS == 0 && + grad_sum_row_stride % D2560_VEC4_CP_ASYNC_VEC_ELEMS == 0 && + grad_norm_row_stride % D2560_VEC4_CP_ASYNC_VEC_ELEMS == 0 && + sum_row_stride % D2560_VEC4_CP_ASYNC_VEC_ELEMS == 0) { + constexpr size_t SMEM_BYTES = 2 * 2 * 2560 * sizeof(T); + addRMSNormBackwardRstdD2560Vec4CpAsyncKernel + <<>>( + static_cast(d_add.data_ptr()), partial_dw, + static_cast(grad_sum.data_ptr()), + static_cast(grad_norm.data_ptr()), + static_cast(sum_input.data_ptr()), + static_cast(weight.data_ptr()), + static_cast(rstd.data_ptr()), + num_tokens, d_add_row_stride, grad_sum_row_stride, + grad_norm_row_stride, sum_row_stride); + return; + } + + constexpr int DIRECT_VEC4_SIZE = 4; + constexpr int DIRECT_VEC4_MAX_VEC_PER_THREAD = 2; + if (hidden_size % DIRECT_VEC4_SIZE == 0 && + hidden_size <= DIRECT_VEC4_MAX_VEC_PER_THREAD * block_size * DIRECT_VEC4_SIZE) { + addRMSNormBackwardRstdDirectVec4Kernel< + T, WeightDType, DIRECT_VEC4_MAX_VEC_PER_THREAD, ZeroCentered> + <<>>( + static_cast(d_add.data_ptr()), partial_dw, + static_cast(grad_sum.data_ptr()), + static_cast(grad_norm.data_ptr()), + static_cast(sum_input.data_ptr()), + static_cast(weight.data_ptr()), + static_cast(rstd.data_ptr()), + hidden_size, num_tokens, d_add_row_stride, grad_sum_row_stride, + grad_norm_row_stride, sum_row_stride); + return; + } + + } + if (vec_hidden <= block_size * 2) { + addRMSNormBackwardRstdDirectKernel + <<>>( + static_cast(d_add.data_ptr()), partial_dw, + static_cast(grad_sum.data_ptr()), + static_cast(grad_norm.data_ptr()), + static_cast(sum_input.data_ptr()), + static_cast(weight.data_ptr()), + static_cast(rstd.data_ptr()), + hidden_size, num_tokens, d_add_row_stride, grad_sum_row_stride, + grad_norm_row_stride, sum_row_stride); + return; + } + if (vec_hidden <= block_size * 3) { + addRMSNormBackwardRstdDirectKernel + <<>>( + static_cast(d_add.data_ptr()), partial_dw, + static_cast(grad_sum.data_ptr()), + static_cast(grad_norm.data_ptr()), + static_cast(sum_input.data_ptr()), + static_cast(weight.data_ptr()), + static_cast(rstd.data_ptr()), + hidden_size, num_tokens, d_add_row_stride, grad_sum_row_stride, + grad_norm_row_stride, sum_row_stride); + return; + } + if (vec_hidden <= block_size * 4) { + addRMSNormBackwardRstdDirectKernel + <<>>( + static_cast(d_add.data_ptr()), partial_dw, + static_cast(grad_sum.data_ptr()), + static_cast(grad_norm.data_ptr()), + static_cast(sum_input.data_ptr()), + static_cast(weight.data_ptr()), + static_cast(rstd.data_ptr()), + hidden_size, num_tokens, d_add_row_stride, grad_sum_row_stride, + grad_norm_row_stride, sum_row_stride); + return; + } + if (vec_hidden <= block_size * 8) { + addRMSNormBackwardRstdDirectKernel + <<>>( + static_cast(d_add.data_ptr()), partial_dw, + static_cast(grad_sum.data_ptr()), + static_cast(grad_norm.data_ptr()), + static_cast(sum_input.data_ptr()), + static_cast(weight.data_ptr()), + static_cast(rstd.data_ptr()), + hidden_size, num_tokens, d_add_row_stride, grad_sum_row_stride, + grad_norm_row_stride, sum_row_stride); + return; + } + if (vec_hidden > block_size * 16) { + throw std::runtime_error("AddRMSNorm backward rstd path hidden size too large"); + } + addRMSNormBackwardRstdDirectKernel + <<>>( + static_cast(d_add.data_ptr()), partial_dw, + static_cast(grad_sum.data_ptr()), + static_cast(grad_norm.data_ptr()), + static_cast(sum_input.data_ptr()), + static_cast(weight.data_ptr()), + static_cast(rstd.data_ptr()), + hidden_size, num_tokens, d_add_row_stride, grad_sum_row_stride, + grad_norm_row_stride, sum_row_stride); +} + +// Final reduction: partial_dw[num_blocks, hidden_size] -> dw[hidden_size] +__global__ void rmsnormReduceDwKernel( + float* __restrict__ dw, const float* __restrict__ partial_dw, + const int64_t hidden_size, const int64_t num_blocks) { + const int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= hidden_size) return; + float sum = 0.0f; + for (int64_t b = 0; b < num_blocks; b++) sum += partial_dw[b * hidden_size + idx]; + dw[idx] = sum; +} + +__global__ void rmsnormReduceDwD128Kernel( + float* __restrict__ dw, const float* __restrict__ partial_dw, + const int64_t num_blocks) { + constexpr int D = 128; + constexpr int COLS_PER_BLOCK = 32; + constexpr int ROW_WARPS = 8; + const int lane = threadIdx.x; + const int row_warp = threadIdx.y; + const int col = static_cast(blockIdx.x) * COLS_PER_BLOCK + lane; + + float sum = 0.0f; + for (int64_t b = row_warp; b < num_blocks; b += ROW_WARPS) { + sum += partial_dw[b * D + col]; + } + + __shared__ float partial[ROW_WARPS][COLS_PER_BLOCK]; + partial[row_warp][lane] = sum; + __syncthreads(); + + if (row_warp == 0) { + float total = partial[0][lane]; + #pragma unroll + for (int i = 1; i < ROW_WARPS; ++i) { + total += partial[i][lane]; + } + dw[col] = total; + } +} + +void LaunchRMSNormReduceDw( + float* dw, const float* partial_dw, int64_t hidden_size, + int64_t num_blocks, cudaStream_t stream) { + if (hidden_size == 128) { + dim3 block(32, 8); + rmsnormReduceDwD128Kernel<<<4, block, 0, stream>>>(dw, partial_dw, num_blocks); + return; + } + const int64_t rt = 256; + rmsnormReduceDwKernel<<<(hidden_size + rt - 1) / rt, rt, 0, stream>>>( + dw, partial_dw, hidden_size, num_blocks); +} + +static int get_sm_count() { + static int sm_count = -1; + if (sm_count < 0) { + int device; cudaGetDevice(&device); + cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device); + } + return sm_count; +} + +// partial_dw_view is pre-allocated from Python side to avoid cudaMallocAsync overhead +template +void RMSNormBackwardImpl( + const at::Tensor& dx, + const at::Tensor& dw, + const at::Tensor& partial_dw_view, + const at::Tensor& dy, + const at::Tensor& input, + const at::Tensor& weight, + double epsilon) { + + ASSERT_CHECK(dw.scalar_type() == at::kFloat); + ASSERT_CHECK(partial_dw_view.scalar_type() == at::kFloat); + ASSERT_CHECK(weight.is_contiguous()); + ASSERT_CHECK(dw.is_contiguous()); + ASSERT_CHECK(partial_dw_view.is_contiguous()); + CheckSameShape(dx, input); + CheckSameShape(dy, input); + const RMSNormShape shape = GetRMSNormShape(input); + const RMSNormShape dx_shape = GetRMSNormShape(dx); + const RMSNormShape dy_shape = GetRMSNormShape(dy); + const int64_t hidden_size = shape.hidden_size; + const int64_t num_tokens = shape.num_tokens; + ASSERT_CHECK(dx_shape.hidden_size == hidden_size); + ASSERT_CHECK(dx_shape.num_tokens == num_tokens); + ASSERT_CHECK(dy_shape.hidden_size == hidden_size); + ASSERT_CHECK(dy_shape.num_tokens == num_tokens); + ASSERT_CHECK(hidden_size == weight.size(0)); + ASSERT_CHECK(hidden_size % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(shape.row_stride % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(dx_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(dy_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + constexpr int MAX_VPT = 8; + // Each thread runs ceil(vec_hidden / block_size) grid-stride iterations and writes one + // local_dw[MAX_VPT * VEC_SIZE] slot per iteration; that count must stay <= MAX_VPT or it + // overflows the per-thread stack array. Grow block_size (hardware cap 1024 threads) to hold + // the invariant, and reject inputs too large even at the cap (mirrors the rstd path guard). + const int64_t vec_hidden = hidden_size / GetVecSize(input.scalar_type()); + int64_t block_size = std::min(hidden_size, 512L); + block_size = std::min(std::max(block_size, (vec_hidden + MAX_VPT - 1) / MAX_VPT), 1024L); + ASSERT_CHECK(vec_hidden <= block_size * MAX_VPT); + const int sm_count = get_sm_count(); + const int64_t num_blocks = std::min(static_cast(sm_count) * 2, num_tokens); + ASSERT_CHECK(partial_dw_view.dim() == 2); + ASSERT_CHECK(partial_dw_view.size(0) == num_blocks); + ASSERT_CHECK(partial_dw_view.size(1) == hidden_size); + ASSERT_CHECK(dw.dim() == 1); + ASSERT_CHECK(dw.size(0) == hidden_size); + float* partial_dw = static_cast(partial_dw_view.data_ptr()); + + if (IsFloat32(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchRMSNormBackwardKernel( + dx, partial_dw, dy, input, weight, static_cast(epsilon), + hidden_size, num_tokens, num_blocks, dx_shape.row_stride, + dy_shape.row_stride, shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat16(weight.scalar_type())) + LaunchRMSNormBackwardKernel( + dx, partial_dw, dy, input, weight, static_cast(epsilon), + hidden_size, num_tokens, num_blocks, dx_shape.row_stride, + dy_shape.row_stride, shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchRMSNormBackwardKernel( + dx, partial_dw, dy, input, weight, static_cast(epsilon), + hidden_size, num_tokens, num_blocks, dx_shape.row_stride, + dy_shape.row_stride, shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchRMSNormBackwardKernel( + dx, partial_dw, dy, input, weight, static_cast(epsilon), + hidden_size, num_tokens, num_blocks, dx_shape.row_stride, + dy_shape.row_stride, shape.row_stride, block_size, stream); + else if (IsBFloat16(input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchRMSNormBackwardKernel<__nv_bfloat16, __nv_bfloat16, MAX_VPT, ZeroCentered>( + dx, partial_dw, dy, input, weight, static_cast(epsilon), + hidden_size, num_tokens, num_blocks, dx_shape.row_stride, + dy_shape.row_stride, shape.row_stride, block_size, stream); + else if (IsBFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchRMSNormBackwardKernel<__nv_bfloat16, float, MAX_VPT, ZeroCentered>( + dx, partial_dw, dy, input, weight, static_cast(epsilon), + hidden_size, num_tokens, num_blocks, dx_shape.row_stride, + dy_shape.row_stride, shape.row_stride, block_size, stream); + else throw std::runtime_error("Unsupported data type for RMSNorm backward"); + + LaunchRMSNormReduceDw( + static_cast(dw.data_ptr()), partial_dw, hidden_size, num_blocks, stream); + sync_check_cuda_error(); +} + +void RMSNormBackward( + const at::Tensor& dx, + const at::Tensor& dw, + const at::Tensor& partial_dw_view, + const at::Tensor& dy, + const at::Tensor& input, + const at::Tensor& weight, + double epsilon) { + RMSNormBackwardImpl(dx, dw, partial_dw_view, dy, input, weight, epsilon); +} + +void RMSNormZeroCenteredBackward( + const at::Tensor& dx, + const at::Tensor& dw, + const at::Tensor& partial_dw_view, + const at::Tensor& dy, + const at::Tensor& input, + const at::Tensor& weight, + double epsilon) { + RMSNormBackwardImpl(dx, dw, partial_dw_view, dy, input, weight, epsilon); +} + +template +void AddRMSNormBackwardImpl( + const at::Tensor& d_add, + const at::Tensor& dw, + const at::Tensor& partial_dw_view, + const at::Tensor& grad_sum, + const at::Tensor& grad_norm, + const at::Tensor& sum_input, + const at::Tensor& weight, + double epsilon) { + + ASSERT_CHECK(dw.scalar_type() == at::kFloat); + ASSERT_CHECK(partial_dw_view.scalar_type() == at::kFloat); + ASSERT_CHECK(weight.is_contiguous()); + ASSERT_CHECK(dw.is_contiguous()); + ASSERT_CHECK(partial_dw_view.is_contiguous()); + CheckSameShape(d_add, sum_input); + CheckSameShape(grad_sum, sum_input); + CheckSameShape(grad_norm, sum_input); + const RMSNormShape shape = GetRMSNormShape(sum_input); + const RMSNormShape d_add_shape = GetRMSNormShape(d_add); + const RMSNormShape grad_sum_shape = GetRMSNormShape(grad_sum); + const RMSNormShape grad_norm_shape = GetRMSNormShape(grad_norm); + const int64_t hidden_size = shape.hidden_size; + const int64_t num_tokens = shape.num_tokens; + ASSERT_CHECK(d_add_shape.hidden_size == hidden_size); + ASSERT_CHECK(d_add_shape.num_tokens == num_tokens); + ASSERT_CHECK(grad_sum_shape.hidden_size == hidden_size); + ASSERT_CHECK(grad_sum_shape.num_tokens == num_tokens); + ASSERT_CHECK(grad_norm_shape.hidden_size == hidden_size); + ASSERT_CHECK(grad_norm_shape.num_tokens == num_tokens); + ASSERT_CHECK(hidden_size == weight.size(0)); + ASSERT_CHECK(sum_input.scalar_type() == d_add.scalar_type()); + ASSERT_CHECK(sum_input.scalar_type() == grad_sum.scalar_type()); + ASSERT_CHECK(sum_input.scalar_type() == grad_norm.scalar_type()); + ASSERT_CHECK(hidden_size % GetVecSize(sum_input.scalar_type()) == 0); + ASSERT_CHECK(shape.row_stride % GetVecSize(sum_input.scalar_type()) == 0); + ASSERT_CHECK(d_add_shape.row_stride % GetVecSize(sum_input.scalar_type()) == 0); + ASSERT_CHECK(grad_sum_shape.row_stride % GetVecSize(sum_input.scalar_type()) == 0); + ASSERT_CHECK(grad_norm_shape.row_stride % GetVecSize(sum_input.scalar_type()) == 0); + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + // LaunchAddRMSNormBackwardKernel instantiates the persistent kernel with MAX_VPT=8; a thread + // writes one local_dw[MAX_VPT * VEC_SIZE] slot per grid-stride iteration, so keep + // ceil(vec_hidden / block_size) <= MAX_VPT or it overflows the per-thread stack array. Grow + // block_size (thread cap 1024) to hold the invariant and reject inputs too large at the cap. + constexpr int ADD_BWD_MAX_VPT = 8; + const int64_t vec_hidden = hidden_size / GetVecSize(sum_input.scalar_type()); + int64_t block_size = std::min(hidden_size, 512L); + block_size = std::min(std::max(block_size, (vec_hidden + ADD_BWD_MAX_VPT - 1) / ADD_BWD_MAX_VPT), 1024L); + ASSERT_CHECK(vec_hidden <= block_size * ADD_BWD_MAX_VPT); + const int sm_count = get_sm_count(); + const int64_t num_blocks = std::min(static_cast(sm_count) * 2, num_tokens); + ASSERT_CHECK(partial_dw_view.dim() == 2); + ASSERT_CHECK(partial_dw_view.size(0) == num_blocks); + ASSERT_CHECK(partial_dw_view.size(1) == hidden_size); + ASSERT_CHECK(dw.dim() == 1); + ASSERT_CHECK(dw.size(0) == hidden_size); + float* partial_dw = static_cast(partial_dw_view.data_ptr()); + + if (IsFloat32(sum_input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchAddRMSNormBackwardKernel( + d_add, partial_dw, grad_sum, grad_norm, sum_input, weight, + static_cast(epsilon), hidden_size, num_tokens, num_blocks, + d_add_shape.row_stride, grad_sum_shape.row_stride, + grad_norm_shape.row_stride, shape.row_stride, block_size, stream); + else if (IsFloat16(sum_input.scalar_type()) && IsFloat16(weight.scalar_type())) + LaunchAddRMSNormBackwardKernel( + d_add, partial_dw, grad_sum, grad_norm, sum_input, weight, + static_cast(epsilon), hidden_size, num_tokens, num_blocks, + d_add_shape.row_stride, grad_sum_shape.row_stride, + grad_norm_shape.row_stride, shape.row_stride, block_size, stream); + else if (IsFloat16(sum_input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchAddRMSNormBackwardKernel( + d_add, partial_dw, grad_sum, grad_norm, sum_input, weight, + static_cast(epsilon), hidden_size, num_tokens, num_blocks, + d_add_shape.row_stride, grad_sum_shape.row_stride, + grad_norm_shape.row_stride, shape.row_stride, block_size, stream); + else if (IsFloat16(sum_input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchAddRMSNormBackwardKernel( + d_add, partial_dw, grad_sum, grad_norm, sum_input, weight, + static_cast(epsilon), hidden_size, num_tokens, num_blocks, + d_add_shape.row_stride, grad_sum_shape.row_stride, + grad_norm_shape.row_stride, shape.row_stride, block_size, stream); + else if (IsBFloat16(sum_input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchAddRMSNormBackwardKernel<__nv_bfloat16, __nv_bfloat16, ZeroCentered>( + d_add, partial_dw, grad_sum, grad_norm, sum_input, weight, + static_cast(epsilon), hidden_size, num_tokens, num_blocks, + d_add_shape.row_stride, grad_sum_shape.row_stride, + grad_norm_shape.row_stride, shape.row_stride, block_size, stream); + else if (IsBFloat16(sum_input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchAddRMSNormBackwardKernel<__nv_bfloat16, float, ZeroCentered>( + d_add, partial_dw, grad_sum, grad_norm, sum_input, weight, + static_cast(epsilon), hidden_size, num_tokens, num_blocks, + d_add_shape.row_stride, grad_sum_shape.row_stride, + grad_norm_shape.row_stride, shape.row_stride, block_size, stream); + else throw std::runtime_error("Unsupported data type for AddRMSNorm backward"); + + LaunchRMSNormReduceDw( + static_cast(dw.data_ptr()), partial_dw, hidden_size, num_blocks, stream); + sync_check_cuda_error(); +} + +void AddRMSNormBackward( + const at::Tensor& d_add, + const at::Tensor& dw, + const at::Tensor& partial_dw_view, + const at::Tensor& grad_sum, + const at::Tensor& grad_norm, + const at::Tensor& sum_input, + const at::Tensor& weight, + double epsilon) { + AddRMSNormBackwardImpl( + d_add, dw, partial_dw_view, grad_sum, grad_norm, sum_input, weight, epsilon); +} + +void AddRMSNormZeroCenteredBackward( + const at::Tensor& d_add, + const at::Tensor& dw, + const at::Tensor& partial_dw_view, + const at::Tensor& grad_sum, + const at::Tensor& grad_norm, + const at::Tensor& sum_input, + const at::Tensor& weight, + double epsilon) { + AddRMSNormBackwardImpl( + d_add, dw, partial_dw_view, grad_sum, grad_norm, sum_input, weight, epsilon); +} + +template +void AddRMSNormBackwardWithRstdImpl( + const at::Tensor& d_add, + const at::Tensor& dw, + const at::Tensor& partial_dw_view, + const at::Tensor& grad_sum, + const at::Tensor& grad_norm, + const at::Tensor& sum_input, + const at::Tensor& weight, + const at::Tensor& rstd) { + + ASSERT_CHECK(dw.scalar_type() == at::kFloat); + ASSERT_CHECK(partial_dw_view.scalar_type() == at::kFloat); + ASSERT_CHECK(rstd.scalar_type() == at::kFloat); + ASSERT_CHECK(weight.is_contiguous()); + ASSERT_CHECK(dw.is_contiguous()); + ASSERT_CHECK(partial_dw_view.is_contiguous()); + ASSERT_CHECK(rstd.is_contiguous()); + CheckSameShape(d_add, sum_input); + CheckSameShape(grad_sum, sum_input); + CheckSameShape(grad_norm, sum_input); + const RMSNormShape shape = GetRMSNormShape(sum_input); + const RMSNormShape d_add_shape = GetRMSNormShape(d_add); + const RMSNormShape grad_sum_shape = GetRMSNormShape(grad_sum); + const RMSNormShape grad_norm_shape = GetRMSNormShape(grad_norm); + const int64_t hidden_size = shape.hidden_size; + const int64_t num_tokens = shape.num_tokens; + ASSERT_CHECK(rstd.dim() == 1); + ASSERT_CHECK(rstd.size(0) == num_tokens); + ASSERT_CHECK(d_add_shape.hidden_size == hidden_size); + ASSERT_CHECK(d_add_shape.num_tokens == num_tokens); + ASSERT_CHECK(grad_sum_shape.hidden_size == hidden_size); + ASSERT_CHECK(grad_sum_shape.num_tokens == num_tokens); + ASSERT_CHECK(grad_norm_shape.hidden_size == hidden_size); + ASSERT_CHECK(grad_norm_shape.num_tokens == num_tokens); + ASSERT_CHECK(hidden_size == weight.size(0)); + ASSERT_CHECK(sum_input.scalar_type() == d_add.scalar_type()); + ASSERT_CHECK(sum_input.scalar_type() == grad_sum.scalar_type()); + ASSERT_CHECK(sum_input.scalar_type() == grad_norm.scalar_type()); + ASSERT_CHECK(hidden_size % GetVecSize(sum_input.scalar_type()) == 0); + ASSERT_CHECK(shape.row_stride % GetVecSize(sum_input.scalar_type()) == 0); + ASSERT_CHECK(d_add_shape.row_stride % GetVecSize(sum_input.scalar_type()) == 0); + ASSERT_CHECK(grad_sum_shape.row_stride % GetVecSize(sum_input.scalar_type()) == 0); + ASSERT_CHECK(grad_norm_shape.row_stride % GetVecSize(sum_input.scalar_type()) == 0); + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + const int64_t block_size = std::min(hidden_size, 512L); + const int sm_count = get_sm_count(); + const int64_t num_blocks = std::min(static_cast(sm_count) * 2, num_tokens); + ASSERT_CHECK(partial_dw_view.dim() == 2); + ASSERT_CHECK(partial_dw_view.size(0) == num_blocks); + ASSERT_CHECK(partial_dw_view.size(1) == hidden_size); + ASSERT_CHECK(dw.dim() == 1); + ASSERT_CHECK(dw.size(0) == hidden_size); + float* partial_dw = static_cast(partial_dw_view.data_ptr()); + + if (IsFloat32(sum_input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchAddRMSNormBackwardWithRstdKernel( + d_add, partial_dw, grad_sum, grad_norm, sum_input, weight, rstd, + hidden_size, num_tokens, num_blocks, d_add_shape.row_stride, + grad_sum_shape.row_stride, grad_norm_shape.row_stride, + shape.row_stride, block_size, stream); + else if (IsFloat16(sum_input.scalar_type()) && IsFloat16(weight.scalar_type())) + LaunchAddRMSNormBackwardWithRstdKernel( + d_add, partial_dw, grad_sum, grad_norm, sum_input, weight, rstd, + hidden_size, num_tokens, num_blocks, d_add_shape.row_stride, + grad_sum_shape.row_stride, grad_norm_shape.row_stride, + shape.row_stride, block_size, stream); + else if (IsFloat16(sum_input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchAddRMSNormBackwardWithRstdKernel( + d_add, partial_dw, grad_sum, grad_norm, sum_input, weight, rstd, + hidden_size, num_tokens, num_blocks, d_add_shape.row_stride, + grad_sum_shape.row_stride, grad_norm_shape.row_stride, + shape.row_stride, block_size, stream); + else if (IsFloat16(sum_input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchAddRMSNormBackwardWithRstdKernel( + d_add, partial_dw, grad_sum, grad_norm, sum_input, weight, rstd, + hidden_size, num_tokens, num_blocks, d_add_shape.row_stride, + grad_sum_shape.row_stride, grad_norm_shape.row_stride, + shape.row_stride, block_size, stream); + else if (IsBFloat16(sum_input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchAddRMSNormBackwardWithRstdKernel<__nv_bfloat16, __nv_bfloat16, ZeroCentered>( + d_add, partial_dw, grad_sum, grad_norm, sum_input, weight, rstd, + hidden_size, num_tokens, num_blocks, d_add_shape.row_stride, + grad_sum_shape.row_stride, grad_norm_shape.row_stride, + shape.row_stride, block_size, stream); + else if (IsBFloat16(sum_input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchAddRMSNormBackwardWithRstdKernel<__nv_bfloat16, float, ZeroCentered>( + d_add, partial_dw, grad_sum, grad_norm, sum_input, weight, rstd, + hidden_size, num_tokens, num_blocks, d_add_shape.row_stride, + grad_sum_shape.row_stride, grad_norm_shape.row_stride, + shape.row_stride, block_size, stream); + else throw std::runtime_error("Unsupported data type for AddRMSNorm backward rstd"); + + LaunchRMSNormReduceDw( + static_cast(dw.data_ptr()), partial_dw, hidden_size, num_blocks, stream); + sync_check_cuda_error(); +} + +void AddRMSNormBackwardWithRstd( + const at::Tensor& d_add, + const at::Tensor& dw, + const at::Tensor& partial_dw_view, + const at::Tensor& grad_sum, + const at::Tensor& grad_norm, + const at::Tensor& sum_input, + const at::Tensor& weight, + const at::Tensor& rstd) { + AddRMSNormBackwardWithRstdImpl( + d_add, dw, partial_dw_view, grad_sum, grad_norm, sum_input, weight, rstd); +} + +void AddRMSNormZeroCenteredBackwardWithRstd( + const at::Tensor& d_add, + const at::Tensor& dw, + const at::Tensor& partial_dw_view, + const at::Tensor& grad_sum, + const at::Tensor& grad_norm, + const at::Tensor& sum_input, + const at::Tensor& weight, + const at::Tensor& rstd) { + AddRMSNormBackwardWithRstdImpl( + d_add, dw, partial_dw_view, grad_sum, grad_norm, sum_input, weight, rstd); +} + +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) + + +// ======================== Gated RMSNorm ======================== + +__device__ __forceinline__ int64_t TensorRowOffset( + int64_t row, const TensorRowLayout layout) { + int64_t offset = 0; + int64_t remaining = row; + for (int i = layout.ndim - 2; i >= 0; --i) { + const int64_t index = remaining % layout.sizes[i]; + remaining /= layout.sizes[i]; + offset += index * layout.strides[i]; + } + return offset; +} + +__device__ __forceinline__ int64_t TensorRowOffsetD128( + int64_t row, const TensorRowLayout layout) { + if (layout.ndim == 2 && layout.stride(1) == 1) { + return row * layout.stride(0); + } + if (layout.ndim == 4 && layout.stride(3) == 1 && + layout.stride(1) == layout.sizes[2] * layout.stride(2) && + layout.stride(0) == layout.sizes[1] * layout.stride(1)) { + return row * layout.stride(2); + } + if (layout.ndim == 4 && layout.sizes[3] == 128 && layout.stride(3) == 1 && + layout.stride(2) == 128 && + layout.stride(0) == layout.sizes[1] * layout.stride(1)) { + if (layout.sizes[2] == 32) { + return (row >> 5) * layout.stride(1) + (row & 31) * 128; + } + return (row / layout.sizes[2]) * layout.stride(1) + + (row % layout.sizes[2]) * layout.stride(2); + } + return TensorRowOffset(row, layout); +} + +constexpr int kD128LayoutGeneric = 0; +constexpr int kD128LayoutLinear = 1; +constexpr int kD128LayoutH32Padded = 2; + +int GetD128LayoutKind(const TensorRowLayout& layout) { + if (layout.ndim == 2 && layout.stride(1) == 1 && layout.stride(0) == 128) { + return kD128LayoutLinear; + } + if (layout.ndim == 4 && layout.stride(3) == 1 && layout.stride(2) == 128 && + layout.stride(1) == layout.sizes[2] * layout.stride(2) && + layout.stride(0) == layout.sizes[1] * layout.stride(1)) { + return kD128LayoutLinear; + } + if (layout.ndim == 4 && layout.sizes[2] == 32 && layout.sizes[3] == 128 && + layout.stride(3) == 1 && layout.stride(2) == 128 && + layout.stride(0) == layout.sizes[1] * layout.stride(1)) { + return kD128LayoutH32Padded; + } + return kD128LayoutGeneric; +} + +__device__ __forceinline__ int64_t TensorRowOffsetD128( + int64_t row, const TensorRowLayout layout, const int layout_kind) { + if (layout_kind == kD128LayoutLinear) { + return row * 128; + } + if (layout_kind == kD128LayoutH32Padded) { + return (row >> 5) * layout.stride(1) + (row & 31) * 128; + } + return TensorRowOffsetD128(row, layout); +} + +template +__device__ __forceinline__ int64_t TensorRowOffsetD128ByKind( + int64_t row, const TensorRowLayout layout) { + if constexpr (LAYOUT_KIND == kD128LayoutLinear) { + return row * 128; + } else if constexpr (LAYOUT_KIND == kD128LayoutH32Padded) { + return (row >> 5) * layout.stride(1) + (row & 31) * 128; + } else { + return TensorRowOffsetD128(row, layout); + } +} + +__device__ __forceinline__ float Silu(float x) { + const float sigmoid = 1.0f / (1.0f + __expf(-x)); + return x * sigmoid; +} + +__device__ __forceinline__ void SiluAndGrad( + float x, float& silu, float& grad) { + const float sigmoid = 1.0f / (1.0f + __expf(-x)); + silu = x * sigmoid; + grad = sigmoid * (1.0f + x * (1.0f - sigmoid)); +} + +template +__global__ void gatedRMSNormKernel( + T* output, + const T* input, + const T* gate, + const WeightDType* weight, + const float epsilon, + const int64_t hidden_size, + const int64_t input_row_stride, + const int64_t output_row_stride, + const TensorRowLayout gate_layout) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + + float square_sum = 0.0f; + __shared__ float inv_rms; + const int64_t tid = threadIdx.x; + const int64_t bid = blockIdx.x; + const int64_t input_offset = bid * input_row_stride; + const int64_t output_offset = bid * output_row_stride; + const int64_t gate_offset = TensorRowOffset(bid, gate_layout); + const int64_t vec_hidden = hidden_size / VEC_SIZE; + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec x_vec = VecTraits::load(input + input_offset + i * VEC_SIZE); + float x_vals[VEC_SIZE]; + VecToFloats(x_vec, x_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + square_sum += x_vals[j] * x_vals[j]; + } + } + square_sum = blockReduceSum(square_sum); + if (tid == 0) { + inv_rms = rsqrtf(square_sum / static_cast(hidden_size) + epsilon); + } + __syncthreads(); + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec x_vec = VecTraits::load(input + input_offset + i * VEC_SIZE); + Vec gate_vec = VecTraits::load(gate + gate_offset + i * VEC_SIZE); + float x_vals[VEC_SIZE], gate_vals[VEC_SIZE], w_vals[VEC_SIZE], out_vals[VEC_SIZE]; + VecToFloats(x_vec, x_vals); + VecToFloats(gate_vec, gate_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + out_vals[j] = x_vals[j] * inv_rms * w_vals[j] * Silu(gate_vals[j]); + } + VecTraits::store(output + output_offset + i * VEC_SIZE, FloatsToVec(out_vals)); + } +} + +template +__global__ void gatedRMSNormD128ForwardKernel( + T* __restrict__ output, + float* __restrict__ rstd, + const T* __restrict__ input, + const T* __restrict__ gate, + const WeightDType* __restrict__ weight, + const float epsilon, + const int64_t num_tokens, + const int64_t input_row_stride, + const int64_t output_row_stride, + const TensorRowLayout gate_layout, + const int gate_layout_kind) { + constexpr int D = 128; + const int lane = threadIdx.x & 31; + const int warp_id = threadIdx.x >> 5; + const int col_base = lane * 4; + const int64_t row = static_cast(blockIdx.x) * WARPS_PER_BLOCK + warp_id; + if (row >= num_tokens) return; + + const int64_t input_off = row * input_row_stride; + const int64_t output_off = row * output_row_stride; + const int64_t gate_off = + TensorRowOffsetD128ByKind(row, gate_layout); + + float x_vals[4]; + float gate_vals[4]; + float w_vals[4]; + Load4Floats(input + input_off + col_base, x_vals); + Load4Floats(gate + gate_off + col_base, gate_vals); + Load4Floats(weight + col_base, w_vals); + + float sq = 0.0f; + #pragma unroll + for (int k = 0; k < 4; ++k) { + sq += x_vals[k] * x_vals[k]; + } + + const float row_sq = warpReduceSum(sq); + float row_rstd = 0.0f; + if (lane == 0) { + row_rstd = rsqrtf(row_sq / static_cast(D) + epsilon); + rstd[row] = row_rstd; + } + row_rstd = __shfl_sync(0xffffffff, row_rstd, 0); + + #pragma unroll + for (int k = 0; k < 4; ++k) { + const float out = x_vals[k] * row_rstd * w_vals[k] * Silu(gate_vals[k]); + x_vals[k] = out; + } + Store4Scalars(output + output_off + col_base, x_vals); +} + +template +__global__ void gatedRMSNormD128ForwardHalfWarpVecKernel( + T* __restrict__ output, + float* __restrict__ rstd, + const T* __restrict__ input, + const T* __restrict__ gate, + const WeightDType* __restrict__ weight, + const float epsilon, + const int64_t num_tokens, + const int64_t input_row_stride, + const int64_t output_row_stride, + const TensorRowLayout gate_layout, + const int gate_layout_kind) { + constexpr int D = 128; + constexpr int LANES_PER_ROW = 16; + constexpr int ROWS_PER_WARP = 2; + constexpr int ROWS_PER_BLOCK = WARPS_PER_BLOCK * ROWS_PER_WARP; + const int lane = threadIdx.x & 31; + const int warp_id = threadIdx.x >> 5; + const int sub_lane = lane & (LANES_PER_ROW - 1); + const int row_group = lane >> 4; + const int col_base0 = sub_lane * 4; + const int col_base1 = col_base0 + 64; + const int64_t row = + static_cast(blockIdx.x) * ROWS_PER_BLOCK + + warp_id * ROWS_PER_WARP + row_group; + + __shared__ float shared_weight[D]; + if (threadIdx.x < D) { + shared_weight[threadIdx.x] = ScalarToFloat(weight[threadIdx.x]); + } + __syncthreads(); + + if (row >= num_tokens) return; + + const int64_t input_off = row * input_row_stride; + const int64_t output_off = row * output_row_stride; + const int64_t gate_off = + TensorRowOffsetD128ByKind(row, gate_layout); + + float x0[4], x1[4]; + float g0[4], g1[4]; + float w0[4], w1[4]; + Load4Floats(input + input_off + col_base0, x0); + Load4Floats(input + input_off + col_base1, x1); + Load4Floats(gate + gate_off + col_base0, g0); + Load4Floats(gate + gate_off + col_base1, g1); + #pragma unroll + for (int k = 0; k < 4; ++k) { + w0[k] = shared_weight[col_base0 + k]; + w1[k] = shared_weight[col_base1 + k]; + } + + float sq = 0.0f; + #pragma unroll + for (int k = 0; k < 4; ++k) { + sq += x0[k] * x0[k] + x1[k] * x1[k]; + } + + // A warp covers ROWS_PER_WARP rows; the other row-group may have exited at the + // num_tokens boundary, so the shuffle mask must name only this row-group's lanes + // (a full 0xffffffff mask naming exited lanes is UB). + const unsigned group_mask = 0xffffu << (row_group * LANES_PER_ROW); + const float row_sq = subwarpReduceSum(sq, group_mask); + float row_rstd = 0.0f; + if (sub_lane == 0) { + row_rstd = rsqrtf(row_sq / static_cast(D) + epsilon); + rstd[row] = row_rstd; + } + row_rstd = __shfl_sync(group_mask, row_rstd, 0, LANES_PER_ROW); + + #pragma unroll + for (int k = 0; k < 4; ++k) { + x0[k] = x0[k] * row_rstd * w0[k] * Silu(g0[k]); + x1[k] = x1[k] * row_rstd * w1[k] * Silu(g1[k]); + } + Store4Scalars(output + output_off + col_base0, x0); + Store4Scalars(output + output_off + col_base1, x1); +} + +template +__global__ void gatedRMSNormD128ForwardQuarterWarpVecKernel( + T* __restrict__ output, + float* __restrict__ rstd, + const T* __restrict__ input, + const T* __restrict__ gate, + const WeightDType* __restrict__ weight, + const float epsilon, + const int64_t num_tokens, + const int64_t input_row_stride, + const int64_t output_row_stride, + const TensorRowLayout gate_layout, + const int gate_layout_kind) { + constexpr int D = 128; + constexpr int LANES_PER_ROW = 8; + constexpr int ROWS_PER_WARP = 4; + constexpr int ROWS_PER_BLOCK = WARPS_PER_BLOCK * ROWS_PER_WARP; + const int lane = threadIdx.x & 31; + const int warp_id = threadIdx.x >> 5; + const int sub_lane = lane & (LANES_PER_ROW - 1); + const int row_group = lane >> 3; + const int64_t row = + static_cast(blockIdx.x) * ROWS_PER_BLOCK + + warp_id * ROWS_PER_WARP + row_group; + if (row >= num_tokens) return; + + const int64_t input_off = row * input_row_stride; + const int64_t output_off = row * output_row_stride; + const int64_t gate_off = + TensorRowOffsetD128ByKind(row, gate_layout); + + float x_vals[4][4]; + float sq = 0.0f; + #pragma unroll + for (int q = 0; q < 4; ++q) { + const int col_base = q * 32 + sub_lane * 4; + Load4Floats(input + input_off + col_base, x_vals[q]); + #pragma unroll + for (int k = 0; k < 4; ++k) { + sq += x_vals[q][k] * x_vals[q][k]; + } + } + + // Same partial-warp hazard as the HalfWarp kernel: 4 row-groups share a warp and + // groups past num_tokens have exited, so restrict shuffles to this group's lanes. + const unsigned group_mask = 0xffu << (row_group * LANES_PER_ROW); + const float row_sq = subwarpReduceSum(sq, group_mask); + float row_rstd = 0.0f; + if (sub_lane == 0) { + row_rstd = rsqrtf(row_sq / static_cast(D) + epsilon); + rstd[row] = row_rstd; + } + row_rstd = __shfl_sync(group_mask, row_rstd, 0, LANES_PER_ROW); + + #pragma unroll + for (int q = 0; q < 4; ++q) { + const int col_base = q * 32 + sub_lane * 4; + float gate_vals[4]; + float w_vals[4]; + Load4Floats(gate + gate_off + col_base, gate_vals); + Load4Floats(weight + col_base, w_vals); + #pragma unroll + for (int k = 0; k < 4; ++k) { + x_vals[q][k] = + x_vals[q][k] * row_rstd * w_vals[k] * Silu(gate_vals[k]); + } + Store4Scalars(output + output_off + col_base, x_vals[q]); + } +} + +template< + typename T, typename WeightDType, int WARPS_PER_BLOCK, + int DGATE_LAYOUT_KIND, int GATE_LAYOUT_KIND> +__global__ void gatedRMSNormD128BackwardKernel( + T* __restrict__ dx, + T* __restrict__ dgate, + float* __restrict__ partial_dw, + const T* __restrict__ dy, + const T* __restrict__ input, + const T* __restrict__ gate, + const float* __restrict__ rstd, + const WeightDType* __restrict__ weight, + const int64_t num_tokens, + const int64_t dx_row_stride, + const int64_t dy_row_stride, + const int64_t input_row_stride, + const TensorRowLayout dgate_layout, + const TensorRowLayout gate_layout, + const int dgate_layout_kind, + const int gate_layout_kind) { + constexpr int D = 128; + const int lane = threadIdx.x & 31; + const int warp_id = threadIdx.x >> 5; + const int col_base = lane * 4; + const int64_t num_blocks = gridDim.x; + const int64_t rows_per_block = (num_tokens + num_blocks - 1) / num_blocks; + const int64_t row_start = static_cast(blockIdx.x) * rows_per_block; + const int64_t row_end = + (row_start + rows_per_block < num_tokens) ? row_start + rows_per_block : num_tokens; + + float local_dw[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + float w_vals[4]; + Load4Floats(weight + col_base, w_vals); + + for (int64_t row = row_start + warp_id; row < row_end; row += WARPS_PER_BLOCK) { + const int64_t dx_off = row * dx_row_stride; + const int64_t dy_off = row * dy_row_stride; + const int64_t input_off = row * input_row_stride; + const int64_t dgate_off = + TensorRowOffsetD128ByKind(row, dgate_layout); + const int64_t gate_off = + TensorRowOffsetD128ByKind(row, gate_layout); + const float row_rstd = rstd[row]; + + float x_vals[4]; + float gate_vals[4]; + float dy_vals[4]; + float xhat_vals[4]; + float grad_vals[4]; + float wdy_vals[4]; + float c1_sum = 0.0f; + Load4Floats(input + input_off + col_base, x_vals); + Load4Floats(gate + gate_off + col_base, gate_vals); + Load4Floats(dy + dy_off + col_base, dy_vals); + + #pragma unroll + for (int k = 0; k < 4; ++k) { + const float x_val = x_vals[k]; + const float gate_val = gate_vals[k]; + const float dy_val = dy_vals[k]; + const float w_val = w_vals[k]; + float gate_silu; + float gate_grad; + SiluAndGrad(gate_val, gate_silu, gate_grad); + + const float xhat = x_val * row_rstd; + const float wdy = dy_val * gate_silu * w_val; + xhat_vals[k] = xhat; + dy_vals[k] = dy_val; + w_vals[k] = w_val; + grad_vals[k] = gate_grad; + wdy_vals[k] = wdy; + c1_sum += xhat * wdy; + local_dw[k] += dy_val * xhat * gate_silu; + } + + const float c1 = __shfl_sync(0xffffffff, warpReduceSum(c1_sum), 0) / + static_cast(D); + + #pragma unroll + for (int k = 0; k < 4; ++k) { + const float dx_val = (wdy_vals[k] - xhat_vals[k] * c1) * row_rstd; + const float dgate_val = dy_vals[k] * xhat_vals[k] * w_vals[k] * grad_vals[k]; + x_vals[k] = dx_val; + gate_vals[k] = dgate_val; + } + Store4Scalars(dx + dx_off + col_base, x_vals); + Store4Scalars(dgate + dgate_off + col_base, gate_vals); + } + + __shared__ float shared_dw[WARPS_PER_BLOCK][32][4]; + #pragma unroll + for (int k = 0; k < 4; ++k) { + shared_dw[warp_id][lane][k] = local_dw[k]; + } + __syncthreads(); + + if (warp_id == 0) { + #pragma unroll + for (int k = 0; k < 4; ++k) { + float sum = 0.0f; + #pragma unroll + for (int w = 0; w < WARPS_PER_BLOCK; ++w) { + sum += shared_dw[w][lane][k]; + } + const int col = col_base + k; + partial_dw[static_cast(blockIdx.x) * D + col] = sum; + } + } +} + +template +void LaunchGatedRMSNormKernel( + const at::Tensor& output, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& weight, + float epsilon, + int64_t hidden_size, + int64_t num_tokens, + int64_t input_row_stride, + int64_t output_row_stride, + TensorRowLayout gate_layout, + int64_t block_size, + cudaStream_t stream) { + gatedRMSNormKernel<<>>( + static_cast(output.data_ptr()), + static_cast(input.data_ptr()), + static_cast(gate.data_ptr()), + static_cast(weight.data_ptr()), + epsilon, hidden_size, input_row_stride, output_row_stride, gate_layout); +} + +template +void LaunchGatedRMSNormD128ForwardKernel( + const at::Tensor& output, + const at::Tensor& rstd, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& weight, + float epsilon, + int64_t num_tokens, + int64_t input_row_stride, + int64_t output_row_stride, + TensorRowLayout gate_layout, + int gate_layout_kind, + cudaStream_t stream) { + const int64_t blocks = (num_tokens + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + gatedRMSNormD128ForwardKernel + <<>>( + static_cast(output.data_ptr()), + static_cast(rstd.data_ptr()), + static_cast(input.data_ptr()), + static_cast(gate.data_ptr()), + static_cast(weight.data_ptr()), + epsilon, num_tokens, input_row_stride, output_row_stride, gate_layout, + gate_layout_kind); +} + +template +void LaunchGatedRMSNormD128ForwardQuarterWarpVecKernel( + const at::Tensor& output, + const at::Tensor& rstd, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& weight, + float epsilon, + int64_t num_tokens, + int64_t input_row_stride, + int64_t output_row_stride, + TensorRowLayout gate_layout, + int gate_layout_kind, + cudaStream_t stream) { + constexpr int ROWS_PER_BLOCK = WARPS_PER_BLOCK * 4; + const int64_t blocks = (num_tokens + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK; + gatedRMSNormD128ForwardQuarterWarpVecKernel< + T, WeightDType, WARPS_PER_BLOCK, GATE_LAYOUT_KIND> + <<>>( + static_cast(output.data_ptr()), + static_cast(rstd.data_ptr()), + static_cast(input.data_ptr()), + static_cast(gate.data_ptr()), + static_cast(weight.data_ptr()), + epsilon, num_tokens, input_row_stride, output_row_stride, gate_layout, + gate_layout_kind); +} + +template +void LaunchGatedRMSNormD128ForwardHalfWarpVecKernel( + const at::Tensor& output, + const at::Tensor& rstd, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& weight, + float epsilon, + int64_t num_tokens, + int64_t input_row_stride, + int64_t output_row_stride, + TensorRowLayout gate_layout, + int gate_layout_kind, + cudaStream_t stream) { + constexpr int ROWS_PER_BLOCK = WARPS_PER_BLOCK * 2; + const int64_t blocks = (num_tokens + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK; + gatedRMSNormD128ForwardHalfWarpVecKernel< + T, WeightDType, WARPS_PER_BLOCK, GATE_LAYOUT_KIND> + <<>>( + static_cast(output.data_ptr()), + static_cast(rstd.data_ptr()), + static_cast(input.data_ptr()), + static_cast(gate.data_ptr()), + static_cast(weight.data_ptr()), + epsilon, num_tokens, input_row_stride, output_row_stride, gate_layout, + gate_layout_kind); +} + +template +void DispatchGatedRMSNormD128ForwardKernelForLayout( + const at::Tensor& output, + const at::Tensor& rstd, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& weight, + float epsilon, + int64_t num_tokens, + int64_t input_row_stride, + int64_t output_row_stride, + TensorRowLayout gate_layout, + int gate_layout_kind, + int64_t warps_per_block, + cudaStream_t stream) { + if (warps_per_block <= 4) { + LaunchGatedRMSNormD128ForwardQuarterWarpVecKernel< + T, WeightDType, 4, GATE_LAYOUT_KIND>( + output, rstd, input, gate, weight, epsilon, num_tokens, + input_row_stride, output_row_stride, gate_layout, gate_layout_kind, stream); + } else if (warps_per_block <= 8) { + LaunchGatedRMSNormD128ForwardHalfWarpVecKernel< + T, WeightDType, 8, GATE_LAYOUT_KIND>( + output, rstd, input, gate, weight, epsilon, num_tokens, + input_row_stride, output_row_stride, gate_layout, gate_layout_kind, stream); + } else { + LaunchGatedRMSNormD128ForwardKernel( + output, rstd, input, gate, weight, epsilon, num_tokens, + input_row_stride, output_row_stride, gate_layout, gate_layout_kind, stream); + } +} + +template +void DispatchGatedRMSNormD128ForwardKernel( + const at::Tensor& output, + const at::Tensor& rstd, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& weight, + float epsilon, + int64_t num_tokens, + int64_t input_row_stride, + int64_t output_row_stride, + TensorRowLayout gate_layout, + int gate_layout_kind, + int64_t warps_per_block, + cudaStream_t stream) { + if (gate_layout_kind == kD128LayoutLinear) { + DispatchGatedRMSNormD128ForwardKernelForLayout< + T, WeightDType, kD128LayoutLinear>( + output, rstd, input, gate, weight, epsilon, num_tokens, + input_row_stride, output_row_stride, gate_layout, gate_layout_kind, + warps_per_block, stream); + } else if (gate_layout_kind == kD128LayoutH32Padded) { + DispatchGatedRMSNormD128ForwardKernelForLayout< + T, WeightDType, kD128LayoutH32Padded>( + output, rstd, input, gate, weight, epsilon, num_tokens, + input_row_stride, output_row_stride, gate_layout, gate_layout_kind, + warps_per_block, stream); + } else { + DispatchGatedRMSNormD128ForwardKernelForLayout< + T, WeightDType, kD128LayoutGeneric>( + output, rstd, input, gate, weight, epsilon, num_tokens, + input_row_stride, output_row_stride, gate_layout, gate_layout_kind, + warps_per_block, stream); + } +} + +void GatedRMSNorm(const at::Tensor& output, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& weight, + double epsilon) { + ASSERT_CHECK(weight.dim() == 1); + ASSERT_CHECK(weight.is_contiguous()); + ASSERT_CHECK(input.scalar_type() == gate.scalar_type()); + ASSERT_CHECK(input.element_size() == gate.element_size()); + CheckSameShape(output, input); + CheckSameShape(gate, input); + const RMSNormShape shape = GetRMSNormShape(input); + const RMSNormShape output_shape = GetRMSNormShape(output); + const TensorRowLayout gate_layout = GetTensorRowLayout(gate); + const int64_t hidden_size = shape.hidden_size; + const int64_t num_tokens = shape.num_tokens; + ASSERT_CHECK(output_shape.hidden_size == hidden_size); + ASSERT_CHECK(output_shape.num_tokens == num_tokens); + ASSERT_CHECK(GetTensorNumTokens(gate) == num_tokens); + ASSERT_CHECK(gate.size(gate.dim() - 1) == hidden_size); + ASSERT_CHECK(hidden_size == weight.size(0)); + const int64_t vec_size = GetVecSize(input.scalar_type()); + ASSERT_CHECK(hidden_size % vec_size == 0); + ASSERT_CHECK(shape.row_stride % vec_size == 0); + ASSERT_CHECK(output_shape.row_stride % vec_size == 0); + CheckTensorRowLayoutVectorized(gate_layout, vec_size); + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + const int64_t block_size = std::min(hidden_size / vec_size, 1024L); + if (IsFloat32(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchGatedRMSNormKernel( + output, input, gate, weight, static_cast(epsilon), hidden_size, + num_tokens, shape.row_stride, output_shape.row_stride, gate_layout, + block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat16(weight.scalar_type())) + LaunchGatedRMSNormKernel( + output, input, gate, weight, static_cast(epsilon), hidden_size, + num_tokens, shape.row_stride, output_shape.row_stride, gate_layout, + block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchGatedRMSNormKernel( + output, input, gate, weight, static_cast(epsilon), hidden_size, + num_tokens, shape.row_stride, output_shape.row_stride, gate_layout, + block_size, stream); + else if (IsBFloat16(input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchGatedRMSNormKernel<__nv_bfloat16, __nv_bfloat16>( + output, input, gate, weight, static_cast(epsilon), hidden_size, + num_tokens, shape.row_stride, output_shape.row_stride, gate_layout, + block_size, stream); + else if (IsBFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchGatedRMSNormKernel<__nv_bfloat16, float>( + output, input, gate, weight, static_cast(epsilon), hidden_size, + num_tokens, shape.row_stride, output_shape.row_stride, gate_layout, + block_size, stream); + else throw std::runtime_error("Unsupported data type for Gated RMSNorm"); + sync_check_cuda_error(); +} + +void GatedRMSNormWithRstd(const at::Tensor& output, + const at::Tensor& rstd, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& weight, + double epsilon, + int64_t warps_per_block) { + ASSERT_CHECK(weight.dim() == 1); + ASSERT_CHECK(weight.is_contiguous()); + ASSERT_CHECK(input.scalar_type() == gate.scalar_type()); + ASSERT_CHECK(input.element_size() == gate.element_size()); + CheckSameShape(output, input); + CheckSameShape(gate, input); + const RMSNormShape shape = GetRMSNormShape(input); + const RMSNormShape output_shape = GetRMSNormShape(output); + const TensorRowLayout gate_layout = GetTensorRowLayout(gate); + const int64_t hidden_size = shape.hidden_size; + const int64_t num_tokens = shape.num_tokens; + ASSERT_CHECK(output_shape.hidden_size == hidden_size); + ASSERT_CHECK(output_shape.num_tokens == num_tokens); + ASSERT_CHECK(GetTensorNumTokens(gate) == num_tokens); + ASSERT_CHECK(gate.size(gate.dim() - 1) == hidden_size); + ASSERT_CHECK(hidden_size == weight.size(0)); + ASSERT_CHECK(rstd.dim() == 1); + ASSERT_CHECK(rstd.size(0) == num_tokens); + ASSERT_CHECK(rstd.is_contiguous()); + ASSERT_CHECK(rstd.scalar_type() == at::kFloat); + const int64_t vec_size = GetVecSize(input.scalar_type()); + ASSERT_CHECK(hidden_size % vec_size == 0); + ASSERT_CHECK(shape.row_stride % vec_size == 0); + ASSERT_CHECK(output_shape.row_stride % vec_size == 0); + CheckTensorRowLayoutVectorized(gate_layout, vec_size); + + if (hidden_size != 128) { + GatedRMSNorm(output, input, gate, weight, epsilon); + return; + } + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + const int gate_layout_kind = GetD128LayoutKind(gate_layout); + if (IsFloat32(input.scalar_type()) && IsFloat32(weight.scalar_type())) + DispatchGatedRMSNormD128ForwardKernel( + output, rstd, input, gate, weight, static_cast(epsilon), num_tokens, + shape.row_stride, output_shape.row_stride, gate_layout, gate_layout_kind, + warps_per_block, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat16(weight.scalar_type())) + DispatchGatedRMSNormD128ForwardKernel( + output, rstd, input, gate, weight, static_cast(epsilon), num_tokens, + shape.row_stride, output_shape.row_stride, gate_layout, gate_layout_kind, + warps_per_block, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + DispatchGatedRMSNormD128ForwardKernel( + output, rstd, input, gate, weight, static_cast(epsilon), num_tokens, + shape.row_stride, output_shape.row_stride, gate_layout, gate_layout_kind, + warps_per_block, stream); + else if (IsBFloat16(input.scalar_type()) && IsBFloat16(weight.scalar_type())) + DispatchGatedRMSNormD128ForwardKernel<__nv_bfloat16, __nv_bfloat16>( + output, rstd, input, gate, weight, static_cast(epsilon), num_tokens, + shape.row_stride, output_shape.row_stride, gate_layout, gate_layout_kind, + warps_per_block, stream); + else if (IsBFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + DispatchGatedRMSNormD128ForwardKernel<__nv_bfloat16, float>( + output, rstd, input, gate, weight, static_cast(epsilon), num_tokens, + shape.row_stride, output_shape.row_stride, gate_layout, gate_layout_kind, + warps_per_block, stream); + else throw std::runtime_error("Unsupported data type for Gated RMSNorm"); + sync_check_cuda_error(); +} + +template +__global__ void gatedRMSNormBackwardPersistentKernel( + T* __restrict__ dx, + T* __restrict__ dgate, + float* __restrict__ partial_dw, + const T* __restrict__ dy, + const T* __restrict__ input, + const T* __restrict__ gate, + const WeightDType* __restrict__ weight, + const float epsilon, + const int64_t hidden_size, + const int64_t num_tokens, + const int64_t dx_row_stride, + const int64_t dy_row_stride, + const int64_t input_row_stride, + const TensorRowLayout dgate_layout, + const TensorRowLayout gate_layout) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + + const int64_t tid = threadIdx.x; + const int64_t num_blocks = gridDim.x; + const int64_t vec_hidden = hidden_size / VEC_SIZE; + const float inv_hidden = 1.0f / static_cast(hidden_size); + + __shared__ float s_inv_rms; + __shared__ float s_c; + + float local_dw[MAX_VEC_PER_THREAD * VEC_SIZE]; + #pragma unroll + for (int j = 0; j < MAX_VEC_PER_THREAD * VEC_SIZE; j++) local_dw[j] = 0.0f; + + for (int64_t row = blockIdx.x; row < num_tokens; row += num_blocks) { + const int64_t dx_off = row * dx_row_stride; + const int64_t dy_off = row * dy_row_stride; + const int64_t input_off = row * input_row_stride; + const int64_t dgate_off = TensorRowOffset(row, dgate_layout); + const int64_t gate_off = TensorRowOffset(row, gate_layout); + float sq = 0.0f; + float dp = 0.0f; + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec x_vec = VecTraits::load(input + input_off + i * VEC_SIZE); + Vec dy_vec = VecTraits::load(dy + dy_off + i * VEC_SIZE); + Vec gate_vec = VecTraits::load(gate + gate_off + i * VEC_SIZE); + float x_vals[VEC_SIZE], dy_vals[VEC_SIZE], gate_vals[VEC_SIZE], w_vals[VEC_SIZE]; + VecToFloats(x_vec, x_vals); + VecToFloats(dy_vec, dy_vals); + VecToFloats(gate_vec, gate_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + const float gate_silu = Silu(gate_vals[j]); + const float p = dy_vals[j] * w_vals[j] * gate_silu; + sq += x_vals[j] * x_vals[j]; + dp += p * x_vals[j]; + } + } + sq = blockReduceSum(sq); + __syncthreads(); + dp = blockReduceSum(dp); + if (tid == 0) { + const float ve = sq * inv_hidden + epsilon; + s_inv_rms = rsqrtf(ve); + s_c = dp / (static_cast(hidden_size) * ve); + } + __syncthreads(); + const float irms = s_inv_rms; + const float cv = s_c; + + int li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x, li++) { + Vec x_vec = VecTraits::load(input + input_off + i * VEC_SIZE); + Vec dy_vec = VecTraits::load(dy + dy_off + i * VEC_SIZE); + Vec gate_vec = VecTraits::load(gate + gate_off + i * VEC_SIZE); + float x_vals[VEC_SIZE], dy_vals[VEC_SIZE], gate_vals[VEC_SIZE], w_vals[VEC_SIZE]; + float dx_vals[VEC_SIZE], dgate_vals[VEC_SIZE]; + VecToFloats(x_vec, x_vals); + VecToFloats(dy_vec, dy_vals); + VecToFloats(gate_vec, gate_vals); + LoadWeightFloats(weight, i * VEC_SIZE, w_vals); + #pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + float gate_silu; + float gate_grad; + SiluAndGrad(gate_vals[j], gate_silu, gate_grad); + const float normed = x_vals[j] * irms; + dx_vals[j] = irms * (dy_vals[j] * w_vals[j] * gate_silu - x_vals[j] * cv); + dgate_vals[j] = dy_vals[j] * normed * w_vals[j] * gate_grad; + local_dw[li * VEC_SIZE + j] += dy_vals[j] * normed * gate_silu; + } + VecTraits::store(dx + dx_off + i * VEC_SIZE, FloatsToVec(dx_vals)); + VecTraits::store(dgate + dgate_off + i * VEC_SIZE, FloatsToVec(dgate_vals)); + } + } + + const int64_t poff = blockIdx.x * hidden_size; + int li = 0; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x, li++) { + if constexpr (VEC_SIZE == 4) { + float4 dw_out; + dw_out.x = local_dw[li * VEC_SIZE + 0]; + dw_out.y = local_dw[li * VEC_SIZE + 1]; + dw_out.z = local_dw[li * VEC_SIZE + 2]; + dw_out.w = local_dw[li * VEC_SIZE + 3]; + *reinterpret_cast(partial_dw + poff + i * VEC_SIZE) = dw_out; + } else { + float2 dw_out; + dw_out.x = local_dw[li * VEC_SIZE + 0]; + dw_out.y = local_dw[li * VEC_SIZE + 1]; + *reinterpret_cast(partial_dw + poff + i * VEC_SIZE) = dw_out; + } + } +} + +template +void LaunchGatedRMSNormBackwardKernel( + const at::Tensor& dx, + const at::Tensor& dgate, + float* partial_dw, + const at::Tensor& dy, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& weight, + float epsilon, + int64_t hidden_size, + int64_t num_tokens, + int64_t num_blocks, + int64_t dx_row_stride, + int64_t dy_row_stride, + int64_t input_row_stride, + TensorRowLayout dgate_layout, + TensorRowLayout gate_layout, + int64_t block_size, + cudaStream_t stream) { + gatedRMSNormBackwardPersistentKernel + <<>>( + static_cast(dx.data_ptr()), + static_cast(dgate.data_ptr()), + partial_dw, + static_cast(dy.data_ptr()), + static_cast(input.data_ptr()), + static_cast(gate.data_ptr()), + static_cast(weight.data_ptr()), + epsilon, hidden_size, num_tokens, dx_row_stride, dy_row_stride, + input_row_stride, dgate_layout, gate_layout); +} + +template< + typename T, typename WeightDType, int WARPS_PER_BLOCK, + int DGATE_LAYOUT_KIND, int GATE_LAYOUT_KIND> +void LaunchGatedRMSNormD128BackwardKernel( + const at::Tensor& dx, + const at::Tensor& dgate, + float* partial_dw, + const at::Tensor& dy, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& rstd, + const at::Tensor& weight, + int64_t num_tokens, + int64_t num_blocks, + int64_t dx_row_stride, + int64_t dy_row_stride, + int64_t input_row_stride, + TensorRowLayout dgate_layout, + TensorRowLayout gate_layout, + int dgate_layout_kind, + int gate_layout_kind, + cudaStream_t stream) { + gatedRMSNormD128BackwardKernel< + T, WeightDType, WARPS_PER_BLOCK, DGATE_LAYOUT_KIND, GATE_LAYOUT_KIND> + <<>>( + static_cast(dx.data_ptr()), + static_cast(dgate.data_ptr()), + partial_dw, + static_cast(dy.data_ptr()), + static_cast(input.data_ptr()), + static_cast(gate.data_ptr()), + static_cast(rstd.data_ptr()), + static_cast(weight.data_ptr()), + num_tokens, dx_row_stride, dy_row_stride, input_row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind); +} + +template +void DispatchGatedRMSNormD128BackwardKernelForLayout( + const at::Tensor& dx, + const at::Tensor& dgate, + float* partial_dw, + const at::Tensor& dy, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& rstd, + const at::Tensor& weight, + int64_t num_tokens, + int64_t num_blocks, + int64_t dx_row_stride, + int64_t dy_row_stride, + int64_t input_row_stride, + TensorRowLayout dgate_layout, + TensorRowLayout gate_layout, + int dgate_layout_kind, + int gate_layout_kind, + int64_t warps_per_block, + cudaStream_t stream) { + if (warps_per_block <= 4) { + LaunchGatedRMSNormD128BackwardKernel< + T, WeightDType, 4, DGATE_LAYOUT_KIND, GATE_LAYOUT_KIND>( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_row_stride, dy_row_stride, input_row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, stream); + } else if (warps_per_block <= 8) { + LaunchGatedRMSNormD128BackwardKernel< + T, WeightDType, 8, DGATE_LAYOUT_KIND, GATE_LAYOUT_KIND>( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_row_stride, dy_row_stride, input_row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, stream); + } else { + LaunchGatedRMSNormD128BackwardKernel< + T, WeightDType, 16, DGATE_LAYOUT_KIND, GATE_LAYOUT_KIND>( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_row_stride, dy_row_stride, input_row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, stream); + } +} + +template +void DispatchGatedRMSNormD128BackwardKernelForDgateLayout( + const at::Tensor& dx, + const at::Tensor& dgate, + float* partial_dw, + const at::Tensor& dy, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& rstd, + const at::Tensor& weight, + int64_t num_tokens, + int64_t num_blocks, + int64_t dx_row_stride, + int64_t dy_row_stride, + int64_t input_row_stride, + TensorRowLayout dgate_layout, + TensorRowLayout gate_layout, + int dgate_layout_kind, + int gate_layout_kind, + int64_t warps_per_block, + cudaStream_t stream) { + if (gate_layout_kind == kD128LayoutLinear) { + DispatchGatedRMSNormD128BackwardKernelForLayout< + T, WeightDType, DGATE_LAYOUT_KIND, kD128LayoutLinear>( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_row_stride, dy_row_stride, input_row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, + warps_per_block, stream); + } else if (gate_layout_kind == kD128LayoutH32Padded) { + DispatchGatedRMSNormD128BackwardKernelForLayout< + T, WeightDType, DGATE_LAYOUT_KIND, kD128LayoutH32Padded>( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_row_stride, dy_row_stride, input_row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, + warps_per_block, stream); + } else { + DispatchGatedRMSNormD128BackwardKernelForLayout< + T, WeightDType, DGATE_LAYOUT_KIND, kD128LayoutGeneric>( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_row_stride, dy_row_stride, input_row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, + warps_per_block, stream); + } +} + +template +void DispatchGatedRMSNormD128BackwardKernel( + const at::Tensor& dx, + const at::Tensor& dgate, + float* partial_dw, + const at::Tensor& dy, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& rstd, + const at::Tensor& weight, + int64_t num_tokens, + int64_t num_blocks, + int64_t dx_row_stride, + int64_t dy_row_stride, + int64_t input_row_stride, + TensorRowLayout dgate_layout, + TensorRowLayout gate_layout, + int dgate_layout_kind, + int gate_layout_kind, + int64_t warps_per_block, + cudaStream_t stream) { + if (dgate_layout_kind == kD128LayoutLinear) { + DispatchGatedRMSNormD128BackwardKernelForDgateLayout< + T, WeightDType, kD128LayoutLinear>( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_row_stride, dy_row_stride, input_row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, + warps_per_block, stream); + } else if (dgate_layout_kind == kD128LayoutH32Padded) { + DispatchGatedRMSNormD128BackwardKernelForDgateLayout< + T, WeightDType, kD128LayoutH32Padded>( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_row_stride, dy_row_stride, input_row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, + warps_per_block, stream); + } else { + DispatchGatedRMSNormD128BackwardKernelForDgateLayout< + T, WeightDType, kD128LayoutGeneric>( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_row_stride, dy_row_stride, input_row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, + warps_per_block, stream); + } +} + +void GatedRMSNormBackward( + const at::Tensor& dx, + const at::Tensor& dgate, + const at::Tensor& dw, + const at::Tensor& partial_dw_view, + const at::Tensor& dy, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& weight, + double epsilon) { + ASSERT_CHECK(dw.scalar_type() == at::kFloat); + ASSERT_CHECK(partial_dw_view.scalar_type() == at::kFloat); + ASSERT_CHECK(weight.is_contiguous()); + ASSERT_CHECK(dw.is_contiguous()); + ASSERT_CHECK(partial_dw_view.is_contiguous()); + ASSERT_CHECK(input.scalar_type() == dy.scalar_type()); + ASSERT_CHECK(input.element_size() == dy.element_size()); + ASSERT_CHECK(input.scalar_type() == gate.scalar_type()); + ASSERT_CHECK(input.element_size() == gate.element_size()); + ASSERT_CHECK(input.scalar_type() == dgate.scalar_type()); + ASSERT_CHECK(input.element_size() == dgate.element_size()); + CheckSameShape(dx, input); + CheckSameShape(dy, input); + CheckSameShape(gate, input); + CheckSameShape(dgate, input); + const RMSNormShape shape = GetRMSNormShape(input); + const RMSNormShape dx_shape = GetRMSNormShape(dx); + const RMSNormShape dy_shape = GetRMSNormShape(dy); + const TensorRowLayout gate_layout = GetTensorRowLayout(gate); + const TensorRowLayout dgate_layout = GetTensorRowLayout(dgate); + const int64_t hidden_size = shape.hidden_size; + const int64_t num_tokens = shape.num_tokens; + ASSERT_CHECK(dx_shape.hidden_size == hidden_size); + ASSERT_CHECK(dx_shape.num_tokens == num_tokens); + ASSERT_CHECK(dy_shape.hidden_size == hidden_size); + ASSERT_CHECK(dy_shape.num_tokens == num_tokens); + ASSERT_CHECK(GetTensorNumTokens(gate) == num_tokens); + ASSERT_CHECK(GetTensorNumTokens(dgate) == num_tokens); + ASSERT_CHECK(hidden_size == weight.size(0)); + const int64_t vec_size = GetVecSize(input.scalar_type()); + ASSERT_CHECK(hidden_size % vec_size == 0); + ASSERT_CHECK(shape.row_stride % vec_size == 0); + ASSERT_CHECK(dx_shape.row_stride % vec_size == 0); + ASSERT_CHECK(dy_shape.row_stride % vec_size == 0); + CheckTensorRowLayoutVectorized(gate_layout, vec_size); + CheckTensorRowLayoutVectorized(dgate_layout, vec_size); + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + const int64_t block_size = std::min(hidden_size / vec_size, 1024L); + ASSERT_CHECK(partial_dw_view.dim() == 2); + ASSERT_CHECK(partial_dw_view.size(1) == hidden_size); + ASSERT_CHECK(dw.dim() == 1); + ASSERT_CHECK(dw.size(0) == hidden_size); + const int64_t num_blocks = partial_dw_view.size(0); + ASSERT_CHECK(num_blocks > 0); + ASSERT_CHECK(num_blocks <= num_tokens); + float* partial_dw = static_cast(partial_dw_view.data_ptr()); + constexpr int MAX_VPT = 8; + // Persistent kernel accumulates into local_dw[MAX_VPT * VEC_SIZE]; each thread runs + // ceil(vec_hidden / block_size) iterations, which must not exceed MAX_VPT. + ASSERT_CHECK(hidden_size / vec_size <= block_size * MAX_VPT); + + if (IsFloat32(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchGatedRMSNormBackwardKernel( + dx, dgate, partial_dw, dy, input, gate, weight, static_cast(epsilon), + hidden_size, num_tokens, num_blocks, dx_shape.row_stride, + dy_shape.row_stride, shape.row_stride, dgate_layout, gate_layout, + block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat16(weight.scalar_type())) + LaunchGatedRMSNormBackwardKernel( + dx, dgate, partial_dw, dy, input, gate, weight, static_cast(epsilon), + hidden_size, num_tokens, num_blocks, dx_shape.row_stride, + dy_shape.row_stride, shape.row_stride, dgate_layout, gate_layout, + block_size, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchGatedRMSNormBackwardKernel( + dx, dgate, partial_dw, dy, input, gate, weight, static_cast(epsilon), + hidden_size, num_tokens, num_blocks, dx_shape.row_stride, + dy_shape.row_stride, shape.row_stride, dgate_layout, gate_layout, + block_size, stream); + else if (IsBFloat16(input.scalar_type()) && IsBFloat16(weight.scalar_type())) + LaunchGatedRMSNormBackwardKernel<__nv_bfloat16, __nv_bfloat16, MAX_VPT>( + dx, dgate, partial_dw, dy, input, gate, weight, static_cast(epsilon), + hidden_size, num_tokens, num_blocks, dx_shape.row_stride, + dy_shape.row_stride, shape.row_stride, dgate_layout, gate_layout, + block_size, stream); + else if (IsBFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + LaunchGatedRMSNormBackwardKernel<__nv_bfloat16, float, MAX_VPT>( + dx, dgate, partial_dw, dy, input, gate, weight, static_cast(epsilon), + hidden_size, num_tokens, num_blocks, dx_shape.row_stride, + dy_shape.row_stride, shape.row_stride, dgate_layout, gate_layout, + block_size, stream); + else throw std::runtime_error("Unsupported data type for Gated RMSNorm backward"); + + LaunchRMSNormReduceDw( + static_cast(dw.data_ptr()), partial_dw, hidden_size, num_blocks, stream); + sync_check_cuda_error(); +} + +void GatedRMSNormBackwardWithRstd( + const at::Tensor& dx, + const at::Tensor& dgate, + const at::Tensor& dw, + const at::Tensor& partial_dw_view, + const at::Tensor& dy, + const at::Tensor& input, + const at::Tensor& gate, + const at::Tensor& rstd, + const at::Tensor& weight, + double epsilon, + int64_t warps_per_block) { + ASSERT_CHECK(dw.scalar_type() == at::kFloat); + ASSERT_CHECK(partial_dw_view.scalar_type() == at::kFloat); + ASSERT_CHECK(weight.is_contiguous()); + ASSERT_CHECK(dw.is_contiguous()); + ASSERT_CHECK(partial_dw_view.is_contiguous()); + ASSERT_CHECK(input.scalar_type() == dy.scalar_type()); + ASSERT_CHECK(input.element_size() == dy.element_size()); + ASSERT_CHECK(input.scalar_type() == gate.scalar_type()); + ASSERT_CHECK(input.element_size() == gate.element_size()); + ASSERT_CHECK(input.scalar_type() == dgate.scalar_type()); + ASSERT_CHECK(input.element_size() == dgate.element_size()); + CheckSameShape(dx, input); + CheckSameShape(dy, input); + CheckSameShape(gate, input); + CheckSameShape(dgate, input); + const RMSNormShape shape = GetRMSNormShape(input); + const RMSNormShape dx_shape = GetRMSNormShape(dx); + const RMSNormShape dy_shape = GetRMSNormShape(dy); + const TensorRowLayout gate_layout = GetTensorRowLayout(gate); + const TensorRowLayout dgate_layout = GetTensorRowLayout(dgate); + const int64_t hidden_size = shape.hidden_size; + const int64_t num_tokens = shape.num_tokens; + ASSERT_CHECK(dx_shape.hidden_size == hidden_size); + ASSERT_CHECK(dx_shape.num_tokens == num_tokens); + ASSERT_CHECK(dy_shape.hidden_size == hidden_size); + ASSERT_CHECK(dy_shape.num_tokens == num_tokens); + ASSERT_CHECK(GetTensorNumTokens(gate) == num_tokens); + ASSERT_CHECK(GetTensorNumTokens(dgate) == num_tokens); + ASSERT_CHECK(hidden_size == weight.size(0)); + ASSERT_CHECK(rstd.dim() == 1); + ASSERT_CHECK(rstd.size(0) == num_tokens); + ASSERT_CHECK(rstd.is_contiguous()); + ASSERT_CHECK(rstd.scalar_type() == at::kFloat); + const int64_t vec_size = GetVecSize(input.scalar_type()); + ASSERT_CHECK(hidden_size % vec_size == 0); + ASSERT_CHECK(shape.row_stride % vec_size == 0); + ASSERT_CHECK(dx_shape.row_stride % vec_size == 0); + ASSERT_CHECK(dy_shape.row_stride % vec_size == 0); + CheckTensorRowLayoutVectorized(gate_layout, vec_size); + CheckTensorRowLayoutVectorized(dgate_layout, vec_size); + + if (hidden_size != 128) { + GatedRMSNormBackward( + dx, dgate, dw, partial_dw_view, dy, input, gate, weight, epsilon); + return; + } + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + ASSERT_CHECK(partial_dw_view.dim() == 2); + ASSERT_CHECK(partial_dw_view.size(1) == hidden_size); + ASSERT_CHECK(dw.dim() == 1); + ASSERT_CHECK(dw.size(0) == hidden_size); + const int64_t num_blocks = partial_dw_view.size(0); + ASSERT_CHECK(num_blocks > 0); + ASSERT_CHECK(num_blocks <= num_tokens); + float* partial_dw = static_cast(partial_dw_view.data_ptr()); + const int dgate_layout_kind = GetD128LayoutKind(dgate_layout); + const int gate_layout_kind = GetD128LayoutKind(gate_layout); + + if (IsFloat32(input.scalar_type()) && IsFloat32(weight.scalar_type())) + DispatchGatedRMSNormD128BackwardKernel( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_shape.row_stride, dy_shape.row_stride, shape.row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, + warps_per_block, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat16(weight.scalar_type())) + DispatchGatedRMSNormD128BackwardKernel( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_shape.row_stride, dy_shape.row_stride, shape.row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, + warps_per_block, stream); + else if (IsFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + DispatchGatedRMSNormD128BackwardKernel( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_shape.row_stride, dy_shape.row_stride, shape.row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, + warps_per_block, stream); + else if (IsBFloat16(input.scalar_type()) && IsBFloat16(weight.scalar_type())) + DispatchGatedRMSNormD128BackwardKernel<__nv_bfloat16, __nv_bfloat16>( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_shape.row_stride, dy_shape.row_stride, shape.row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, + warps_per_block, stream); + else if (IsBFloat16(input.scalar_type()) && IsFloat32(weight.scalar_type())) + DispatchGatedRMSNormD128BackwardKernel<__nv_bfloat16, float>( + dx, dgate, partial_dw, dy, input, gate, rstd, weight, num_tokens, + num_blocks, dx_shape.row_stride, dy_shape.row_stride, shape.row_stride, + dgate_layout, gate_layout, dgate_layout_kind, gate_layout_kind, + warps_per_block, stream); + else throw std::runtime_error("Unsupported data type for Gated RMSNorm backward"); + + LaunchRMSNormReduceDw( + static_cast(dw.data_ptr()), partial_dw, hidden_size, num_blocks, stream); + sync_check_cuda_error(); +} + +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) + + +// ======================== Adaptive RMSNorm ======================== + +struct AdaRMSNormShape { + int64_t num_tokens; + int64_t hidden_size; + int64_t num_mod_rows; + int64_t tokens_per_mod; +}; + +bool SameDType(at::ScalarType lhs, at::ScalarType rhs) { + return lhs == rhs; +} + +AdaRMSNormShape GetAdaRMSNormShape(const at::Tensor& input, + const at::Tensor& modulation) { + const RMSNormShape input_shape = GetRMSNormShape(input); + const int64_t hidden_size = input_shape.hidden_size; + ASSERT_CHECK(modulation.dim() >= 1); + ASSERT_CHECK(modulation.is_contiguous()); + const int mod_ndim = modulation.dim(); + ASSERT_CHECK(modulation.size(mod_ndim - 1) == hidden_size * 3); + + if (mod_ndim == input.dim()) { + bool same_prefix = true; + for (int i = 0; i < input.dim() - 1; ++i) { + if (modulation.size(i) != input.size(i)) { + same_prefix = false; + break; + } + } + if (same_prefix) { + return {input_shape.num_tokens, hidden_size, input_shape.num_tokens, 1}; + } + + if (input.dim() == 3 && modulation.size(0) == input.size(0) && + modulation.size(1) == 1) { + return {input_shape.num_tokens, hidden_size, input.size(0), input.size(1)}; + } + } + + if (input.dim() == 3 && mod_ndim == 2 && modulation.size(0) == input.size(0)) { + return {input_shape.num_tokens, hidden_size, input.size(0), input.size(1)}; + } + + throw std::runtime_error( + "Unsupported AdaRMSNorm shapes: modulation must match input prefix or be " + "[batch, 3 * hidden] for a [batch, seq, hidden] input"); +} + +void CheckAdaGateShape(const at::Tensor& gate, + const AdaRMSNormShape& shape) { + ASSERT_CHECK(gate.dim() >= 1); + ASSERT_CHECK(gate.is_contiguous()); + ASSERT_CHECK(gate.size(gate.dim() - 1) == shape.hidden_size); + int64_t gate_rows = 1; + for (int i = 0; i < gate.dim() - 1; ++i) { + gate_rows *= gate.size(i); + } + ASSERT_CHECK(gate_rows == shape.num_mod_rows); +} + +template +__global__ void adarmsnormKernel( + T* __restrict__ output, + T* __restrict__ gate, + const T* __restrict__ input, + const ModDType* __restrict__ modulation, + const float epsilon, + const int64_t hidden_size, + const int64_t tokens_per_mod, + const int64_t input_row_stride, + const int64_t output_row_stride) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + + float square_sum = 0.0f; + __shared__ float inv_rms; + const int64_t tid = threadIdx.x; + const int64_t row = blockIdx.x; + const int64_t input_offset = row * input_row_stride; + const int64_t output_offset = row * output_row_stride; + const int64_t mod_row = row / tokens_per_mod; + const int64_t mod_base = mod_row * hidden_size * 3; + const int64_t vec_hidden = hidden_size / VEC_SIZE; + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec x_vec = VecTraits::load(input + input_offset + i * VEC_SIZE); + float x_vals[VEC_SIZE]; + VecToFloats(x_vec, x_vals); +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + square_sum += x_vals[j] * x_vals[j]; + } + } + square_sum = blockReduceSum(square_sum); + if (tid == 0) inv_rms = rsqrtf(square_sum / static_cast(hidden_size) + epsilon); + __syncthreads(); + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec in_vec = VecTraits::load(input + input_offset + i * VEC_SIZE); + float x_vals[VEC_SIZE], scale_vals[VEC_SIZE], shift_vals[VEC_SIZE], out_vals[VEC_SIZE]; + VecToFloats(in_vec, x_vals); + LoadWeightFloats(modulation, mod_base + i * VEC_SIZE, scale_vals); + LoadWeightFloats( + modulation, mod_base + hidden_size + i * VEC_SIZE, shift_vals); +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + out_vals[j] = x_vals[j] * inv_rms * (1.0f + scale_vals[j]) + shift_vals[j]; + } + VecTraits::store(output + output_offset + i * VEC_SIZE, FloatsToVec(out_vals)); + } + + if ((row % tokens_per_mod) == 0) { + const int64_t gate_offset = mod_row * hidden_size; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + float gate_vals[VEC_SIZE]; + LoadWeightFloats( + modulation, mod_base + hidden_size * 2 + i * VEC_SIZE, gate_vals); + VecTraits::store(gate + gate_offset + i * VEC_SIZE, FloatsToVec(gate_vals)); + } + } +} + +template +void LaunchAdaRMSNormKernel( + const at::Tensor& output, + const at::Tensor& gate, + const at::Tensor& input, + const at::Tensor& modulation, + float epsilon, + int64_t hidden_size, + int64_t num_tokens, + int64_t tokens_per_mod, + int64_t input_row_stride, + int64_t output_row_stride, + int64_t block_size, + cudaStream_t stream) { + adarmsnormKernel<<>>( + static_cast(output.data_ptr()), static_cast(gate.data_ptr()), + static_cast(input.data_ptr()), + static_cast(modulation.data_ptr()), epsilon, hidden_size, + tokens_per_mod, input_row_stride, output_row_stride); +} + +template +void DispatchAdaRMSNormKernel( + const at::Tensor& output, + const at::Tensor& gate, + const at::Tensor& input, + const at::Tensor& modulation, + float epsilon, + int64_t hidden_size, + int64_t num_tokens, + int64_t tokens_per_mod, + int64_t input_row_stride, + int64_t output_row_stride, + int64_t block_size, + cudaStream_t stream) { + if (IsFloat32(modulation.scalar_type())) + LaunchAdaRMSNormKernel( + output, gate, input, modulation, epsilon, hidden_size, num_tokens, + tokens_per_mod, input_row_stride, output_row_stride, block_size, stream); + else if (IsFloat16(modulation.scalar_type())) + LaunchAdaRMSNormKernel( + output, gate, input, modulation, epsilon, hidden_size, num_tokens, + tokens_per_mod, input_row_stride, output_row_stride, block_size, stream); + else if (IsBFloat16(modulation.scalar_type())) + LaunchAdaRMSNormKernel( + output, gate, input, modulation, epsilon, hidden_size, num_tokens, + tokens_per_mod, input_row_stride, output_row_stride, block_size, stream); + else throw std::runtime_error("Unsupported modulation data type for AdaRMSNorm"); +} + +void AdaRMSNorm(const at::Tensor& output, const at::Tensor& gate, + const at::Tensor& input, const at::Tensor& modulation, + double epsilon) { + ASSERT_CHECK(SameDType(output.scalar_type(), input.scalar_type())); + ASSERT_CHECK(SameDType(gate.scalar_type(), input.scalar_type())); + CheckSameShape(output, input); + const AdaRMSNormShape shape = GetAdaRMSNormShape(input, modulation); + const RMSNormShape input_shape = GetRMSNormShape(input); + const RMSNormShape output_shape = GetRMSNormShape(output); + ASSERT_CHECK(output_shape.hidden_size == input_shape.hidden_size); + ASSERT_CHECK(output_shape.num_tokens == input_shape.num_tokens); + CheckAdaGateShape(gate, shape); + ASSERT_CHECK(shape.hidden_size % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(input_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(output_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + const int64_t block_size = std::min(shape.hidden_size, 1024L); + + if (IsFloat32(input.scalar_type())) + DispatchAdaRMSNormKernel( + output, gate, input, modulation, static_cast(epsilon), + shape.hidden_size, shape.num_tokens, shape.tokens_per_mod, + input_shape.row_stride, output_shape.row_stride, block_size, stream); + else if (IsFloat16(input.scalar_type())) + DispatchAdaRMSNormKernel( + output, gate, input, modulation, static_cast(epsilon), + shape.hidden_size, shape.num_tokens, shape.tokens_per_mod, + input_shape.row_stride, output_shape.row_stride, block_size, stream); + else if (IsBFloat16(input.scalar_type())) + DispatchAdaRMSNormKernel<__nv_bfloat16>( + output, gate, input, modulation, static_cast(epsilon), + shape.hidden_size, shape.num_tokens, shape.tokens_per_mod, + input_shape.row_stride, output_shape.row_stride, block_size, stream); + else throw std::runtime_error("Unsupported data type for AdaRMSNorm"); + sync_check_cuda_error(); +} + +// Removed TVM FFI export (see pybind11 registration) + +template +__global__ void adarmsnormBackwardKernel( + T* __restrict__ dx, + float* __restrict__ dmod, + const T* __restrict__ dy, + const T* __restrict__ dgate, + const T* __restrict__ input, + const ModDType* __restrict__ modulation, + const float epsilon, + const int64_t hidden_size, + const int64_t tokens_per_mod, + const int64_t dx_row_stride, + const int64_t dy_row_stride, + const int64_t input_row_stride) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + + const int64_t tid = threadIdx.x; + const int64_t row = blockIdx.x; + const int64_t dx_offset = row * dx_row_stride; + const int64_t dy_offset = row * dy_row_stride; + const int64_t input_offset = row * input_row_stride; + const int64_t mod_row = row / tokens_per_mod; + const int64_t mod_base = mod_row * hidden_size * 3; + const int64_t gate_offset = mod_row * hidden_size; + const int64_t vec_hidden = hidden_size / VEC_SIZE; + const float inv_hidden = 1.0f / static_cast(hidden_size); + + __shared__ float s_inv_rms; + __shared__ float s_c; + + float sq = 0.0f, dp = 0.0f; + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec xv = VecTraits::load(input + input_offset + i * VEC_SIZE); + Vec dv = VecTraits::load(dy + dy_offset + i * VEC_SIZE); + float x_vals[VEC_SIZE], d_vals[VEC_SIZE], scale_vals[VEC_SIZE]; + VecToFloats(xv, x_vals); + VecToFloats(dv, d_vals); + LoadWeightFloats(modulation, mod_base + i * VEC_SIZE, scale_vals); +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + const float z = d_vals[j] * (1.0f + scale_vals[j]); + sq += x_vals[j] * x_vals[j]; + dp += z * x_vals[j]; + } + } + sq = blockReduceSum(sq); __syncthreads(); + dp = blockReduceSum(dp); + if (tid == 0) { + float ve = sq * inv_hidden + epsilon; + s_inv_rms = rsqrtf(ve); + s_c = dp / (static_cast(hidden_size) * ve); + } + __syncthreads(); + const float irms = s_inv_rms, cv = s_c; + + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec xv = VecTraits::load(input + input_offset + i * VEC_SIZE); + Vec dv = VecTraits::load(dy + dy_offset + i * VEC_SIZE); + float x_vals[VEC_SIZE], d_vals[VEC_SIZE], scale_vals[VEC_SIZE], dx_vals[VEC_SIZE]; + VecToFloats(xv, x_vals); + VecToFloats(dv, d_vals); + LoadWeightFloats(modulation, mod_base + i * VEC_SIZE, scale_vals); +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + const int64_t elem = i * VEC_SIZE + j; + const float normed = x_vals[j] * irms; + const float z = d_vals[j] * (1.0f + scale_vals[j]); + dx_vals[j] = irms * (z - x_vals[j] * cv); + const float dscale = d_vals[j] * normed; + const float dshift = d_vals[j]; + if (tokens_per_mod == 1) { + dmod[mod_base + elem] = dscale; + dmod[mod_base + hidden_size + elem] = dshift; + } else { + atomicAdd(dmod + mod_base + elem, dscale); + atomicAdd(dmod + mod_base + hidden_size + elem, dshift); + } + } + VecTraits::store(dx + dx_offset + i * VEC_SIZE, FloatsToVec(dx_vals)); + } + + if ((row % tokens_per_mod) == 0) { + for (int64_t i = tid; i < vec_hidden; i += blockDim.x) { + Vec gv = VecTraits::load(dgate + gate_offset + i * VEC_SIZE); + float gate_vals[VEC_SIZE]; + VecToFloats(gv, gate_vals); +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + dmod[mod_base + hidden_size * 2 + i * VEC_SIZE + j] = gate_vals[j]; + } + } + } +} + +template +void LaunchAdaRMSNormBackwardKernel( + const at::Tensor& dx, + const at::Tensor& dmod, + const at::Tensor& dy, + const at::Tensor& dgate, + const at::Tensor& input, + const at::Tensor& modulation, + float epsilon, + int64_t hidden_size, + int64_t num_tokens, + int64_t tokens_per_mod, + int64_t dx_row_stride, + int64_t dy_row_stride, + int64_t input_row_stride, + int64_t block_size, + cudaStream_t stream) { + adarmsnormBackwardKernel<<>>( + static_cast(dx.data_ptr()), static_cast(dmod.data_ptr()), + static_cast(dy.data_ptr()), static_cast(dgate.data_ptr()), + static_cast(input.data_ptr()), + static_cast(modulation.data_ptr()), epsilon, hidden_size, + tokens_per_mod, dx_row_stride, dy_row_stride, input_row_stride); +} + +template +void DispatchAdaRMSNormBackwardKernel( + const at::Tensor& dx, + const at::Tensor& dmod, + const at::Tensor& dy, + const at::Tensor& dgate, + const at::Tensor& input, + const at::Tensor& modulation, + float epsilon, + int64_t hidden_size, + int64_t num_tokens, + int64_t tokens_per_mod, + int64_t dx_row_stride, + int64_t dy_row_stride, + int64_t input_row_stride, + int64_t block_size, + cudaStream_t stream) { + if (IsFloat32(modulation.scalar_type())) + LaunchAdaRMSNormBackwardKernel( + dx, dmod, dy, dgate, input, modulation, epsilon, hidden_size, + num_tokens, tokens_per_mod, dx_row_stride, dy_row_stride, + input_row_stride, block_size, stream); + else if (IsFloat16(modulation.scalar_type())) + LaunchAdaRMSNormBackwardKernel( + dx, dmod, dy, dgate, input, modulation, epsilon, hidden_size, + num_tokens, tokens_per_mod, dx_row_stride, dy_row_stride, + input_row_stride, block_size, stream); + else if (IsBFloat16(modulation.scalar_type())) + LaunchAdaRMSNormBackwardKernel( + dx, dmod, dy, dgate, input, modulation, epsilon, hidden_size, + num_tokens, tokens_per_mod, dx_row_stride, dy_row_stride, + input_row_stride, block_size, stream); + else throw std::runtime_error("Unsupported modulation data type for AdaRMSNorm backward"); +} + +void AdaRMSNormBackward( + const at::Tensor& dx, + const at::Tensor& dmod, + const at::Tensor& dy, + const at::Tensor& dgate, + const at::Tensor& input, + const at::Tensor& modulation, + double epsilon) { + ASSERT_CHECK(dmod.scalar_type() == at::kFloat); + ASSERT_CHECK(dmod.is_contiguous()); + ASSERT_CHECK(dgate.is_contiguous()); + ASSERT_CHECK(modulation.is_contiguous()); + ASSERT_CHECK(SameDType(dx.scalar_type(), input.scalar_type())); + ASSERT_CHECK(SameDType(dy.scalar_type(), input.scalar_type())); + ASSERT_CHECK(SameDType(dgate.scalar_type(), input.scalar_type())); + CheckSameShape(dx, input); + CheckSameShape(dy, input); + CheckSameShape(dmod, modulation); + const AdaRMSNormShape shape = GetAdaRMSNormShape(input, modulation); + const RMSNormShape input_shape = GetRMSNormShape(input); + const RMSNormShape dx_shape = GetRMSNormShape(dx); + const RMSNormShape dy_shape = GetRMSNormShape(dy); + ASSERT_CHECK(dx_shape.hidden_size == input_shape.hidden_size); + ASSERT_CHECK(dx_shape.num_tokens == input_shape.num_tokens); + ASSERT_CHECK(dy_shape.hidden_size == input_shape.hidden_size); + ASSERT_CHECK(dy_shape.num_tokens == input_shape.num_tokens); + CheckAdaGateShape(dgate, shape); + ASSERT_CHECK(shape.hidden_size % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(input_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(dx_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + ASSERT_CHECK(dy_shape.row_stride % GetVecSize(input.scalar_type()) == 0); + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + const int64_t block_size = std::min(shape.hidden_size, 1024L); + + if (IsFloat32(input.scalar_type())) + DispatchAdaRMSNormBackwardKernel( + dx, dmod, dy, dgate, input, modulation, static_cast(epsilon), + shape.hidden_size, shape.num_tokens, shape.tokens_per_mod, + dx_shape.row_stride, dy_shape.row_stride, input_shape.row_stride, + block_size, stream); + else if (IsFloat16(input.scalar_type())) + DispatchAdaRMSNormBackwardKernel( + dx, dmod, dy, dgate, input, modulation, static_cast(epsilon), + shape.hidden_size, shape.num_tokens, shape.tokens_per_mod, + dx_shape.row_stride, dy_shape.row_stride, input_shape.row_stride, + block_size, stream); + else if (IsBFloat16(input.scalar_type())) + DispatchAdaRMSNormBackwardKernel<__nv_bfloat16>( + dx, dmod, dy, dgate, input, modulation, static_cast(epsilon), + shape.hidden_size, shape.num_tokens, shape.tokens_per_mod, + dx_shape.row_stride, dy_shape.row_stride, input_shape.row_stride, + block_size, stream); + else throw std::runtime_error("Unsupported data type for AdaRMSNorm backward"); + sync_check_cuda_error(); +} + +// Removed TVM FFI export (see pybind11 registration) + +} // namespace tvm_ffi_rmsnorm diff --git a/wall_x/model/core/ops/csrc/rope/rope_kernel.cu b/wall_x/model/core/ops/csrc/rope/rope_kernel.cu index f9ef0a69..858b1957 100644 --- a/wall_x/model/core/ops/csrc/rope/rope_kernel.cu +++ b/wall_x/model/core/ops/csrc/rope/rope_kernel.cu @@ -7,15 +7,36 @@ #include #include +#include #include - #include +#include #include #include #include +#include #include "../common/cuda_utils.h" namespace wallx_cuda_rope { +constexpr int64_t kMaxRopeGridDimYZ = 65535; + +inline dim3 MakeRopeGrid(int64_t batch, int64_t seq_len) { + ASSERT_CHECK(batch > 0); + ASSERT_CHECK(seq_len > 0); + ASSERT_CHECK(batch <= 2147483647LL); + + const int64_t grid_z = (seq_len + kMaxRopeGridDimYZ - 1) / kMaxRopeGridDimYZ; + ASSERT_CHECK(grid_z <= kMaxRopeGridDimYZ); + const int64_t grid_y = (seq_len + grid_z - 1) / grid_z; + return dim3(static_cast(batch), + static_cast(grid_y), + static_cast(grid_z)); +} + +__device__ __forceinline__ int64_t RopeSeqBlockIdx() { + return static_cast(blockIdx.y) + static_cast(blockIdx.z) * gridDim.y; +} + // Helper: convert vector of T to float4 template __device__ __forceinline__ float4 to_float4(const T* ptr) { @@ -66,6 +87,21 @@ __device__ __forceinline__ void from_float4(T* ptr, const float4& f4) { } } +template +__device__ __forceinline__ void copy4(T* dst, const T* src) { + if constexpr (std::is_same_v) { + reinterpret_cast(dst)[0] = reinterpret_cast(src)[0]; + reinterpret_cast(dst)[1] = reinterpret_cast(src)[1]; + } else if constexpr (std::is_same_v) { + reinterpret_cast<__nv_bfloat162*>(dst)[0] = reinterpret_cast(src)[0]; + reinterpret_cast<__nv_bfloat162*>(dst)[1] = reinterpret_cast(src)[1]; + } else if constexpr (std::is_same_v) { + reinterpret_cast(dst)[0] = reinterpret_cast(src)[0]; + } else { + static_assert(sizeof(T) == 0, "Unsupported type"); + } +} + template __global__ void RopeKernel(const float* cos, const float* sin, @@ -76,6 +112,7 @@ __global__ void RopeKernel(const float* cos, T* q_embed, T* k_embed, const int d, + const int64_t seq_len, const int qb_stride, const int qs_stride, const int qh_stride, @@ -93,9 +130,10 @@ __global__ void RopeKernel(const float* cos, float* cos_smem = reinterpret_cast(cos_sin); float* sin_smem = cos_smem + half_dim; int b = blockIdx.x; - int s = blockIdx.y; + int64_t s = RopeSeqBlockIdx(); + if (s >= seq_len) return; - int64_t cos_sin_b_stride = gridDim.y * half_dim; + int64_t cos_sin_b_stride = seq_len * half_dim; int64_t cos_sin_s_stride = half_dim; #define SIN_GMEM(b, c, d) sin[(b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] #define COS_GMEM(b, c, d) cos[(b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] @@ -210,10 +248,412 @@ __global__ void RopeKernel(const float* cos, } } +template +__device__ __forceinline__ float scalar_to_float(T value) { + if constexpr (std::is_same_v) { + return __half2float(value); + } else if constexpr (std::is_same_v) { + return __bfloat162float(value); + } else if constexpr (std::is_same_v) { + return value; + } else { + static_assert(sizeof(T) == 0, "Unsupported type"); + } +} + +template +__global__ void RopeFullKernel(const C* cos, + const C* sin, + const T* q, + const int q_h, + const T* k, + const int k_h, + T* q_embed, + T* k_embed, + const int full_d, + const int rotary_d, + const int qb_stride, + const int qs_stride, + const int qh_stride, + const int kb_stride, + const int ks_stride, + const int kh_stride, + const int qeb_stride, + const int qes_stride, + const int qeh_stride, + const int keb_stride, + const int kes_stride, + const int keh_stride, + const int64_t cos_b_stride, + const int64_t cos_s_stride, + const int64_t cos_d_stride, + const int64_t sin_b_stride, + const int64_t sin_s_stride, + const int64_t sin_d_stride) { + extern __shared__ char cos_sin[]; + const int half_dim = rotary_d / 2; + float* cos_smem = reinterpret_cast(cos_sin); + float* sin_smem = cos_smem + half_dim; + const int b = blockIdx.x; + const int s = blockIdx.y; + + for (int i = threadIdx.x; i < half_dim; i += blockDim.x) { + const int64_t cos_offset = + static_cast(b) * cos_b_stride + + static_cast(s) * cos_s_stride + + static_cast(i) * cos_d_stride; + const int64_t sin_offset = + static_cast(b) * sin_b_stride + + static_cast(s) * sin_s_stride + + static_cast(i) * sin_d_stride; + cos_smem[i] = scalar_to_float(cos[cos_offset]); + sin_smem[i] = scalar_to_float(sin[sin_offset]); + } + __syncthreads(); + + for (int j = threadIdx.x * 4; j < q_h * half_dim; j += blockDim.x * 4) { + const int h_idx = j / half_dim; + const int base_d = j % half_dim; + if (base_d + 3 >= half_dim) { + continue; + } + + const int64_t q_base = + static_cast(b) * qb_stride + + static_cast(s) * qs_stride + + static_cast(h_idx) * qh_stride; + const int64_t qe_base = + static_cast(b) * qeb_stride + + static_cast(s) * qes_stride + + static_cast(h_idx) * qeh_stride; + const float4 cos_vec = to_float4(&cos_smem[base_d]); + const float4 sin_vec = to_float4(&sin_smem[base_d]); + + if constexpr (interleave) { + const float4 q_x0 = to_float4(&q[q_base + 2 * base_d]); + const float4 q_x1 = to_float4(&q[q_base + 2 * base_d + 4]); + const float4 q0_rot = make_float4( + q_x0.x * cos_vec.x - q_x0.y * sin_vec.x, + q_x0.y * cos_vec.x + q_x0.x * sin_vec.x, + q_x0.z * cos_vec.y - q_x0.w * sin_vec.y, + q_x0.w * cos_vec.y + q_x0.z * sin_vec.y); + const float4 q1_rot = make_float4( + q_x1.x * cos_vec.z - q_x1.y * sin_vec.z, + q_x1.y * cos_vec.z + q_x1.x * sin_vec.z, + q_x1.z * cos_vec.w - q_x1.w * sin_vec.w, + q_x1.w * cos_vec.w + q_x1.z * sin_vec.w); + from_float4(&q_embed[qe_base + 2 * base_d], q0_rot); + from_float4(&q_embed[qe_base + 2 * base_d + 4], q1_rot); + } else { + const float4 q_x0 = to_float4(&q[q_base + base_d]); + const float4 q_x1 = to_float4(&q[q_base + base_d + half_dim]); + const float4 q0_rot = make_float4( + q_x0.x * cos_vec.x - q_x1.x * sin_vec.x, + q_x0.y * cos_vec.y - q_x1.y * sin_vec.y, + q_x0.z * cos_vec.z - q_x1.z * sin_vec.z, + q_x0.w * cos_vec.w - q_x1.w * sin_vec.w); + const float4 q1_rot = make_float4( + q_x1.x * cos_vec.x + q_x0.x * sin_vec.x, + q_x1.y * cos_vec.y + q_x0.y * sin_vec.y, + q_x1.z * cos_vec.z + q_x0.z * sin_vec.z, + q_x1.w * cos_vec.w + q_x0.w * sin_vec.w); + from_float4(&q_embed[qe_base + base_d], q0_rot); + from_float4(&q_embed[qe_base + base_d + half_dim], q1_rot); + } + } + + for (int j = threadIdx.x * 4; j < k_h * half_dim; j += blockDim.x * 4) { + const int h_idx = j / half_dim; + const int base_d = j % half_dim; + if (base_d + 3 >= half_dim) { + continue; + } + + const int64_t k_base = + static_cast(b) * kb_stride + + static_cast(s) * ks_stride + + static_cast(h_idx) * kh_stride; + const int64_t ke_base = + static_cast(b) * keb_stride + + static_cast(s) * kes_stride + + static_cast(h_idx) * keh_stride; + const float4 cos_vec = to_float4(&cos_smem[base_d]); + const float4 sin_vec = to_float4(&sin_smem[base_d]); + + if constexpr (interleave) { + const float4 k_x0 = to_float4(&k[k_base + 2 * base_d]); + const float4 k_x1 = to_float4(&k[k_base + 2 * base_d + 4]); + const float4 k0_rot = make_float4( + k_x0.x * cos_vec.x - k_x0.y * sin_vec.x, + k_x0.y * cos_vec.x + k_x0.x * sin_vec.x, + k_x0.z * cos_vec.y - k_x0.w * sin_vec.y, + k_x0.w * cos_vec.y + k_x0.z * sin_vec.y); + const float4 k1_rot = make_float4( + k_x1.x * cos_vec.z - k_x1.y * sin_vec.z, + k_x1.y * cos_vec.z + k_x1.x * sin_vec.z, + k_x1.z * cos_vec.w - k_x1.w * sin_vec.w, + k_x1.w * cos_vec.w + k_x1.z * sin_vec.w); + from_float4(&k_embed[ke_base + 2 * base_d], k0_rot); + from_float4(&k_embed[ke_base + 2 * base_d + 4], k1_rot); + } else { + const float4 k_x0 = to_float4(&k[k_base + base_d]); + const float4 k_x1 = to_float4(&k[k_base + base_d + half_dim]); + const float4 k0_rot = make_float4( + k_x0.x * cos_vec.x - k_x1.x * sin_vec.x, + k_x0.y * cos_vec.y - k_x1.y * sin_vec.y, + k_x0.z * cos_vec.z - k_x1.z * sin_vec.z, + k_x0.w * cos_vec.w - k_x1.w * sin_vec.w); + const float4 k1_rot = make_float4( + k_x1.x * cos_vec.x + k_x0.x * sin_vec.x, + k_x1.y * cos_vec.y + k_x0.y * sin_vec.y, + k_x1.z * cos_vec.z + k_x0.z * sin_vec.z, + k_x1.w * cos_vec.w + k_x0.w * sin_vec.w); + from_float4(&k_embed[ke_base + base_d], k0_rot); + from_float4(&k_embed[ke_base + base_d + half_dim], k1_rot); + } + } + + const int pass_dim = full_d - rotary_d; + if (pass_dim <= 0) { + return; + } + const int pass_vecs = (pass_dim + 3) / 4; + for (int linear = threadIdx.x; linear < q_h * pass_vecs; linear += blockDim.x) { + const int h_idx = linear / pass_vecs; + const int rel_vec = linear - h_idx * pass_vecs; + const int rel_d = rel_vec * 4; + const int64_t q_base = + static_cast(b) * qb_stride + + static_cast(s) * qs_stride + + static_cast(h_idx) * qh_stride + rotary_d + rel_d; + const int64_t qe_base = + static_cast(b) * qeb_stride + + static_cast(s) * qes_stride + + static_cast(h_idx) * qeh_stride + rotary_d + rel_d; + if (rel_d + 3 < pass_dim) { + copy4(&q_embed[qe_base], &q[q_base]); + } else { + for (int t = rel_d; t < pass_dim; ++t) { + q_embed[qe_base + (t - rel_d)] = q[q_base + (t - rel_d)]; + } + } + } + for (int linear = threadIdx.x; linear < k_h * pass_vecs; linear += blockDim.x) { + const int h_idx = linear / pass_vecs; + const int rel_vec = linear - h_idx * pass_vecs; + const int rel_d = rel_vec * 4; + const int64_t k_base = + static_cast(b) * kb_stride + + static_cast(s) * ks_stride + + static_cast(h_idx) * kh_stride + rotary_d + rel_d; + const int64_t ke_base = + static_cast(b) * keb_stride + + static_cast(s) * kes_stride + + static_cast(h_idx) * keh_stride + rotary_d + rel_d; + if (rel_d + 3 < pass_dim) { + copy4(&k_embed[ke_base], &k[k_base]); + } else { + for (int t = rel_d; t < pass_dim; ++t) { + k_embed[ke_base + (t - rel_d)] = k[k_base + (t - rel_d)]; + } + } + } +} + +template +__global__ void RopeFullKernelBackward( + const C* cos, + const C* sin, + const T* grad_q_embed, + const T* grad_k_embed, + T* grad_q, + T* grad_k, + const int q_h, + const int k_h, + const int full_d, + const int rotary_d, + const int qb_stride, const int qs_stride, const int qh_stride, + const int kb_stride, const int ks_stride, const int kh_stride, + const int qgb_stride, const int qgs_stride, const int qgh_stride, + const int kgb_stride, const int kgs_stride, const int kgh_stride, + const int64_t cos_b_stride, + const int64_t cos_s_stride, + const int64_t cos_d_stride, + const int64_t sin_b_stride, + const int64_t sin_s_stride, + const int64_t sin_d_stride +) { + extern __shared__ char cos_sin[]; + const int half_dim = rotary_d / 2; + float* cos_smem = reinterpret_cast(cos_sin); + float* sin_smem = cos_smem + half_dim; + const int b = blockIdx.x; + const int s = blockIdx.y; + + for (int i = threadIdx.x; i < half_dim; i += blockDim.x) { + const int64_t cos_offset = + static_cast(b) * cos_b_stride + + static_cast(s) * cos_s_stride + + static_cast(i) * cos_d_stride; + const int64_t sin_offset = + static_cast(b) * sin_b_stride + + static_cast(s) * sin_s_stride + + static_cast(i) * sin_d_stride; + cos_smem[i] = scalar_to_float(cos[cos_offset]); + sin_smem[i] = scalar_to_float(sin[sin_offset]); + } + __syncthreads(); + + for (int j = threadIdx.x * 4; j < q_h * half_dim; j += blockDim.x * 4) { + const int h_idx = j / half_dim; + const int base_d = j % half_dim; + if (base_d + 3 >= half_dim) { + continue; + } + const int64_t gqe_base = + static_cast(b) * qb_stride + + static_cast(s) * qs_stride + + static_cast(h_idx) * qh_stride; + const int64_t gq_base = + static_cast(b) * qgb_stride + + static_cast(s) * qgs_stride + + static_cast(h_idx) * qgh_stride; + const float4 cos_vec = to_float4(&cos_smem[base_d]); + const float4 sin_vec = to_float4(&sin_smem[base_d]); + if constexpr (interleave) { + const float4 gq0_embed = to_float4(&grad_q_embed[gqe_base + 2 * base_d]); + const float4 gq1_embed = to_float4(&grad_q_embed[gqe_base + 2 * base_d + 4]); + const float4 gq0 = make_float4( + gq0_embed.x * cos_vec.x + gq0_embed.y * sin_vec.x, + gq0_embed.y * cos_vec.x - gq0_embed.x * sin_vec.x, + gq0_embed.z * cos_vec.y + gq0_embed.w * sin_vec.y, + gq0_embed.w * cos_vec.y - gq0_embed.z * sin_vec.y); + const float4 gq1 = make_float4( + gq1_embed.x * cos_vec.z + gq1_embed.y * sin_vec.z, + gq1_embed.y * cos_vec.z - gq1_embed.x * sin_vec.z, + gq1_embed.z * cos_vec.w + gq1_embed.w * sin_vec.w, + gq1_embed.w * cos_vec.w - gq1_embed.z * sin_vec.w); + from_float4(&grad_q[gq_base + 2 * base_d], gq0); + from_float4(&grad_q[gq_base + 2 * base_d + 4], gq1); + } else { + const float4 gq0_rot = to_float4(&grad_q_embed[gqe_base + base_d]); + const float4 gq1_rot = to_float4(&grad_q_embed[gqe_base + base_d + half_dim]); + const float4 gq0 = make_float4( + gq0_rot.x * cos_vec.x + gq1_rot.x * sin_vec.x, + gq0_rot.y * cos_vec.y + gq1_rot.y * sin_vec.y, + gq0_rot.z * cos_vec.z + gq1_rot.z * sin_vec.z, + gq0_rot.w * cos_vec.w + gq1_rot.w * sin_vec.w); + const float4 gq1 = make_float4( + gq1_rot.x * cos_vec.x - gq0_rot.x * sin_vec.x, + gq1_rot.y * cos_vec.y - gq0_rot.y * sin_vec.y, + gq1_rot.z * cos_vec.z - gq0_rot.z * sin_vec.z, + gq1_rot.w * cos_vec.w - gq0_rot.w * sin_vec.w); + from_float4(&grad_q[gq_base + base_d], gq0); + from_float4(&grad_q[gq_base + base_d + half_dim], gq1); + } + } + + for (int j = threadIdx.x * 4; j < k_h * half_dim; j += blockDim.x * 4) { + const int h_idx = j / half_dim; + const int base_d = j % half_dim; + if (base_d + 3 >= half_dim) { + continue; + } + const int64_t gke_base = + static_cast(b) * kb_stride + + static_cast(s) * ks_stride + + static_cast(h_idx) * kh_stride; + const int64_t gk_base = + static_cast(b) * kgb_stride + + static_cast(s) * kgs_stride + + static_cast(h_idx) * kgh_stride; + const float4 cos_vec = to_float4(&cos_smem[base_d]); + const float4 sin_vec = to_float4(&sin_smem[base_d]); + if constexpr (interleave) { + const float4 gk0_embed = to_float4(&grad_k_embed[gke_base + 2 * base_d]); + const float4 gk1_embed = to_float4(&grad_k_embed[gke_base + 2 * base_d + 4]); + const float4 gk0 = make_float4( + gk0_embed.x * cos_vec.x + gk0_embed.y * sin_vec.x, + gk0_embed.y * cos_vec.x - gk0_embed.x * sin_vec.x, + gk0_embed.z * cos_vec.y + gk0_embed.w * sin_vec.y, + gk0_embed.w * cos_vec.y - gk0_embed.z * sin_vec.y); + const float4 gk1 = make_float4( + gk1_embed.x * cos_vec.z + gk1_embed.y * sin_vec.z, + gk1_embed.y * cos_vec.z - gk1_embed.x * sin_vec.z, + gk1_embed.z * cos_vec.w + gk1_embed.w * sin_vec.w, + gk1_embed.w * cos_vec.w - gk1_embed.z * sin_vec.w); + from_float4(&grad_k[gk_base + 2 * base_d], gk0); + from_float4(&grad_k[gk_base + 2 * base_d + 4], gk1); + } else { + const float4 gk0_rot = to_float4(&grad_k_embed[gke_base + base_d]); + const float4 gk1_rot = to_float4(&grad_k_embed[gke_base + base_d + half_dim]); + const float4 gk0 = make_float4( + gk0_rot.x * cos_vec.x + gk1_rot.x * sin_vec.x, + gk0_rot.y * cos_vec.y + gk1_rot.y * sin_vec.y, + gk0_rot.z * cos_vec.z + gk1_rot.z * sin_vec.z, + gk0_rot.w * cos_vec.w + gk1_rot.w * sin_vec.w); + const float4 gk1 = make_float4( + gk1_rot.x * cos_vec.x - gk0_rot.x * sin_vec.x, + gk1_rot.y * cos_vec.y - gk0_rot.y * sin_vec.y, + gk1_rot.z * cos_vec.z - gk0_rot.z * sin_vec.z, + gk1_rot.w * cos_vec.w - gk0_rot.w * sin_vec.w); + from_float4(&grad_k[gk_base + base_d], gk0); + from_float4(&grad_k[gk_base + base_d + half_dim], gk1); + } + } + + const int pass_dim = full_d - rotary_d; + if (pass_dim <= 0) { + return; + } + const int pass_vecs = (pass_dim + 3) / 4; + for (int linear = threadIdx.x; linear < q_h * pass_vecs; linear += blockDim.x) { + const int h_idx = linear / pass_vecs; + const int rel_vec = linear - h_idx * pass_vecs; + const int rel_d = rel_vec * 4; + const int64_t gqe_base = + static_cast(b) * qb_stride + + static_cast(s) * qs_stride + + static_cast(h_idx) * qh_stride + rotary_d + rel_d; + const int64_t gq_base = + static_cast(b) * qgb_stride + + static_cast(s) * qgs_stride + + static_cast(h_idx) * qgh_stride + rotary_d + rel_d; + if (rel_d + 3 < pass_dim) { + copy4(&grad_q[gq_base], &grad_q_embed[gqe_base]); + } else { + for (int t = rel_d; t < pass_dim; ++t) { + grad_q[gq_base + (t - rel_d)] = grad_q_embed[gqe_base + (t - rel_d)]; + } + } + } + for (int linear = threadIdx.x; linear < k_h * pass_vecs; linear += blockDim.x) { + const int h_idx = linear / pass_vecs; + const int rel_vec = linear - h_idx * pass_vecs; + const int rel_d = rel_vec * 4; + const int64_t gke_base = + static_cast(b) * kb_stride + + static_cast(s) * ks_stride + + static_cast(h_idx) * kh_stride + rotary_d + rel_d; + const int64_t gk_base = + static_cast(b) * kgb_stride + + static_cast(s) * kgs_stride + + static_cast(h_idx) * kgh_stride + rotary_d + rel_d; + if (rel_d + 3 < pass_dim) { + copy4(&grad_k[gk_base], &grad_k_embed[gke_base]); + } else { + for (int t = rel_d; t < pass_dim; ++t) { + grad_k[gk_base + (t - rel_d)] = grad_k_embed[gke_base + (t - rel_d)]; + } + } + } +} + template __global__ void RopeInplaceKernel(float* cos, float* sin, T* q, const int q_h, T* k, const int k_h, const int d, + const int64_t seq_len, const int qb_stride, const int qs_stride, const int qh_stride, const int kb_stride, const int ks_stride, const int kh_stride) { extern __shared__ char cos_sin[]; @@ -221,9 +661,10 @@ __global__ void RopeInplaceKernel(float* cos, float* sin, float* cos_smem = reinterpret_cast(cos_sin); float* sin_smem = cos_smem + half_dim; int b = blockIdx.x; - int s = blockIdx.y; + int64_t s = RopeSeqBlockIdx(); + if (s >= seq_len) return; - int64_t cos_sin_b_stride = gridDim.y * half_dim; + int64_t cos_sin_b_stride = seq_len * half_dim; int64_t cos_sin_s_stride = half_dim; #define SIN_GMEM(b, c, d) sin[(b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] #define COS_GMEM(b, c, d) cos[(b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] @@ -347,6 +788,7 @@ __global__ void RopeKernelBackward( const int q_h, const int k_h, const int d, + const int64_t seq_len, const int qb_stride, const int qs_stride, const int qh_stride, const int kb_stride, const int ks_stride, const int kh_stride ) { @@ -356,9 +798,10 @@ __global__ void RopeKernelBackward( float* sin_smem = cos_smem + half_dim; int b = blockIdx.x; - int s = blockIdx.y; + int64_t s = RopeSeqBlockIdx(); + if (s >= seq_len) return; - int64_t cos_sin_b_stride = gridDim.y * half_dim; + int64_t cos_sin_b_stride = seq_len * half_dim; int64_t cos_sin_s_stride = half_dim; #define SIN_GMEM(b, c, d) sin[(b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] #define COS_GMEM(b, c, d) cos[(b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)] @@ -548,7 +991,7 @@ void Rope(const at::Tensor& q, // [b, s, h, d] or [s, h, d] keb_stride = k_embed.stride(0); kes_stride = k_embed.stride(1); keh_stride = k_embed.stride(2); } - dim3 grid(batch, seq_len); + dim3 grid = MakeRopeGrid(batch, seq_len); if (q.scalar_type() == at::kHalf) { const half* q_data = static_cast(q.data_ptr()); @@ -563,7 +1006,7 @@ void Rope(const at::Tensor& q, // [b, s, h, d] or [s, h, d] q_data, q_head_num, k_data, k_head_num, q_embed_data, k_embed_data, - dim, + dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride); } else { @@ -572,7 +1015,7 @@ void Rope(const at::Tensor& q, // [b, s, h, d] or [s, h, d] q_data, q_head_num, k_data, k_head_num, q_embed_data, k_embed_data, - dim, + dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride); } @@ -589,7 +1032,7 @@ void Rope(const at::Tensor& q, // [b, s, h, d] or [s, h, d] q_data, q_head_num, k_data, k_head_num, q_embed_data, k_embed_data, - dim, + dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride); } else { @@ -598,7 +1041,7 @@ void Rope(const at::Tensor& q, // [b, s, h, d] or [s, h, d] q_data, q_head_num, k_data, k_head_num, q_embed_data, k_embed_data, - dim, + dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride); } @@ -615,7 +1058,7 @@ void Rope(const at::Tensor& q, // [b, s, h, d] or [s, h, d] q_data, q_head_num, k_data, k_head_num, q_embed_data, k_embed_data, - dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride ); } else { @@ -624,7 +1067,7 @@ void Rope(const at::Tensor& q, // [b, s, h, d] or [s, h, d] q_data, q_head_num, k_data, k_head_num, q_embed_data, k_embed_data, - dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride ); } @@ -635,6 +1078,291 @@ void Rope(const at::Tensor& q, // [b, s, h, d] or [s, h, d] sync_check_cuda_error(); } +struct CosSinLayout { + int64_t b_stride; + int64_t s_stride; + int64_t d_stride; +}; + +CosSinLayout GetCosSinLayout( + const at::Tensor& tensor, + bool q_3d, + int64_t batch, + int64_t seq_len, + int64_t half_dim +) { + ASSERT_CHECK(tensor.size(tensor.dim() - 1) >= half_dim); + + CosSinLayout layout; + layout.d_stride = tensor.stride(tensor.dim() - 1); + + if (q_3d) { + if (tensor.dim() == 2) { + ASSERT_CHECK(tensor.size(0) == seq_len); + layout.b_stride = 0; + layout.s_stride = tensor.stride(0); + } else if (tensor.dim() == 3) { + ASSERT_CHECK(tensor.size(0) == 1); + ASSERT_CHECK(tensor.size(1) == seq_len); + layout.b_stride = 0; + layout.s_stride = tensor.stride(1); + } else { + ASSERT_CHECK(false); + } + } else { + ASSERT_CHECK(tensor.dim() == 3); + ASSERT_CHECK(tensor.size(0) == batch || tensor.size(0) == 1); + ASSERT_CHECK(tensor.size(1) == seq_len); + layout.b_stride = (tensor.size(0) == 1) ? 0 : tensor.stride(0); + layout.s_stride = tensor.stride(1); + } + + return layout; +} + +template +void LaunchRopeFullTyped( + const C* cos_data, + const C* sin_data, + const T* q_data, + const int64_t q_head_num, + const T* k_data, + const int64_t k_head_num, + T* q_embed_data, + T* k_embed_data, + const int64_t full_dim, + const int64_t rotary_dim, + const int64_t qb_stride, + const int64_t qs_stride, + const int64_t qh_stride, + const int64_t kb_stride, + const int64_t ks_stride, + const int64_t kh_stride, + const int64_t qeb_stride, + const int64_t qes_stride, + const int64_t qeh_stride, + const int64_t keb_stride, + const int64_t kes_stride, + const int64_t keh_stride, + const CosSinLayout cos_layout, + const CosSinLayout sin_layout, + const bool interleave, + const dim3 grid, + const int nthreads, + const size_t smem, + cudaStream_t stream +) { + if (interleave) { + RopeFullKernel<<>>( + cos_data, sin_data, q_data, q_head_num, k_data, k_head_num, + q_embed_data, k_embed_data, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride, + cos_layout.b_stride, cos_layout.s_stride, cos_layout.d_stride, + sin_layout.b_stride, sin_layout.s_stride, sin_layout.d_stride); + } else { + RopeFullKernel<<>>( + cos_data, sin_data, q_data, q_head_num, k_data, k_head_num, + q_embed_data, k_embed_data, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride, + cos_layout.b_stride, cos_layout.s_stride, cos_layout.d_stride, + sin_layout.b_stride, sin_layout.s_stride, sin_layout.d_stride); + } +} + +template +void DispatchRopeFullByCosDtype( + const at::Tensor& cos, + const at::Tensor& sin, + const T* q_data, + const int64_t q_head_num, + const T* k_data, + const int64_t k_head_num, + T* q_embed_data, + T* k_embed_data, + const int64_t full_dim, + const int64_t rotary_dim, + const int64_t qb_stride, + const int64_t qs_stride, + const int64_t qh_stride, + const int64_t kb_stride, + const int64_t ks_stride, + const int64_t kh_stride, + const int64_t qeb_stride, + const int64_t qes_stride, + const int64_t qeh_stride, + const int64_t keb_stride, + const int64_t kes_stride, + const int64_t keh_stride, + const CosSinLayout cos_layout, + const CosSinLayout sin_layout, + const bool interleave, + const dim3 grid, + const int nthreads, + const size_t smem, + cudaStream_t stream +) { + if (cos.scalar_type() == at::kFloat) { + LaunchRopeFullTyped( + static_cast(cos.data_ptr()), + static_cast(sin.data_ptr()), + q_data, q_head_num, k_data, k_head_num, q_embed_data, k_embed_data, + full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride, + cos_layout, sin_layout, interleave, grid, nthreads, smem, stream); + } else if (cos.scalar_type() == at::kHalf) { + LaunchRopeFullTyped( + static_cast(cos.data_ptr()), + static_cast(sin.data_ptr()), + q_data, q_head_num, k_data, k_head_num, q_embed_data, k_embed_data, + full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride, + cos_layout, sin_layout, interleave, grid, nthreads, smem, stream); + } else if (cos.scalar_type() == at::kBFloat16) { + LaunchRopeFullTyped( + static_cast(cos.data_ptr()), + static_cast(sin.data_ptr()), + q_data, q_head_num, k_data, k_head_num, q_embed_data, k_embed_data, + full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride, + cos_layout, sin_layout, interleave, grid, nthreads, smem, stream); + } else { + throw std::runtime_error("Unsupported cos/sin data type for rope_full"); + } +} + +void RopeFull(const at::Tensor& q, // [b, s, h, full_d] or [s, h, full_d] + const at::Tensor& k, // [b, s, h_k, full_d] or [s, h_k, full_d] + const at::Tensor& q_embed, // same shape as q + const at::Tensor& k_embed, // same shape as k + const at::Tensor& cos, // compact [*, rotary_d / 2] or duplicated [*, rotary_d] + const at::Tensor& sin, // compact [*, rotary_d / 2] or duplicated [*, rotary_d] + bool interleave, + int64_t rotary_dim + ) { + int Nthreads = 256; + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + bool q_3d = (q.dim() == 3); + + ASSERT_CHECK(q.scalar_type() == k.scalar_type()); + ASSERT_CHECK(q.scalar_type() == q_embed.scalar_type()); + ASSERT_CHECK(k.scalar_type() == k_embed.scalar_type()); + ASSERT_CHECK(cos.scalar_type() == sin.scalar_type()); + ASSERT_CHECK(cos.dim() == sin.dim()); + ASSERT_CHECK(cos.size(cos.dim() - 1) == sin.size(sin.dim() - 1)); + + ASSERT_CHECK(rotary_dim > 0); + ASSERT_CHECK(rotary_dim % 8 == 0); + const int64_t half_dim = rotary_dim / 2; + ASSERT_CHECK(cos.size(cos.dim() - 1) >= half_dim); + + int64_t batch, q_head_num, k_head_num, seq_len, full_dim; + int64_t qb_stride, kb_stride, qs_stride, ks_stride, qh_stride, kh_stride; + int64_t qeb_stride, keb_stride, qes_stride, kes_stride, qeh_stride, keh_stride; + + if (q_3d) { + ASSERT_CHECK(q_embed.size(0) == q.size(0)); + ASSERT_CHECK(q_embed.size(1) == q.size(1)); + ASSERT_CHECK(q_embed.size(2) == q.size(2)); + ASSERT_CHECK(k_embed.size(0) == k.size(0)); + ASSERT_CHECK(k_embed.size(1) == k.size(1)); + ASSERT_CHECK(k_embed.size(2) == k.size(2)); + ASSERT_CHECK(k.size(2) == q.size(2)); + ASSERT_CHECK(rotary_dim <= q.size(2)); + ASSERT_CHECK(rotary_dim <= k.size(2)); + ASSERT_CHECK(q.stride(2) == 1); + ASSERT_CHECK(k.stride(2) == 1); + ASSERT_CHECK(q_embed.stride(2) == 1); + ASSERT_CHECK(k_embed.stride(2) == 1); + batch = 1; + seq_len = q.size(0); + q_head_num = q.size(1); + k_head_num = k.size(1); + full_dim = q.size(2); + qb_stride = 0; qs_stride = q.stride(0); qh_stride = q.stride(1); + kb_stride = 0; ks_stride = k.stride(0); kh_stride = k.stride(1); + qeb_stride = 0; qes_stride = q_embed.stride(0); qeh_stride = q_embed.stride(1); + keb_stride = 0; kes_stride = k_embed.stride(0); keh_stride = k_embed.stride(1); + } else { + ASSERT_CHECK(q_embed.size(0) == q.size(0)); + ASSERT_CHECK(q_embed.size(1) == q.size(1)); + ASSERT_CHECK(q_embed.size(2) == q.size(2)); + ASSERT_CHECK(q_embed.size(3) == q.size(3)); + ASSERT_CHECK(k_embed.size(0) == k.size(0)); + ASSERT_CHECK(k_embed.size(1) == k.size(1)); + ASSERT_CHECK(k_embed.size(2) == k.size(2)); + ASSERT_CHECK(k_embed.size(3) == k.size(3)); + ASSERT_CHECK(k.size(3) == q.size(3)); + ASSERT_CHECK(rotary_dim <= q.size(3)); + ASSERT_CHECK(rotary_dim <= k.size(3)); + ASSERT_CHECK(q.stride(3) == 1); + ASSERT_CHECK(k.stride(3) == 1); + ASSERT_CHECK(q_embed.stride(3) == 1); + ASSERT_CHECK(k_embed.stride(3) == 1); + batch = q.size(0); + seq_len = q.size(1); + q_head_num = q.size(2); + k_head_num = k.size(2); + full_dim = q.size(3); + qb_stride = q.stride(0); qs_stride = q.stride(1); qh_stride = q.stride(2); + kb_stride = k.stride(0); ks_stride = k.stride(1); kh_stride = k.stride(2); + qeb_stride = q_embed.stride(0); qes_stride = q_embed.stride(1); qeh_stride = q_embed.stride(2); + keb_stride = k_embed.stride(0); kes_stride = k_embed.stride(1); keh_stride = k_embed.stride(2); + } + + const CosSinLayout cos_layout = + GetCosSinLayout(cos, q_3d, batch, seq_len, half_dim); + const CosSinLayout sin_layout = + GetCosSinLayout(sin, q_3d, batch, seq_len, half_dim); + dim3 grid(batch, seq_len); + const size_t smem = static_cast(rotary_dim) * sizeof(float); + + if (q.scalar_type() == at::kHalf) { + const half* q_data = static_cast(q.data_ptr()); + const half* k_data = static_cast(k.data_ptr()); + half* q_embed_data = static_cast(q_embed.data_ptr()); + half* k_embed_data = static_cast(k_embed.data_ptr()); + DispatchRopeFullByCosDtype( + cos, sin, q_data, q_head_num, k_data, k_head_num, + q_embed_data, k_embed_data, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride, + cos_layout, sin_layout, interleave, grid, Nthreads, smem, stream); + } else if (q.scalar_type() == at::kBFloat16) { + const __nv_bfloat16* q_data = static_cast(q.data_ptr()); + const __nv_bfloat16* k_data = static_cast(k.data_ptr()); + __nv_bfloat16* q_embed_data = static_cast<__nv_bfloat16*>(q_embed.data_ptr()); + __nv_bfloat16* k_embed_data = static_cast<__nv_bfloat16*>(k_embed.data_ptr()); + DispatchRopeFullByCosDtype<__nv_bfloat16>( + cos, sin, q_data, q_head_num, k_data, k_head_num, + q_embed_data, k_embed_data, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride, + cos_layout, sin_layout, interleave, grid, Nthreads, smem, stream); + } else if (q.scalar_type() == at::kFloat) { + const float* q_data = static_cast(q.data_ptr()); + const float* k_data = static_cast(k.data_ptr()); + float* q_embed_data = static_cast(q_embed.data_ptr()); + float* k_embed_data = static_cast(k_embed.data_ptr()); + DispatchRopeFullByCosDtype( + cos, sin, q_data, q_head_num, k_data, k_head_num, + q_embed_data, k_embed_data, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride, + cos_layout, sin_layout, interleave, grid, Nthreads, smem, stream); + } else { + throw std::runtime_error("Unsupported data type for rope_full"); + } + + sync_check_cuda_error(); +} + void RopeInplace(const at::Tensor& q, // [b, s, h, d] or [s, h, d] const at::Tensor& k, // [b, s, h_k, d] or [s, h_k, d] const at::Tensor& cos, // [b, s, d / 2] or [s, d / 2] @@ -682,7 +1410,7 @@ void RopeInplace(const at::Tensor& q, // [b, s, h, d] or [s, h, d] kb_stride = k.stride(0); ks_stride = k.stride(1); kh_stride = k.stride(2); } - dim3 grid(batch, seq_len); + dim3 grid = MakeRopeGrid(batch, seq_len); if (q.scalar_type() == at::kHalf) { half* q_data = static_cast(q.data_ptr()); @@ -694,14 +1422,14 @@ void RopeInplace(const at::Tensor& q, // [b, s, h, d] or [s, h, d] cos_data, sin_data, q_data, q_head_num, k_data, k_head_num, - dim, + dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); } else { RopeInplaceKernel<<>>( cos_data, sin_data, q_data, q_head_num, k_data, k_head_num, - dim, + dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); } } else if (q.scalar_type() == at::kBFloat16) { @@ -714,14 +1442,14 @@ void RopeInplace(const at::Tensor& q, // [b, s, h, d] or [s, h, d] cos_data, sin_data, q_data, q_head_num, k_data, k_head_num, - dim, + dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); } else { RopeInplaceKernel<__nv_bfloat16, false><<>>( cos_data, sin_data, q_data, q_head_num, k_data, k_head_num, - dim, + dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride); } } else if (q.scalar_type() == at::kFloat) { @@ -734,14 +1462,14 @@ void RopeInplace(const at::Tensor& q, // [b, s, h, d] or [s, h, d] cos_data, sin_data, q_data, q_head_num, k_data, k_head_num, - dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride ); } else { RopeInplaceKernel<<>>( cos_data, sin_data, q_data, q_head_num, k_data, k_head_num, - dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride + dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride ); } } else { @@ -814,7 +1542,7 @@ void RopeBackward( kb_stride = grad_k_embed.stride(0); ks_stride = grad_k_embed.stride(1); kh_stride = grad_k_embed.stride(2); } - dim3 grid(batch, seq_len); + dim3 grid = MakeRopeGrid(batch, seq_len); if (grad_q_embed.scalar_type() == at::kHalf) { const half* gq_embed = static_cast(grad_q_embed.data_ptr()); @@ -828,7 +1556,7 @@ void RopeBackward( cos_data, sin_data, gq_embed, gk_embed, gq, gk, - q_head_num, k_head_num, dim, + q_head_num, k_head_num, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride ); } else { @@ -836,7 +1564,7 @@ void RopeBackward( cos_data, sin_data, gq_embed, gk_embed, gq, gk, - q_head_num, k_head_num, dim, + q_head_num, k_head_num, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride ); } @@ -852,7 +1580,7 @@ void RopeBackward( cos_data, sin_data, gq_embed, gk_embed, gq, gk, - q_head_num, k_head_num, dim, + q_head_num, k_head_num, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride ); } else { @@ -860,7 +1588,7 @@ void RopeBackward( cos_data, sin_data, gq_embed, gk_embed, gq, gk, - q_head_num, k_head_num, dim, + q_head_num, k_head_num, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride ); } @@ -876,7 +1604,7 @@ void RopeBackward( cos_data, sin_data, gq_embed, gk_embed, gq, gk, - q_head_num, k_head_num, dim, + q_head_num, k_head_num, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride ); } else { @@ -884,7 +1612,7 @@ void RopeBackward( cos_data, sin_data, gq_embed, gk_embed, gq, gk, - q_head_num, k_head_num, dim, + q_head_num, k_head_num, dim, seq_len, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride ); } @@ -895,6 +1623,248 @@ void RopeBackward( sync_check_cuda_error(); } +template +void LaunchRopeFullBackwardTyped( + const C* cos_data, + const C* sin_data, + const T* gq_embed, + const T* gk_embed, + T* gq, + T* gk, + const int64_t q_head_num, + const int64_t k_head_num, + const int64_t full_dim, + const int64_t rotary_dim, + const int64_t qb_stride, + const int64_t qs_stride, + const int64_t qh_stride, + const int64_t kb_stride, + const int64_t ks_stride, + const int64_t kh_stride, + const int64_t qgb_stride, + const int64_t qgs_stride, + const int64_t qgh_stride, + const int64_t kgb_stride, + const int64_t kgs_stride, + const int64_t kgh_stride, + const CosSinLayout cos_layout, + const CosSinLayout sin_layout, + const bool interleave, + const dim3 grid, + const int nthreads, + const size_t smem, + cudaStream_t stream +) { + if (interleave) { + RopeFullKernelBackward<<>>( + cos_data, sin_data, gq_embed, gk_embed, gq, gk, + q_head_num, k_head_num, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qgb_stride, qgs_stride, qgh_stride, kgb_stride, kgs_stride, kgh_stride, + cos_layout.b_stride, cos_layout.s_stride, cos_layout.d_stride, + sin_layout.b_stride, sin_layout.s_stride, sin_layout.d_stride); + } else { + RopeFullKernelBackward<<>>( + cos_data, sin_data, gq_embed, gk_embed, gq, gk, + q_head_num, k_head_num, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qgb_stride, qgs_stride, qgh_stride, kgb_stride, kgs_stride, kgh_stride, + cos_layout.b_stride, cos_layout.s_stride, cos_layout.d_stride, + sin_layout.b_stride, sin_layout.s_stride, sin_layout.d_stride); + } +} + +template +void DispatchRopeFullBackwardByCosDtype( + const at::Tensor& cos, + const at::Tensor& sin, + const T* gq_embed, + const T* gk_embed, + T* gq, + T* gk, + const int64_t q_head_num, + const int64_t k_head_num, + const int64_t full_dim, + const int64_t rotary_dim, + const int64_t qb_stride, + const int64_t qs_stride, + const int64_t qh_stride, + const int64_t kb_stride, + const int64_t ks_stride, + const int64_t kh_stride, + const int64_t qgb_stride, + const int64_t qgs_stride, + const int64_t qgh_stride, + const int64_t kgb_stride, + const int64_t kgs_stride, + const int64_t kgh_stride, + const CosSinLayout cos_layout, + const CosSinLayout sin_layout, + const bool interleave, + const dim3 grid, + const int nthreads, + const size_t smem, + cudaStream_t stream +) { + if (cos.scalar_type() == at::kFloat) { + LaunchRopeFullBackwardTyped( + static_cast(cos.data_ptr()), + static_cast(sin.data_ptr()), + gq_embed, gk_embed, gq, gk, + q_head_num, k_head_num, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qgb_stride, qgs_stride, qgh_stride, kgb_stride, kgs_stride, kgh_stride, + cos_layout, sin_layout, interleave, grid, nthreads, smem, stream); + } else if (cos.scalar_type() == at::kHalf) { + LaunchRopeFullBackwardTyped( + static_cast(cos.data_ptr()), + static_cast(sin.data_ptr()), + gq_embed, gk_embed, gq, gk, + q_head_num, k_head_num, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qgb_stride, qgs_stride, qgh_stride, kgb_stride, kgs_stride, kgh_stride, + cos_layout, sin_layout, interleave, grid, nthreads, smem, stream); + } else if (cos.scalar_type() == at::kBFloat16) { + LaunchRopeFullBackwardTyped( + static_cast(cos.data_ptr()), + static_cast(sin.data_ptr()), + gq_embed, gk_embed, gq, gk, + q_head_num, k_head_num, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qgb_stride, qgs_stride, qgh_stride, kgb_stride, kgs_stride, kgh_stride, + cos_layout, sin_layout, interleave, grid, nthreads, smem, stream); + } else { + throw std::runtime_error("Unsupported cos/sin data type for rope_full backward"); + } +} + +void RopeFullBackward( + const at::Tensor& grad_q_embed, + const at::Tensor& grad_k_embed, + const at::Tensor& grad_q, + const at::Tensor& grad_k, + const at::Tensor& cos, + const at::Tensor& sin, + bool interleave, + int64_t rotary_dim +) { + int Nthreads = 256; + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + bool q_3d = (grad_q_embed.dim() == 3); + + ASSERT_CHECK(grad_q_embed.scalar_type() == grad_k_embed.scalar_type()); + ASSERT_CHECK(grad_q.scalar_type() == grad_k.scalar_type()); + ASSERT_CHECK(grad_q.scalar_type() == grad_q_embed.scalar_type()); + ASSERT_CHECK(cos.scalar_type() == sin.scalar_type()); + ASSERT_CHECK(cos.dim() == sin.dim()); + ASSERT_CHECK(cos.size(cos.dim() - 1) == sin.size(sin.dim() - 1)); + + ASSERT_CHECK(rotary_dim > 0); + ASSERT_CHECK(rotary_dim % 8 == 0); + const int64_t half_dim = rotary_dim / 2; + ASSERT_CHECK(cos.size(cos.dim() - 1) >= half_dim); + + int64_t batch, seq_len, q_head_num, k_head_num, full_dim; + int64_t qb_stride, kb_stride, qs_stride, ks_stride, qh_stride, kh_stride; + int64_t qgb_stride, kgb_stride, qgs_stride, kgs_stride, qgh_stride, kgh_stride; + + if (q_3d) { + ASSERT_CHECK(grad_q_embed.size(0) == grad_q.size(0)); + ASSERT_CHECK(grad_q_embed.size(1) == grad_q.size(1)); + ASSERT_CHECK(grad_q_embed.size(2) == grad_q.size(2)); + ASSERT_CHECK(grad_k_embed.size(0) == grad_k.size(0)); + ASSERT_CHECK(grad_k_embed.size(1) == grad_k.size(1)); + ASSERT_CHECK(grad_k_embed.size(2) == grad_k.size(2)); + ASSERT_CHECK(grad_k_embed.size(2) == grad_q_embed.size(2)); + ASSERT_CHECK(rotary_dim <= grad_q_embed.size(2)); + ASSERT_CHECK(grad_q_embed.stride(2) == 1); + ASSERT_CHECK(grad_k_embed.stride(2) == 1); + ASSERT_CHECK(grad_q.stride(2) == 1); + ASSERT_CHECK(grad_k.stride(2) == 1); + batch = 1; + seq_len = grad_q_embed.size(0); + q_head_num = grad_q_embed.size(1); + k_head_num = grad_k_embed.size(1); + full_dim = grad_q_embed.size(2); + qb_stride = 0; qs_stride = grad_q_embed.stride(0); qh_stride = grad_q_embed.stride(1); + kb_stride = 0; ks_stride = grad_k_embed.stride(0); kh_stride = grad_k_embed.stride(1); + qgb_stride = 0; qgs_stride = grad_q.stride(0); qgh_stride = grad_q.stride(1); + kgb_stride = 0; kgs_stride = grad_k.stride(0); kgh_stride = grad_k.stride(1); + } else { + ASSERT_CHECK(grad_q_embed.size(0) == grad_q.size(0)); + ASSERT_CHECK(grad_q_embed.size(1) == grad_q.size(1)); + ASSERT_CHECK(grad_q_embed.size(2) == grad_q.size(2)); + ASSERT_CHECK(grad_q_embed.size(3) == grad_q.size(3)); + ASSERT_CHECK(grad_k_embed.size(0) == grad_k.size(0)); + ASSERT_CHECK(grad_k_embed.size(1) == grad_k.size(1)); + ASSERT_CHECK(grad_k_embed.size(2) == grad_k.size(2)); + ASSERT_CHECK(grad_k_embed.size(3) == grad_k.size(3)); + ASSERT_CHECK(grad_k_embed.size(3) == grad_q_embed.size(3)); + ASSERT_CHECK(rotary_dim <= grad_q_embed.size(3)); + ASSERT_CHECK(grad_q_embed.stride(3) == 1); + ASSERT_CHECK(grad_k_embed.stride(3) == 1); + ASSERT_CHECK(grad_q.stride(3) == 1); + ASSERT_CHECK(grad_k.stride(3) == 1); + batch = grad_q_embed.size(0); + seq_len = grad_q_embed.size(1); + q_head_num = grad_q_embed.size(2); + k_head_num = grad_k_embed.size(2); + full_dim = grad_q_embed.size(3); + qb_stride = grad_q_embed.stride(0); qs_stride = grad_q_embed.stride(1); qh_stride = grad_q_embed.stride(2); + kb_stride = grad_k_embed.stride(0); ks_stride = grad_k_embed.stride(1); kh_stride = grad_k_embed.stride(2); + qgb_stride = grad_q.stride(0); qgs_stride = grad_q.stride(1); qgh_stride = grad_q.stride(2); + kgb_stride = grad_k.stride(0); kgs_stride = grad_k.stride(1); kgh_stride = grad_k.stride(2); + } + + const CosSinLayout cos_layout = + GetCosSinLayout(cos, q_3d, batch, seq_len, half_dim); + const CosSinLayout sin_layout = + GetCosSinLayout(sin, q_3d, batch, seq_len, half_dim); + dim3 grid(batch, seq_len); + const size_t smem = static_cast(rotary_dim) * sizeof(float); + + if (grad_q_embed.scalar_type() == at::kHalf) { + const half* gq_embed = static_cast(grad_q_embed.data_ptr()); + const half* gk_embed = static_cast(grad_k_embed.data_ptr()); + half* gq = static_cast(grad_q.data_ptr()); + half* gk = static_cast(grad_k.data_ptr()); + DispatchRopeFullBackwardByCosDtype( + cos, sin, gq_embed, gk_embed, gq, gk, + q_head_num, k_head_num, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qgb_stride, qgs_stride, qgh_stride, kgb_stride, kgs_stride, kgh_stride, + cos_layout, sin_layout, interleave, grid, Nthreads, smem, stream); + } else if (grad_q_embed.scalar_type() == at::kBFloat16) { + const __nv_bfloat16* gq_embed = static_cast(grad_q_embed.data_ptr()); + const __nv_bfloat16* gk_embed = static_cast(grad_k_embed.data_ptr()); + __nv_bfloat16* gq = static_cast<__nv_bfloat16*>(grad_q.data_ptr()); + __nv_bfloat16* gk = static_cast<__nv_bfloat16*>(grad_k.data_ptr()); + DispatchRopeFullBackwardByCosDtype<__nv_bfloat16>( + cos, sin, gq_embed, gk_embed, gq, gk, + q_head_num, k_head_num, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qgb_stride, qgs_stride, qgh_stride, kgb_stride, kgs_stride, kgh_stride, + cos_layout, sin_layout, interleave, grid, Nthreads, smem, stream); + } else if (grad_q_embed.scalar_type() == at::kFloat) { + const float* gq_embed = static_cast(grad_q_embed.data_ptr()); + const float* gk_embed = static_cast(grad_k_embed.data_ptr()); + float* gq = static_cast(grad_q.data_ptr()); + float* gk = static_cast(grad_k.data_ptr()); + DispatchRopeFullBackwardByCosDtype( + cos, sin, gq_embed, gk_embed, gq, gk, + q_head_num, k_head_num, full_dim, rotary_dim, + qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride, + qgb_stride, qgs_stride, qgh_stride, kgb_stride, kgs_stride, kgh_stride, + cos_layout, sin_layout, interleave, grid, Nthreads, smem, stream); + } else { + throw std::runtime_error("Unsupported data type for rope_full backward"); + } + + sync_check_cuda_error(); +} + // Pack interface: accepts qkv [seq_len, q_dim + 2*kv_dim], applies rope inplace // on the q and k slices. Supports GQA (q_num_heads != kv_num_heads). // CUDA kernels are not modified; only host-side pointer arithmetic is added. @@ -930,7 +1900,7 @@ void RopeInplacePack(const at::Tensor& qkv, // [seq_len, q_dim + 2*kv_dim] int64_t ks_stride = qkv_row_stride; int64_t kh_stride = head_dim; - dim3 grid(1, seq_len); + dim3 grid = MakeRopeGrid(1, seq_len); if (qkv.scalar_type() == at::kHalf) { half* q_data = static_cast(qkv.data_ptr()); @@ -940,12 +1910,12 @@ void RopeInplacePack(const at::Tensor& qkv, // [seq_len, q_dim + 2*kv_dim] if (interleave) { RopeInplaceKernel<<>>( cos_data, sin_data, - q_data, q_num_heads, k_data, kv_num_heads, head_dim, + q_data, q_num_heads, k_data, kv_num_heads, head_dim, seq_len, 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); } else { RopeInplaceKernel<<>>( cos_data, sin_data, - q_data, q_num_heads, k_data, kv_num_heads, head_dim, + q_data, q_num_heads, k_data, kv_num_heads, head_dim, seq_len, 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); } } else if (qkv.scalar_type() == at::kBFloat16) { @@ -956,12 +1926,12 @@ void RopeInplacePack(const at::Tensor& qkv, // [seq_len, q_dim + 2*kv_dim] if (interleave) { RopeInplaceKernel<__nv_bfloat16, true><<>>( cos_data, sin_data, - q_data, q_num_heads, k_data, kv_num_heads, head_dim, + q_data, q_num_heads, k_data, kv_num_heads, head_dim, seq_len, 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); } else { RopeInplaceKernel<__nv_bfloat16, false><<>>( cos_data, sin_data, - q_data, q_num_heads, k_data, kv_num_heads, head_dim, + q_data, q_num_heads, k_data, kv_num_heads, head_dim, seq_len, 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); } } else if (qkv.scalar_type() == at::kFloat) { @@ -972,12 +1942,12 @@ void RopeInplacePack(const at::Tensor& qkv, // [seq_len, q_dim + 2*kv_dim] if (interleave) { RopeInplaceKernel<<>>( cos_data, sin_data, - q_data, q_num_heads, k_data, kv_num_heads, head_dim, + q_data, q_num_heads, k_data, kv_num_heads, head_dim, seq_len, 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); } else { RopeInplaceKernel<<>>( cos_data, sin_data, - q_data, q_num_heads, k_data, kv_num_heads, head_dim, + q_data, q_num_heads, k_data, kv_num_heads, head_dim, seq_len, 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); } } else { @@ -1020,7 +1990,7 @@ void RopeInplacePackBackward(const at::Tensor& dqkv, // [seq_len, q_dim + 2*kv_d int64_t ks_stride = dqkv_row_stride; int64_t kh_stride = head_dim; - dim3 grid(1, seq_len); + dim3 grid = MakeRopeGrid(1, seq_len); if (dqkv.scalar_type() == at::kHalf) { half* dq_data = static_cast(dqkv.data_ptr()); @@ -1032,14 +2002,14 @@ void RopeInplacePackBackward(const at::Tensor& dqkv, // [seq_len, q_dim + 2*kv_d cos_data, sin_data, dq_data, dk_data, dq_data, dk_data, - q_num_heads, kv_num_heads, head_dim, + q_num_heads, kv_num_heads, head_dim, seq_len, 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); } else { RopeKernelBackward<<>>( cos_data, sin_data, dq_data, dk_data, dq_data, dk_data, - q_num_heads, kv_num_heads, head_dim, + q_num_heads, kv_num_heads, head_dim, seq_len, 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); } } else if (dqkv.scalar_type() == at::kBFloat16) { @@ -1052,14 +2022,14 @@ void RopeInplacePackBackward(const at::Tensor& dqkv, // [seq_len, q_dim + 2*kv_d cos_data, sin_data, dq_data, dk_data, dq_data, dk_data, - q_num_heads, kv_num_heads, head_dim, + q_num_heads, kv_num_heads, head_dim, seq_len, 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); } else { RopeKernelBackward<__nv_bfloat16, false><<>>( cos_data, sin_data, dq_data, dk_data, dq_data, dk_data, - q_num_heads, kv_num_heads, head_dim, + q_num_heads, kv_num_heads, head_dim, seq_len, 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); } } else if (dqkv.scalar_type() == at::kFloat) { @@ -1072,14 +2042,14 @@ void RopeInplacePackBackward(const at::Tensor& dqkv, // [seq_len, q_dim + 2*kv_d cos_data, sin_data, dq_data, dk_data, dq_data, dk_data, - q_num_heads, kv_num_heads, head_dim, + q_num_heads, kv_num_heads, head_dim, seq_len, 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); } else { RopeKernelBackward<<>>( cos_data, sin_data, dq_data, dk_data, dq_data, dk_data, - q_num_heads, kv_num_heads, head_dim, + q_num_heads, kv_num_heads, head_dim, seq_len, 0, qs_stride, qh_stride, 0, ks_stride, kh_stride); } } else { @@ -1094,5 +2064,7 @@ void RopeInplacePackBackward(const at::Tensor& dqkv, // [seq_len, q_dim + 2*kv_d // Removed TVM FFI export (see pybind11 registration) // Removed TVM FFI export (see pybind11 registration) // Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) } // wallx_cuda_rope diff --git a/wall_x/model/core/ops/csrc/silu_and_mul/silu_and_mul.cu b/wall_x/model/core/ops/csrc/silu_and_mul/silu_and_mul.cu new file mode 100644 index 00000000..b25a92f1 --- /dev/null +++ b/wall_x/model/core/ops/csrc/silu_and_mul/silu_and_mul.cu @@ -0,0 +1,1012 @@ +// Re-enable CUDA half operators (PyTorch disables them) +#undef __CUDA_NO_HALF_OPERATORS__ +#undef __CUDA_NO_HALF_CONVERSIONS__ +#undef __CUDA_NO_HALF2_OPERATORS__ +#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include +#include + +#include +#include +#include + +#include "../common/cuda_utils.h" + +namespace wallx_cuda_silu_and_mul { + +static int get_sm_count() { + static int sm_count = -1; + if (sm_count < 0) { + int device; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device); + } + return sm_count; +} + +// Returns the stride (in elements) of dimension d. +// When strides is null (fully contiguous, row-major), computes from shape. +// When strides is non-null, reads the stored value directly. +static int64_t safe_stride(const at::Tensor& t, int d) { + if (t.is_contiguous()) { + int64_t s = 1; + for (int i = t.dim() - 1; i > d; --i) s *= t.size(i); + return s; + } + return t.stride(d); +} + +// ============================================================ +// Kernel 1: fully-contiguous path (used by non-inplace SiluAndMul) +// Persistent: launched with sm_count blocks, each loops with stride = gridDim.x * blockDim.x. +// output[i] = silu(input1[i]) * input2[i] +// ============================================================ +template +__global__ void siluAndMulPersistentKernel( + T* __restrict__ output, + const T* __restrict__ input1, + const T* __restrict__ input2, + int64_t vec_n // total elements / VEC_SIZE +) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + + const int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + const int64_t stride = (int64_t)gridDim.x * blockDim.x; + + for (int64_t i = tid; i < vec_n; i += stride) { + Vec x = VecTraits::load(input1 + i * VEC_SIZE); + Vec g = VecTraits::load(input2 + i * VEC_SIZE); + Vec out; + + if constexpr (std::is_same_v) { + out.x = x.x / (1.0f + __expf(-x.x)) * g.x; + out.y = x.y / (1.0f + __expf(-x.y)) * g.y; + out.z = x.z / (1.0f + __expf(-x.z)) * g.z; + out.w = x.w / (1.0f + __expf(-x.w)) * g.w; + } else if constexpr (std::is_same_v) { + const float x0 = __half2float(x.x), x1 = __half2float(x.y); + const float g0 = __half2float(g.x), g1 = __half2float(g.y); + out.x = __float2half(x0 / (1.0f + __expf(-x0)) * g0); + out.y = __float2half(x1 / (1.0f + __expf(-x1)) * g1); + } else { + const float x0 = __bfloat162float(x.x), x1 = __bfloat162float(x.y); + const float g0 = __bfloat162float(g.x), g1 = __bfloat162float(g.y); + out.x = __float2bfloat16(x0 / (1.0f + __expf(-x0)) * g0); + out.y = __float2bfloat16(x1 / (1.0f + __expf(-x1)) * g1); + } + + VecTraits::store(output + i * VEC_SIZE, out); + } +} + +// ============================================================ +// Kernel 2: last-dim-contiguous path (used by inplace SiluAndMulInplace) +// Requirements: stride(ndim-1) == 1 for both tensors; outer dims may be non-contiguous. +// Each work item maps to a (row, col_vec) pair. +// Row offset for each tensor is computed from its outer strides. +// ============================================================ +constexpr int MAX_OUTER_DIMS = 7; // supports tensors up to 8-D + +struct SiluMulInplaceArgs { + int64_t strides1[MAX_OUTER_DIMS]; // outer-dim strides of inout1, in elements + int64_t strides2[MAX_OUTER_DIMS]; // outer-dim strides of input2, in elements + int64_t shapes[MAX_OUTER_DIMS]; // outer-dim sizes (same for both tensors) + int outer_ndim; // ndim - 1 (0 for 1-D tensors) +}; + +struct SiluMulSplitArgs { + int64_t output_strides[MAX_OUTER_DIMS]; // outer-dim strides of output + int64_t gate_up_strides[MAX_OUTER_DIMS]; // outer-dim strides of gate_up + int64_t shapes[MAX_OUTER_DIMS]; // outer-dim sizes + int outer_ndim; // ndim - 1 +}; + +template +__global__ void siluAndMulInplacePersistentKernel( + T* __restrict__ inout1, + const T* __restrict__ input2, + SiluMulInplaceArgs args, + int64_t num_rows, // product of outer dims + int64_t vec_last // last_dim / VEC_SIZE +) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + + const int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + const int64_t grid_stride = (int64_t)gridDim.x * blockDim.x; + const int64_t total_vecs = num_rows * vec_last; + + for (int64_t vi = tid; vi < total_vecs; vi += grid_stride) { + const int64_t row = vi / vec_last; + const int64_t col_vec = vi % vec_last; + + // Compute per-row pointer offsets by decomposing the flat row index + // into per-dimension indices and accumulating strides. + int64_t offset1 = 0, offset2 = 0; + int64_t r = row; + for (int d = args.outer_ndim - 1; d >= 0; --d) { + const int64_t idx = r % args.shapes[d]; + r /= args.shapes[d]; + offset1 += idx * args.strides1[d]; + offset2 += idx * args.strides2[d]; + } + + T* ptr1 = inout1 + offset1 + col_vec * VEC_SIZE; + const T* ptr2 = input2 + offset2 + col_vec * VEC_SIZE; + + Vec x = VecTraits::load(ptr1); + Vec g = VecTraits::load(ptr2); + Vec out; + + if constexpr (std::is_same_v) { + out.x = x.x / (1.0f + __expf(-x.x)) * g.x; + out.y = x.y / (1.0f + __expf(-x.y)) * g.y; + out.z = x.z / (1.0f + __expf(-x.z)) * g.z; + out.w = x.w / (1.0f + __expf(-x.w)) * g.w; + } else if constexpr (std::is_same_v) { + const float x0 = __half2float(x.x), x1 = __half2float(x.y); + const float g0 = __half2float(g.x), g1 = __half2float(g.y); + out.x = __float2half(x0 / (1.0f + __expf(-x0)) * g0); + out.y = __float2half(x1 / (1.0f + __expf(-x1)) * g1); + } else { + const float x0 = __bfloat162float(x.x), x1 = __bfloat162float(x.y); + const float g0 = __bfloat162float(g.x), g1 = __bfloat162float(g.y); + out.x = __float2bfloat16(x0 / (1.0f + __expf(-x0)) * g0); + out.y = __float2bfloat16(x1 / (1.0f + __expf(-x1)) * g1); + } + + VecTraits::store(ptr1, out); + } +} + +// ============================================================ +// Helpers +// ============================================================ + +static int64_t get_total_elements(const at::Tensor& t) { + int64_t n = 1; + for (int i = 0; i < t.dim(); ++i) n *= t.size(i); + return n; +} + +static void dispatch_contiguous( + void* output_ptr, + const void* input1_ptr, + const void* input2_ptr, + int64_t n, + const at::Tensor& ref_tensor, + cudaStream_t stream) +{ + const int blockSize = 256; + const int num_blocks = get_sm_count(); + + if (ref_tensor.scalar_type() == at::kFloat) { + ASSERT_CHECK(n % VecTraits::vec_size == 0); + siluAndMulPersistentKernel<<>>( + static_cast(output_ptr), + static_cast(input1_ptr), + static_cast(input2_ptr), + n / VecTraits::vec_size); + } else if (ref_tensor.scalar_type() == at::kHalf) { + ASSERT_CHECK(n % VecTraits<__half>::vec_size == 0); + siluAndMulPersistentKernel<__half><<>>( + static_cast<__half*>(output_ptr), + static_cast(input1_ptr), + static_cast(input2_ptr), + n / VecTraits<__half>::vec_size); + } else if (ref_tensor.scalar_type() == at::kBFloat16) { + ASSERT_CHECK(n % VecTraits<__nv_bfloat16>::vec_size == 0); + siluAndMulPersistentKernel<__nv_bfloat16><<>>( + static_cast<__nv_bfloat16*>(output_ptr), + static_cast(input1_ptr), + static_cast(input2_ptr), + n / VecTraits<__nv_bfloat16>::vec_size); + } else { + throw std::runtime_error("silu_and_mul: unsupported dtype"); + } +} + +template +static void dispatch_inplace_strided( + T* inout1, + const T* input2, + const SiluMulInplaceArgs& args, + int64_t num_rows, + int64_t vec_last, + cudaStream_t stream) +{ + const int blockSize = 256; + const int num_blocks = get_sm_count(); + siluAndMulInplacePersistentKernel<<>>( + inout1, input2, args, num_rows, vec_last); +} + +template +__global__ void siluAndMulSplitPersistentKernel( + T* __restrict__ output, + const T* __restrict__ gate_up, + SiluMulSplitArgs args, + int64_t num_rows, + int64_t half_vec) +{ + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + const int64_t intermediate_size = half_vec * VEC_SIZE; + + const int64_t col_vec = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (col_vec >= half_vec) return; + + #pragma unroll 4 + for (int64_t row = (int64_t)blockIdx.y; row < num_rows; row += (int64_t)gridDim.y) { + int64_t output_offset = 0; + int64_t gate_up_offset = 0; + if (args.outer_ndim == 1) { + output_offset = row * args.output_strides[0]; + gate_up_offset = row * args.gate_up_strides[0]; + } else if (args.outer_ndim > 1) { + int64_t r = row; + #pragma unroll MAX_OUTER_DIMS + for (int d = MAX_OUTER_DIMS - 1; d >= 0; --d) { + if (d < args.outer_ndim) { + const int64_t idx = r % args.shapes[d]; + r /= args.shapes[d]; + output_offset += idx * args.output_strides[d]; + gate_up_offset += idx * args.gate_up_strides[d]; + } + } + } + + T* out_ptr = output + output_offset + col_vec * VEC_SIZE; + const T* gate_ptr = gate_up + gate_up_offset + col_vec * VEC_SIZE; + const T* up_ptr = gate_up + gate_up_offset + intermediate_size + col_vec * VEC_SIZE; + + Vec x = VecTraits::load(gate_ptr); + Vec g = VecTraits::load(up_ptr); + Vec out; + + if constexpr (std::is_same_v) { + out.x = x.x / (1.0f + __expf(-x.x)) * g.x; + out.y = x.y / (1.0f + __expf(-x.y)) * g.y; + out.z = x.z / (1.0f + __expf(-x.z)) * g.z; + out.w = x.w / (1.0f + __expf(-x.w)) * g.w; + } else if constexpr (std::is_same_v) { + const float x0 = __half2float(x.x), x1 = __half2float(x.y); + const float g0 = __half2float(g.x), g1 = __half2float(g.y); + out.x = __float2half(x0 / (1.0f + __expf(-x0)) * g0); + out.y = __float2half(x1 / (1.0f + __expf(-x1)) * g1); + } else { + const float x0 = __bfloat162float(x.x), x1 = __bfloat162float(x.y); + const float g0 = __bfloat162float(g.x), g1 = __bfloat162float(g.y); + out.x = __float2bfloat16(x0 / (1.0f + __expf(-x0)) * g0); + out.y = __float2bfloat16(x1 / (1.0f + __expf(-x1)) * g1); + } + + VecTraits::store(out_ptr, out); + } +} + +template +static void dispatch_split_strided( + T* output, + const T* gate_up, + const SiluMulSplitArgs& args, + int64_t num_rows, + int64_t half_vec, + cudaStream_t stream) +{ + constexpr int BLOCK_X = 256; + const int grid_x = (int)((half_vec + BLOCK_X - 1) / BLOCK_X); + const int grid_y = (int)std::min(num_rows, (int64_t)get_sm_count()); + siluAndMulSplitPersistentKernel<<>>( + output, gate_up, args, num_rows, half_vec); +} + +// ============================================================ +// Public functions +// ============================================================ + +// Non-inplace: output = silu(input1) * input2 (all tensors must be fully contiguous) +void SiluAndMul( + const at::Tensor& output, + const at::Tensor& input1, + const at::Tensor& input2) +{ + ASSERT_CHECK(output.is_contiguous()); + ASSERT_CHECK(input1.is_contiguous()); + ASSERT_CHECK(input2.is_contiguous()); + ASSERT_CHECK(output.dim() == input1.dim()); + ASSERT_CHECK(input1.dim() == input2.dim()); + for (int i = 0; i < input1.dim(); ++i) { + ASSERT_CHECK(output.size(i) == input1.size(i)); + ASSERT_CHECK(input1.size(i) == input2.size(i)); + } + + const int64_t n = get_total_elements(input1); + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + dispatch_contiguous(output.data_ptr(), input1.data_ptr(), input2.data_ptr(), n, input1, stream); + sync_check_cuda_error(); +} + +// Inplace: input1 = silu(input1) * input2 +// Only requires the last dimension to be contiguous (stride[-1] == 1). +// Outer dimensions may be non-contiguous; their offsets are computed via strides. +void SiluAndMulInplace( + const at::Tensor& input1, + const at::Tensor& input2) +{ + const int ndim = input1.dim(); + ASSERT_CHECK(ndim >= 1); + ASSERT_CHECK(ndim == input2.dim()); + ASSERT_CHECK(ndim - 1 <= MAX_OUTER_DIMS); + + // Shape must match on every dimension. + for (int i = 0; i < ndim; ++i) { + ASSERT_CHECK(input1.size(i) == input2.size(i)); + } + + // Only the last dimension must be contiguous. + ASSERT_CHECK(safe_stride(input1, ndim - 1) == 1); + ASSERT_CHECK(safe_stride(input2, ndim - 1) == 1); + + const int64_t last_dim = input1.size(ndim - 1); + const int outer_ndim = ndim - 1; + + // Build kernel args: collect outer-dim strides and shapes. + SiluMulInplaceArgs args{}; + args.outer_ndim = outer_ndim; + int64_t num_rows = 1; + for (int d = 0; d < outer_ndim; ++d) { + args.strides1[d] = safe_stride(input1, d); + args.strides2[d] = safe_stride(input2, d); + args.shapes[d] = input1.size(d); + num_rows *= input1.size(d); + } + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + if (input1.scalar_type() == at::kFloat) { + ASSERT_CHECK(last_dim % VecTraits::vec_size == 0); + dispatch_inplace_strided( + static_cast(input1.data_ptr()), + static_cast(input2.data_ptr()), + args, num_rows, last_dim / VecTraits::vec_size, stream); + } else if (input1.scalar_type() == at::kHalf) { + ASSERT_CHECK(last_dim % VecTraits<__half>::vec_size == 0); + dispatch_inplace_strided( + static_cast<__half*>(input1.data_ptr()), + static_cast(input2.data_ptr()), + args, num_rows, last_dim / VecTraits<__half>::vec_size, stream); + } else if (input1.scalar_type() == at::kBFloat16) { + ASSERT_CHECK(last_dim % VecTraits<__nv_bfloat16>::vec_size == 0); + dispatch_inplace_strided( + static_cast<__nv_bfloat16*>(input1.data_ptr()), + static_cast(input2.data_ptr()), + args, num_rows, last_dim / VecTraits<__nv_bfloat16>::vec_size, stream); + } else { + throw std::runtime_error("silu_and_mul_inplace: unsupported dtype"); + } + + sync_check_cuda_error(); +} + +// ============================================================ +// Kernel 3: split+silu+mul inplace (2-D grid, division-free hot path) +// +// Grid layout: +// blockIdx.x * blockDim.x + threadIdx.x → col_vec ∈ [0, half_vec) +// blockIdx.y (grid-stride over rows) → row ∈ [0, num_rows) +// +// This eliminates the per-iteration vi/half_vec and vi%half_vec +// 64-bit integer divisions present in the original 1-D scheme. +// +// Input: gate_up [..., 2*N], last dim contiguous (stride[-1] == 1). +// Output (inplace): gate_up[..., :N] = silu(gate_up[..., :N]) * gate_up[..., N:] +// ============================================================ +struct SiluMulSplitInplaceArgs { + int64_t outer_strides[MAX_OUTER_DIMS]; // outer-dim strides, in elements + int64_t shapes[MAX_OUTER_DIMS]; // outer-dim sizes + int outer_ndim; // ndim - 1 +}; + +struct SiluMulSplitBackwardArgs { + int64_t grad_gate_up_strides[MAX_OUTER_DIMS]; // outer-dim strides of grad_gate_up + int64_t grad_output_strides[MAX_OUTER_DIMS]; // outer-dim strides of grad_output + int64_t gate_up_strides[MAX_OUTER_DIMS]; // outer-dim strides of original gate_up + int64_t shapes[MAX_OUTER_DIMS]; // outer-dim sizes + int outer_ndim; // ndim - 1 +}; + +// Load a vector through the read-only data cache (__ldg). +// up_ptr shares the same base allocation as inout but is never written, +// so routing it through the texture/RO path is both safe and faster. +template +__device__ __forceinline__ typename VecTraits::Type ldg_vec(const T* ptr) { + using Vec = typename VecTraits::Type; + return __ldg(reinterpret_cast(ptr)); +} + +template +__global__ void siluAndMulSplitInplacePersistentKernel( + T* __restrict__ inout, + SiluMulSplitInplaceArgs args, + int64_t num_rows, + int64_t half_vec // intermediate_size / VEC_SIZE +) { + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + // Derive intermediate_size inside the kernel — avoids one extra parameter. + const int64_t intermediate_size = half_vec * VEC_SIZE; + + const int64_t col_vec = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (col_vec >= half_vec) return; // guard last partial block + + // Unroll by 2: allows the compiler to interleave memory loads from two + // consecutive rows, hiding global-memory latency at the cost of ~2x + // register pressure on the row-local variables. + #pragma unroll 4 + for (int64_t row = (int64_t)blockIdx.y; row < num_rows; row += (int64_t)gridDim.y) { + // ------------------------------------------------------------------ + // Row offset computation + // outer_ndim == 0 → always row_offset = 0 (1-D tensor) + // outer_ndim == 1 → single multiply, no division (most common) + // outer_ndim > 1 → fully-unrolled decomposition loop + // ------------------------------------------------------------------ + int64_t row_offset = 0; + if (args.outer_ndim == 1) { + row_offset = row * args.outer_strides[0]; + } else if (args.outer_ndim > 1) { + int64_t r = row; + // Unroll over the compile-time maximum depth; the runtime guard + // `d < args.outer_ndim` is loop-invariant and uniform across the + // warp, so the compiler emits predicated instructions with no + // divergence. + #pragma unroll MAX_OUTER_DIMS + for (int d = MAX_OUTER_DIMS - 1; d >= 0; --d) { + if (d < args.outer_ndim) { + const int64_t idx = r % args.shapes[d]; + r /= args.shapes[d]; + row_offset += idx * args.outer_strides[d]; + } + } + } + + T* gate_ptr = inout + row_offset + col_vec * VEC_SIZE; + const T* up_ptr = inout + row_offset + intermediate_size + col_vec * VEC_SIZE; + + Vec x = VecTraits::load(gate_ptr); + Vec g = ldg_vec(up_ptr); // read-only: steer through RO cache + Vec out; + + if constexpr (std::is_same_v) { + // float4 — 4 lanes, explicitly unrolled + out.x = x.x / (1.0f + __expf(-x.x)) * g.x; + out.y = x.y / (1.0f + __expf(-x.y)) * g.y; + out.z = x.z / (1.0f + __expf(-x.z)) * g.z; + out.w = x.w / (1.0f + __expf(-x.w)) * g.w; + } else if constexpr (std::is_same_v) { + // half2 — 2 lanes, explicitly unrolled + const float x0 = __half2float(x.x), x1 = __half2float(x.y); + const float g0 = __half2float(g.x), g1 = __half2float(g.y); + out.x = __float2half(x0 / (1.0f + __expf(-x0)) * g0); + out.y = __float2half(x1 / (1.0f + __expf(-x1)) * g1); + } else { + // bfloat162 — 2 lanes, explicitly unrolled + const float x0 = __bfloat162float(x.x), x1 = __bfloat162float(x.y); + const float g0 = __bfloat162float(g.x), g1 = __bfloat162float(g.y); + out.x = __float2bfloat16(x0 / (1.0f + __expf(-x0)) * g0); + out.y = __float2bfloat16(x1 / (1.0f + __expf(-x1)) * g1); + } + + VecTraits::store(gate_ptr, out); + } +} + +template +static void dispatch_split_inplace_strided( + T* inout, + const SiluMulSplitInplaceArgs& args, + int64_t num_rows, + int64_t half_vec, + cudaStream_t stream) +{ + constexpr int BLOCK_X = 256; + // grid_x covers the col_vec (intermediate) dimension exactly. + // grid_y covers rows up to sm_count; the kernel loops with stride gridDim.y. + const int grid_x = (int)((half_vec + BLOCK_X - 1) / BLOCK_X); + const int grid_y = (int)std::min(num_rows, (int64_t)get_sm_count()); + siluAndMulSplitInplacePersistentKernel<<>>( + inout, args, num_rows, half_vec); +} + +template +__device__ __forceinline__ float to_float(T x); + +template<> +__device__ __forceinline__ float to_float(float x) { + return x; +} + +template<> +__device__ __forceinline__ float to_float<__half>(__half x) { + return __half2float(x); +} + +template<> +__device__ __forceinline__ float to_float<__nv_bfloat16>(__nv_bfloat16 x) { + return __bfloat162float(x); +} + +template +__device__ __forceinline__ T from_float(float x); + +template<> +__device__ __forceinline__ float from_float(float x) { + return x; +} + +template<> +__device__ __forceinline__ __half from_float<__half>(float x) { + return __float2half(x); +} + +template<> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float x) { + return __float2bfloat16(x); +} + +template +__device__ __forceinline__ void silu_mul_backward_lane( + T gate, + T up, + T grad_output, + T& grad_gate, + T& grad_up) +{ + const float gate_f = to_float(gate); + const float up_f = to_float(up); + const float grad_f = to_float(grad_output); + const float sigmoid = 1.0f / (1.0f + __expf(-gate_f)); + const float silu = gate_f * sigmoid; + const float silu_grad = sigmoid * (1.0f + gate_f * (1.0f - sigmoid)); + + grad_gate = from_float(grad_f * up_f * silu_grad); + grad_up = from_float(grad_f * silu); +} + +template +__device__ __forceinline__ void silu_mul_backward_vec( + typename VecTraits::Type gate, + typename VecTraits::Type up, + typename VecTraits::Type grad_output, + typename VecTraits::Type& grad_gate, + typename VecTraits::Type& grad_up) +{ + silu_mul_backward_lane(gate.x, up.x, grad_output.x, grad_gate.x, grad_up.x); + silu_mul_backward_lane(gate.y, up.y, grad_output.y, grad_gate.y, grad_up.y); + + if constexpr (std::is_same_v) { + silu_mul_backward_lane(gate.z, up.z, grad_output.z, grad_gate.z, grad_up.z); + silu_mul_backward_lane(gate.w, up.w, grad_output.w, grad_gate.w, grad_up.w); + } +} + +// Backward for split+silu+mul inplace. +// Given original gate_up [..., 2*N] and grad_output [..., N], writes +// grad_gate_up[..., :N] = grad_output * up * d_silu(gate) +// grad_gate_up[..., N:] = grad_output * silu(gate) +template +__global__ void siluAndMulSplitInplaceBackwardPersistentKernel( + T* __restrict__ grad_gate_up, + const T* __restrict__ grad_output, + const T* __restrict__ gate_up, + SiluMulSplitBackwardArgs args, + int64_t num_rows, + int64_t half_vec) +{ + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + const int64_t intermediate_size = half_vec * VEC_SIZE; + + const int64_t col_vec = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (col_vec >= half_vec) return; + + #pragma unroll 4 + for (int64_t row = (int64_t)blockIdx.y; row < num_rows; row += (int64_t)gridDim.y) { + int64_t grad_gate_up_offset = 0; + int64_t grad_output_offset = 0; + int64_t gate_up_offset = 0; + + if (args.outer_ndim == 1) { + grad_gate_up_offset = row * args.grad_gate_up_strides[0]; + grad_output_offset = row * args.grad_output_strides[0]; + gate_up_offset = row * args.gate_up_strides[0]; + } else if (args.outer_ndim > 1) { + int64_t r = row; + #pragma unroll MAX_OUTER_DIMS + for (int d = MAX_OUTER_DIMS - 1; d >= 0; --d) { + if (d < args.outer_ndim) { + const int64_t idx = r % args.shapes[d]; + r /= args.shapes[d]; + grad_gate_up_offset += idx * args.grad_gate_up_strides[d]; + grad_output_offset += idx * args.grad_output_strides[d]; + gate_up_offset += idx * args.gate_up_strides[d]; + } + } + } + + T* grad_gate_ptr = grad_gate_up + grad_gate_up_offset + col_vec * VEC_SIZE; + T* grad_up_ptr = grad_gate_up + grad_gate_up_offset + intermediate_size + col_vec * VEC_SIZE; + const T* grad_output_ptr = grad_output + grad_output_offset + col_vec * VEC_SIZE; + const T* gate_ptr = gate_up + gate_up_offset + col_vec * VEC_SIZE; + const T* up_ptr = gate_up + gate_up_offset + intermediate_size + col_vec * VEC_SIZE; + + Vec gate = VecTraits::load(gate_ptr); + Vec up = VecTraits::load(up_ptr); + Vec grad = VecTraits::load(grad_output_ptr); + Vec grad_gate; + Vec grad_up; + + silu_mul_backward_vec(gate, up, grad, grad_gate, grad_up); + + VecTraits::store(grad_gate_ptr, grad_gate); + VecTraits::store(grad_up_ptr, grad_up); + } +} + +template +__global__ void siluAndMulBackwardPersistentKernel( + T* __restrict__ grad_input1, + T* __restrict__ grad_input2, + const T* __restrict__ grad_output, + const T* __restrict__ input1, + const T* __restrict__ input2, + int64_t vec_n) +{ + using Vec = typename VecTraits::Type; + constexpr int VEC_SIZE = VecTraits::vec_size; + + const int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + const int64_t stride = (int64_t)gridDim.x * blockDim.x; + + for (int64_t i = tid; i < vec_n; i += stride) { + Vec gate = VecTraits::load(input1 + i * VEC_SIZE); + Vec up = VecTraits::load(input2 + i * VEC_SIZE); + Vec grad = VecTraits::load(grad_output + i * VEC_SIZE); + Vec grad_gate; + Vec grad_up; + + silu_mul_backward_vec(gate, up, grad, grad_gate, grad_up); + + VecTraits::store(grad_input1 + i * VEC_SIZE, grad_gate); + VecTraits::store(grad_input2 + i * VEC_SIZE, grad_up); + } +} + +static void dispatch_contiguous_backward( + void* grad_input1_ptr, + void* grad_input2_ptr, + const void* grad_output_ptr, + const void* input1_ptr, + const void* input2_ptr, + int64_t n, + const at::Tensor& ref_tensor, + cudaStream_t stream) +{ + const int blockSize = 256; + const int num_blocks = get_sm_count(); + + if (ref_tensor.scalar_type() == at::kFloat) { + ASSERT_CHECK(n % VecTraits::vec_size == 0); + siluAndMulBackwardPersistentKernel<<>>( + static_cast(grad_input1_ptr), + static_cast(grad_input2_ptr), + static_cast(grad_output_ptr), + static_cast(input1_ptr), + static_cast(input2_ptr), + n / VecTraits::vec_size); + } else if (ref_tensor.scalar_type() == at::kHalf) { + ASSERT_CHECK(n % VecTraits<__half>::vec_size == 0); + siluAndMulBackwardPersistentKernel<__half><<>>( + static_cast<__half*>(grad_input1_ptr), + static_cast<__half*>(grad_input2_ptr), + static_cast(grad_output_ptr), + static_cast(input1_ptr), + static_cast(input2_ptr), + n / VecTraits<__half>::vec_size); + } else if (ref_tensor.scalar_type() == at::kBFloat16) { + ASSERT_CHECK(n % VecTraits<__nv_bfloat16>::vec_size == 0); + siluAndMulBackwardPersistentKernel<__nv_bfloat16><<>>( + static_cast<__nv_bfloat16*>(grad_input1_ptr), + static_cast<__nv_bfloat16*>(grad_input2_ptr), + static_cast(grad_output_ptr), + static_cast(input1_ptr), + static_cast(input2_ptr), + n / VecTraits<__nv_bfloat16>::vec_size); + } else { + throw std::runtime_error("silu_and_mul_backward: unsupported dtype"); + } +} + +template +static void dispatch_split_inplace_backward_strided( + T* grad_gate_up, + const T* grad_output, + const T* gate_up, + const SiluMulSplitBackwardArgs& args, + int64_t num_rows, + int64_t half_vec, + cudaStream_t stream) +{ + constexpr int BLOCK_X = 256; + const int grid_x = (int)((half_vec + BLOCK_X - 1) / BLOCK_X); + const int grid_y = (int)std::min(num_rows, (int64_t)get_sm_count()); + siluAndMulSplitInplaceBackwardPersistentKernel<<>>( + grad_gate_up, grad_output, gate_up, args, num_rows, half_vec); +} + +// SiluAndMulSplit: fused split + silu + mul, out-of-place. +// gate_up has shape [..., 2*N], output has shape [..., N]. +// Only the last dimension must be contiguous; outer dimensions may be strided. +void SiluAndMulSplit( + const at::Tensor& output, + const at::Tensor& gate_up) +{ + const int ndim = gate_up.dim(); + ASSERT_CHECK(ndim >= 1); + ASSERT_CHECK(ndim == output.dim()); + ASSERT_CHECK(ndim - 1 <= MAX_OUTER_DIMS); + ASSERT_CHECK(safe_stride(gate_up, ndim - 1) == 1); + ASSERT_CHECK(safe_stride(output, ndim - 1) == 1); + + const int64_t last_dim = gate_up.size(ndim - 1); + ASSERT_CHECK(last_dim % 2 == 0); + const int64_t intermediate_size = last_dim / 2; + ASSERT_CHECK(output.size(ndim - 1) == intermediate_size); + + const int outer_ndim = ndim - 1; + SiluMulSplitArgs args{}; + args.outer_ndim = outer_ndim; + int64_t num_rows = 1; + for (int d = 0; d < outer_ndim; ++d) { + ASSERT_CHECK(output.size(d) == gate_up.size(d)); + args.output_strides[d] = safe_stride(output, d); + args.gate_up_strides[d] = safe_stride(gate_up, d); + args.shapes[d] = gate_up.size(d); + num_rows *= gate_up.size(d); + } + + ASSERT_CHECK(output.scalar_type() == gate_up.scalar_type()); + ASSERT_CHECK(output.element_size() == gate_up.element_size()); + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + if (gate_up.scalar_type() == at::kFloat) { + ASSERT_CHECK(intermediate_size % VecTraits::vec_size == 0); + dispatch_split_strided( + static_cast(output.data_ptr()), + static_cast(gate_up.data_ptr()), + args, num_rows, intermediate_size / VecTraits::vec_size, stream); + } else if (gate_up.scalar_type() == at::kHalf) { + ASSERT_CHECK(intermediate_size % VecTraits<__half>::vec_size == 0); + dispatch_split_strided( + static_cast<__half*>(output.data_ptr()), + static_cast(gate_up.data_ptr()), + args, num_rows, intermediate_size / VecTraits<__half>::vec_size, stream); + } else if (gate_up.scalar_type() == at::kBFloat16) { + ASSERT_CHECK(intermediate_size % VecTraits<__nv_bfloat16>::vec_size == 0); + dispatch_split_strided( + static_cast<__nv_bfloat16*>(output.data_ptr()), + static_cast(gate_up.data_ptr()), + args, num_rows, intermediate_size / VecTraits<__nv_bfloat16>::vec_size, stream); + } else { + throw std::runtime_error("silu_and_mul_split: unsupported dtype"); + } + + sync_check_cuda_error(); +} + +// SiluAndMulSplitInplace: fused split + silu + mul, inplace. +// gate_up has shape [..., 2*N], last dim must be contiguous and even. +// After this call, gate_up[..., :N] = silu(gate_up[..., :N]) * gate_up[..., N:]. +void SiluAndMulSplitInplace(const at::Tensor& gate_up) +{ + const int ndim = gate_up.dim(); + ASSERT_CHECK(ndim >= 1); + ASSERT_CHECK(ndim - 1 <= MAX_OUTER_DIMS); + ASSERT_CHECK(safe_stride(gate_up, ndim - 1) == 1); + + const int64_t last_dim = gate_up.size(ndim - 1); + ASSERT_CHECK(last_dim % 2 == 0); + const int64_t intermediate_size = last_dim / 2; + + const int outer_ndim = ndim - 1; + SiluMulSplitInplaceArgs args{}; + args.outer_ndim = outer_ndim; + int64_t num_rows = 1; + for (int d = 0; d < outer_ndim; ++d) { + args.outer_strides[d] = safe_stride(gate_up, d); + args.shapes[d] = gate_up.size(d); + num_rows *= gate_up.size(d); + } + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + if (gate_up.scalar_type() == at::kFloat) { + ASSERT_CHECK(intermediate_size % VecTraits::vec_size == 0); + dispatch_split_inplace_strided( + static_cast(gate_up.data_ptr()), + args, num_rows, intermediate_size / VecTraits::vec_size, stream); + } else if (gate_up.scalar_type() == at::kHalf) { + ASSERT_CHECK(intermediate_size % VecTraits<__half>::vec_size == 0); + dispatch_split_inplace_strided( + static_cast<__half*>(gate_up.data_ptr()), + args, num_rows, intermediate_size / VecTraits<__half>::vec_size, stream); + } else if (gate_up.scalar_type() == at::kBFloat16) { + ASSERT_CHECK(intermediate_size % VecTraits<__nv_bfloat16>::vec_size == 0); + dispatch_split_inplace_strided( + static_cast<__nv_bfloat16*>(gate_up.data_ptr()), + args, num_rows, intermediate_size / VecTraits<__nv_bfloat16>::vec_size, stream); + } else { + throw std::runtime_error("silu_and_mul_split_inplace: unsupported dtype"); + } + + sync_check_cuda_error(); +} + +void SiluAndMulSplitInplaceBackward( + const at::Tensor& grad_gate_up, + const at::Tensor& grad_output, + const at::Tensor& gate_up) +{ + const int ndim = gate_up.dim(); + ASSERT_CHECK(ndim >= 1); + ASSERT_CHECK(ndim == grad_gate_up.dim()); + ASSERT_CHECK(ndim == grad_output.dim()); + ASSERT_CHECK(ndim - 1 <= MAX_OUTER_DIMS); + + ASSERT_CHECK(safe_stride(grad_gate_up, ndim - 1) == 1); + ASSERT_CHECK(safe_stride(grad_output, ndim - 1) == 1); + ASSERT_CHECK(safe_stride(gate_up, ndim - 1) == 1); + + const int64_t last_dim = gate_up.size(ndim - 1); + ASSERT_CHECK(last_dim % 2 == 0); + const int64_t intermediate_size = last_dim / 2; + + for (int d = 0; d < ndim - 1; ++d) { + ASSERT_CHECK(grad_gate_up.size(d) == gate_up.size(d)); + ASSERT_CHECK(grad_output.size(d) == gate_up.size(d)); + } + ASSERT_CHECK(grad_gate_up.size(ndim - 1) == last_dim); + ASSERT_CHECK(grad_output.size(ndim - 1) == intermediate_size); + + ASSERT_CHECK(grad_gate_up.scalar_type() == gate_up.scalar_type()); + ASSERT_CHECK(grad_gate_up.element_size() == gate_up.element_size()); + ASSERT_CHECK(grad_output.scalar_type() == gate_up.scalar_type()); + ASSERT_CHECK(grad_output.element_size() == gate_up.element_size()); + + const int outer_ndim = ndim - 1; + SiluMulSplitBackwardArgs args{}; + args.outer_ndim = outer_ndim; + int64_t num_rows = 1; + for (int d = 0; d < outer_ndim; ++d) { + args.grad_gate_up_strides[d] = safe_stride(grad_gate_up, d); + args.grad_output_strides[d] = safe_stride(grad_output, d); + args.gate_up_strides[d] = safe_stride(gate_up, d); + args.shapes[d] = gate_up.size(d); + num_rows *= gate_up.size(d); + } + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + if (gate_up.scalar_type() == at::kFloat) { + ASSERT_CHECK(intermediate_size % VecTraits::vec_size == 0); + dispatch_split_inplace_backward_strided( + static_cast(grad_gate_up.data_ptr()), + static_cast(grad_output.data_ptr()), + static_cast(gate_up.data_ptr()), + args, num_rows, intermediate_size / VecTraits::vec_size, stream); + } else if (gate_up.scalar_type() == at::kHalf) { + ASSERT_CHECK(intermediate_size % VecTraits<__half>::vec_size == 0); + dispatch_split_inplace_backward_strided( + static_cast<__half*>(grad_gate_up.data_ptr()), + static_cast(grad_output.data_ptr()), + static_cast(gate_up.data_ptr()), + args, num_rows, intermediate_size / VecTraits<__half>::vec_size, stream); + } else if (gate_up.scalar_type() == at::kBFloat16) { + ASSERT_CHECK(intermediate_size % VecTraits<__nv_bfloat16>::vec_size == 0); + dispatch_split_inplace_backward_strided( + static_cast<__nv_bfloat16*>(grad_gate_up.data_ptr()), + static_cast(grad_output.data_ptr()), + static_cast(gate_up.data_ptr()), + args, num_rows, intermediate_size / VecTraits<__nv_bfloat16>::vec_size, stream); + } else { + throw std::runtime_error("silu_and_mul_split_inplace_backward: unsupported dtype"); + } + + sync_check_cuda_error(); +} + +void SiluAndMulBackward( + const at::Tensor& grad_input1, + const at::Tensor& grad_input2, + const at::Tensor& grad_output, + const at::Tensor& input1, + const at::Tensor& input2) +{ + ASSERT_CHECK(grad_input1.is_contiguous()); + ASSERT_CHECK(grad_input2.is_contiguous()); + ASSERT_CHECK(grad_output.is_contiguous()); + ASSERT_CHECK(input1.is_contiguous()); + ASSERT_CHECK(input2.is_contiguous()); + + ASSERT_CHECK(input1.dim() == input2.dim()); + ASSERT_CHECK(input1.dim() == grad_output.dim()); + ASSERT_CHECK(input1.dim() == grad_input1.dim()); + ASSERT_CHECK(input1.dim() == grad_input2.dim()); + for (int i = 0; i < input1.dim(); ++i) { + ASSERT_CHECK(input1.size(i) == input2.size(i)); + ASSERT_CHECK(input1.size(i) == grad_output.size(i)); + ASSERT_CHECK(input1.size(i) == grad_input1.size(i)); + ASSERT_CHECK(input1.size(i) == grad_input2.size(i)); + } + + ASSERT_CHECK(grad_input1.scalar_type() == input1.scalar_type()); + ASSERT_CHECK(grad_input1.element_size() == input1.element_size()); + ASSERT_CHECK(grad_input2.scalar_type() == input1.scalar_type()); + ASSERT_CHECK(grad_input2.element_size() == input1.element_size()); + ASSERT_CHECK(grad_output.scalar_type() == input1.scalar_type()); + ASSERT_CHECK(grad_output.element_size() == input1.element_size()); + ASSERT_CHECK(input2.scalar_type() == input1.scalar_type()); + ASSERT_CHECK(input2.element_size() == input1.element_size()); + + const int64_t n = get_total_elements(input1); + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + dispatch_contiguous_backward( + grad_input1.data_ptr(), + grad_input2.data_ptr(), + grad_output.data_ptr(), + input1.data_ptr(), + input2.data_ptr(), + n, + input1, + stream); + sync_check_cuda_error(); +} + +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) + +} // namespace wallx_cuda_silu_and_mul diff --git a/wall_x/model/core/ops/csrc/two_expert_scatter/two_expert_scatter.cu b/wall_x/model/core/ops/csrc/two_expert_scatter/two_expert_scatter.cu new file mode 100644 index 00000000..0b2f0931 --- /dev/null +++ b/wall_x/model/core/ops/csrc/two_expert_scatter/two_expert_scatter.cu @@ -0,0 +1,520 @@ +// Re-enable CUDA half operators (PyTorch disables them) +#undef __CUDA_NO_HALF_OPERATORS__ +#undef __CUDA_NO_HALF_CONVERSIONS__ +#undef __CUDA_NO_HALF2_OPERATORS__ +#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "../common/cuda_utils.h" + +namespace wallx_cuda_two_expert_scatter { + +static int checked_row_blocks(int64_t batch_size, int64_t seq_length) { + ASSERT_CHECK(batch_size >= 0); + ASSERT_CHECK(seq_length >= 0); + const int64_t max_blocks = std::numeric_limits::max(); + if (seq_length != 0 && batch_size > max_blocks / seq_length) { + throw std::runtime_error("two_expert_scatter: batch_size * seq_length exceeds INT_MAX"); + } + return static_cast(batch_size * seq_length); +} + +static void check_index_tensor( + const at::Tensor& tensor, + const char* name, + int64_t expected_size, + c10::DeviceType device_type, + c10::DeviceIndex device_id) { + const std::string prefix = std::string(name) + ": "; + if (tensor.dim() != 1) { + throw std::runtime_error( + prefix + "expected 1D tensor, got " + std::to_string(tensor.dim()) + "D"); + } + if (tensor.size(0) != expected_size) { + throw std::runtime_error( + prefix + "expected size " + std::to_string(expected_size) + ", got " + + std::to_string(tensor.size(0))); + } + if (!tensor.is_contiguous()) { + throw std::runtime_error(prefix + "expected contiguous tensor"); + } + if (!(tensor.scalar_type() == at::kLong)) { + throw std::runtime_error(prefix + "expected int64 tensor"); + } + if (tensor.device().type() != device_type || + tensor.device().index() != device_id) { + throw std::runtime_error(prefix + "expected tensor on the same device as input"); + } +} + +static bool can_use_vec16(const void* output, const void* input, int64_t row_bytes) { + return row_bytes % static_cast(sizeof(uint4)) == 0 && + (reinterpret_cast(output) % alignof(uint4)) == 0 && + (reinterpret_cast(input) % alignof(uint4)) == 0; +} + +static int row_copy_threads(int64_t row_items) { + if (row_items <= 64) { + return 64; + } + if (row_items <= 128) { + return 128; + } + return 256; +} + +__device__ __forceinline__ int64_t find_batch_for_packed_row( + const int64_t* __restrict__ offsets, + int64_t batch_size, + int64_t row) { + int64_t low = 0; + int64_t high = batch_size; + while (low < high) { + const int64_t mid = (low + high) >> 1; + if (offsets[mid + 1] <= row) { + low = mid + 1; + } else { + high = mid; + } + } + return low; +} + +__device__ __forceinline__ int64_t source_row_from_packed_row( + int64_t packed_row, + int64_t batch_size, + int64_t seq_length, + const int64_t* __restrict__ action_starts, + const int64_t* __restrict__ action_ends, + const int64_t* __restrict__ llm_offsets, + const int64_t* __restrict__ action_offsets) { + const int64_t total_llm = llm_offsets[batch_size]; + + if (packed_row < total_llm) { + const int64_t batch = find_batch_for_packed_row(llm_offsets, batch_size, packed_row); + const int64_t local_llm = packed_row - llm_offsets[batch]; + const int64_t action_start = action_starts[batch]; + const int64_t action_len = action_ends[batch] - action_start; + const int64_t seq = (local_llm < action_start) ? local_llm : (local_llm + action_len); + return batch * seq_length + seq; + } + + const int64_t action_row = packed_row - total_llm; + const int64_t batch = find_batch_for_packed_row(action_offsets, batch_size, action_row); + const int64_t seq = action_starts[batch] + (action_row - action_offsets[batch]); + return batch * seq_length + seq; +} + +__device__ __forceinline__ int64_t packed_row_from_source_row( + int64_t source_row, + int64_t batch_size, + int64_t seq_length, + const int64_t* __restrict__ action_starts, + const int64_t* __restrict__ action_ends, + const int64_t* __restrict__ llm_offsets, + const int64_t* __restrict__ action_offsets) { + const int64_t batch = source_row / seq_length; + const int64_t seq = source_row - batch * seq_length; + const int64_t action_start = action_starts[batch]; + const int64_t action_end = action_ends[batch]; + const int64_t action_len = action_end - action_start; + const int64_t total_llm = llm_offsets[batch_size]; + + if (seq >= action_start && seq < action_end) { + return total_llm + action_offsets[batch] + (seq - action_start); + } + + const int64_t local_llm = (seq < action_start) ? seq : (seq - action_len); + return llm_offsets[batch] + local_llm; +} + +template +__global__ void twoExpertPackRowsKernel( + T* __restrict__ output, + const T* __restrict__ input, + int64_t batch_size, + int64_t seq_length, + int64_t hidden_dim, + const int64_t* __restrict__ action_starts, + const int64_t* __restrict__ action_ends, + const int64_t* __restrict__ llm_offsets, + const int64_t* __restrict__ action_offsets) { + const int64_t packed_row = blockIdx.x; + const int64_t src_row = source_row_from_packed_row( + packed_row, batch_size, seq_length, action_starts, action_ends, llm_offsets, action_offsets); + + for (int64_t col = threadIdx.x; col < hidden_dim; col += blockDim.x) { + output[packed_row * hidden_dim + col] = input[src_row * hidden_dim + col]; + } +} + +__global__ void twoExpertPackVec16Kernel( + uint4* __restrict__ output, + const uint4* __restrict__ input, + int64_t batch_size, + int64_t seq_length, + int64_t row_vecs, + const int64_t* __restrict__ action_starts, + const int64_t* __restrict__ action_ends, + const int64_t* __restrict__ llm_offsets, + const int64_t* __restrict__ action_offsets) { + const int64_t packed_row = blockIdx.x; + const int64_t src_row = source_row_from_packed_row( + packed_row, batch_size, seq_length, action_starts, action_ends, llm_offsets, action_offsets); + + for (int64_t col = threadIdx.x; col < row_vecs; col += blockDim.x) { + output[packed_row * row_vecs + col] = input[src_row * row_vecs + col]; + } +} + +template +__global__ void twoExpertUnpackRowsKernel( + T* __restrict__ output, + const T* __restrict__ input, + int64_t batch_size, + int64_t seq_length, + int64_t hidden_dim, + const int64_t* __restrict__ action_starts, + const int64_t* __restrict__ action_ends, + const int64_t* __restrict__ llm_offsets, + const int64_t* __restrict__ action_offsets) { + const int64_t src_row = blockIdx.x; + const int64_t packed_row = packed_row_from_source_row( + src_row, batch_size, seq_length, action_starts, action_ends, llm_offsets, action_offsets); + + for (int64_t col = threadIdx.x; col < hidden_dim; col += blockDim.x) { + output[src_row * hidden_dim + col] = input[packed_row * hidden_dim + col]; + } +} + +__global__ void twoExpertUnpackVec16Kernel( + uint4* __restrict__ output, + const uint4* __restrict__ input, + int64_t batch_size, + int64_t seq_length, + int64_t row_vecs, + const int64_t* __restrict__ action_starts, + const int64_t* __restrict__ action_ends, + const int64_t* __restrict__ llm_offsets, + const int64_t* __restrict__ action_offsets) { + const int64_t src_row = blockIdx.x; + const int64_t packed_row = packed_row_from_source_row( + src_row, batch_size, seq_length, action_starts, action_ends, llm_offsets, action_offsets); + + for (int64_t col = threadIdx.x; col < row_vecs; col += blockDim.x) { + output[src_row * row_vecs + col] = input[packed_row * row_vecs + col]; + } +} + +static void launch_pack_vec16( + void* output, + const void* input, + int64_t batch_size, + int64_t seq_length, + int64_t row_bytes, + const int64_t* action_starts, + const int64_t* action_ends, + const int64_t* llm_offsets, + const int64_t* action_offsets, + cudaStream_t stream) { + const int64_t row_vecs = row_bytes / static_cast(sizeof(uint4)); + const int blocks = checked_row_blocks(batch_size, seq_length); + const int threads = row_copy_threads(row_vecs); + twoExpertPackVec16Kernel<<>>( + static_cast(output), + static_cast(input), + batch_size, + seq_length, + row_vecs, + action_starts, + action_ends, + llm_offsets, + action_offsets); +} + +static void launch_unpack_vec16( + void* output, + const void* input, + int64_t batch_size, + int64_t seq_length, + int64_t row_bytes, + const int64_t* action_starts, + const int64_t* action_ends, + const int64_t* llm_offsets, + const int64_t* action_offsets, + cudaStream_t stream) { + const int64_t row_vecs = row_bytes / static_cast(sizeof(uint4)); + const int blocks = checked_row_blocks(batch_size, seq_length); + const int threads = row_copy_threads(row_vecs); + twoExpertUnpackVec16Kernel<<>>( + static_cast(output), + static_cast(input), + batch_size, + seq_length, + row_vecs, + action_starts, + action_ends, + llm_offsets, + action_offsets); +} + +template +static void launch_pack( + T* output, + const T* input, + int64_t batch_size, + int64_t seq_length, + int64_t hidden_dim, + const int64_t* action_starts, + const int64_t* action_ends, + const int64_t* llm_offsets, + const int64_t* action_offsets, + cudaStream_t stream) { + const int64_t row_bytes = hidden_dim * static_cast(sizeof(T)); + if (can_use_vec16(output, input, row_bytes)) { + launch_pack_vec16( + output, + input, + batch_size, + seq_length, + row_bytes, + action_starts, + action_ends, + llm_offsets, + action_offsets, + stream); + return; + } + + const int blocks = checked_row_blocks(batch_size, seq_length); + const int threads = row_copy_threads(hidden_dim); + twoExpertPackRowsKernel<<>>( + output, + input, + batch_size, + seq_length, + hidden_dim, + action_starts, + action_ends, + llm_offsets, + action_offsets); +} + +template +static void launch_unpack( + T* output, + const T* input, + int64_t batch_size, + int64_t seq_length, + int64_t hidden_dim, + const int64_t* action_starts, + const int64_t* action_ends, + const int64_t* llm_offsets, + const int64_t* action_offsets, + cudaStream_t stream) { + const int64_t row_bytes = hidden_dim * static_cast(sizeof(T)); + if (can_use_vec16(output, input, row_bytes)) { + launch_unpack_vec16( + output, + input, + batch_size, + seq_length, + row_bytes, + action_starts, + action_ends, + llm_offsets, + action_offsets, + stream); + return; + } + + const int blocks = checked_row_blocks(batch_size, seq_length); + const int threads = row_copy_threads(hidden_dim); + twoExpertUnpackRowsKernel<<>>( + output, + input, + batch_size, + seq_length, + hidden_dim, + action_starts, + action_ends, + llm_offsets, + action_offsets); +} + +void TwoExpertPack( + const at::Tensor& output, + const at::Tensor& input, + const at::Tensor& action_starts, + const at::Tensor& action_ends, + const at::Tensor& llm_offsets, + const at::Tensor& action_offsets) { + ASSERT_CHECK(input.dim() == 3); + ASSERT_CHECK(output.dim() == 2); + ASSERT_CHECK(input.is_contiguous()); + ASSERT_CHECK(output.is_contiguous()); + + const int64_t batch_size = input.size(0); + const int64_t seq_length = input.size(1); + const int64_t hidden_dim = input.size(2); + check_index_tensor( + action_starts, "action_starts", batch_size, input.device().type(), input.device().index()); + check_index_tensor( + action_ends, "action_ends", batch_size, input.device().type(), input.device().index()); + check_index_tensor( + llm_offsets, "llm_offsets", batch_size + 1, input.device().type(), input.device().index()); + check_index_tensor( + action_offsets, "action_offsets", batch_size + 1, input.device().type(), input.device().index()); + ASSERT_CHECK(output.size(0) == batch_size * seq_length); + ASSERT_CHECK(output.size(1) == hidden_dim); + ASSERT_CHECK(output.scalar_type() == input.scalar_type()); + ASSERT_CHECK(output.element_size() == input.element_size()); + ASSERT_CHECK(output.device().type() == input.device().type()); + ASSERT_CHECK(output.device().index() == input.device().index()); + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + const int64_t* action_starts_ptr = static_cast(action_starts.data_ptr()); + const int64_t* action_ends_ptr = static_cast(action_ends.data_ptr()); + const int64_t* llm_offsets_ptr = static_cast(llm_offsets.data_ptr()); + const int64_t* action_offsets_ptr = static_cast(action_offsets.data_ptr()); + + if (input.scalar_type() == at::kFloat) { + launch_pack( + static_cast(output.data_ptr()), + static_cast(input.data_ptr()), + batch_size, + seq_length, + hidden_dim, + action_starts_ptr, + action_ends_ptr, + llm_offsets_ptr, + action_offsets_ptr, + stream); + } else if (input.scalar_type() == at::kHalf) { + launch_pack<__half>( + static_cast<__half*>(output.data_ptr()), + static_cast(input.data_ptr()), + batch_size, + seq_length, + hidden_dim, + action_starts_ptr, + action_ends_ptr, + llm_offsets_ptr, + action_offsets_ptr, + stream); + } else if (input.scalar_type() == at::kBFloat16) { + launch_pack<__nv_bfloat16>( + static_cast<__nv_bfloat16*>(output.data_ptr()), + static_cast(input.data_ptr()), + batch_size, + seq_length, + hidden_dim, + action_starts_ptr, + action_ends_ptr, + llm_offsets_ptr, + action_offsets_ptr, + stream); + } else { + throw std::runtime_error("two_expert_pack: unsupported dtype"); + } + sync_check_cuda_error(); +} + +void TwoExpertUnpack( + const at::Tensor& output, + const at::Tensor& input, + int64_t batch_size, + int64_t seq_length, + const at::Tensor& action_starts, + const at::Tensor& action_ends, + const at::Tensor& llm_offsets, + const at::Tensor& action_offsets) { + ASSERT_CHECK(input.dim() == 2); + ASSERT_CHECK(output.dim() == 3); + ASSERT_CHECK(input.is_contiguous()); + ASSERT_CHECK(output.is_contiguous()); + + const int64_t hidden_dim = input.size(1); + check_index_tensor( + action_starts, "action_starts", batch_size, input.device().type(), input.device().index()); + check_index_tensor( + action_ends, "action_ends", batch_size, input.device().type(), input.device().index()); + check_index_tensor( + llm_offsets, "llm_offsets", batch_size + 1, input.device().type(), input.device().index()); + check_index_tensor( + action_offsets, "action_offsets", batch_size + 1, input.device().type(), input.device().index()); + ASSERT_CHECK(input.size(0) == batch_size * seq_length); + ASSERT_CHECK(output.size(0) == batch_size); + ASSERT_CHECK(output.size(1) == seq_length); + ASSERT_CHECK(output.size(2) == hidden_dim); + ASSERT_CHECK(output.scalar_type() == input.scalar_type()); + ASSERT_CHECK(output.element_size() == input.element_size()); + ASSERT_CHECK(output.device().type() == input.device().type()); + ASSERT_CHECK(output.device().index() == input.device().index()); + + cudaStream_t stream = static_cast( + at::cuda::getCurrentCUDAStream().stream()); + + const int64_t* action_starts_ptr = static_cast(action_starts.data_ptr()); + const int64_t* action_ends_ptr = static_cast(action_ends.data_ptr()); + const int64_t* llm_offsets_ptr = static_cast(llm_offsets.data_ptr()); + const int64_t* action_offsets_ptr = static_cast(action_offsets.data_ptr()); + + if (input.scalar_type() == at::kFloat) { + launch_unpack( + static_cast(output.data_ptr()), + static_cast(input.data_ptr()), + batch_size, + seq_length, + hidden_dim, + action_starts_ptr, + action_ends_ptr, + llm_offsets_ptr, + action_offsets_ptr, + stream); + } else if (input.scalar_type() == at::kHalf) { + launch_unpack<__half>( + static_cast<__half*>(output.data_ptr()), + static_cast(input.data_ptr()), + batch_size, + seq_length, + hidden_dim, + action_starts_ptr, + action_ends_ptr, + llm_offsets_ptr, + action_offsets_ptr, + stream); + } else if (input.scalar_type() == at::kBFloat16) { + launch_unpack<__nv_bfloat16>( + static_cast<__nv_bfloat16*>(output.data_ptr()), + static_cast(input.data_ptr()), + batch_size, + seq_length, + hidden_dim, + action_starts_ptr, + action_ends_ptr, + llm_offsets_ptr, + action_offsets_ptr, + stream); + } else { + throw std::runtime_error("two_expert_unpack: unsupported dtype"); + } + sync_check_cuda_error(); +} + +// Removed TVM FFI export (see pybind11 registration) +// Removed TVM FFI export (see pybind11 registration) + +} // namespace wallx_cuda_two_expert_scatter diff --git a/wall_x/model/core/ops/gated_delta.py b/wall_x/model/core/ops/gated_delta.py new file mode 100644 index 00000000..2f6d165e --- /dev/null +++ b/wall_x/model/core/ops/gated_delta.py @@ -0,0 +1,137 @@ +"""GatedDeltaNet post-conv qkv pack / preprocess with backend dispatch. + +Model code calls :func:`pack_gated_delta_qkv` / :func:`preprocess_gated_delta_inputs` +and stays oblivious to which backend runs. Backend selection is capability-based +and lives entirely here: the fused external_accel kernels are used when the loaded +build provides them (and inputs are CUDA), otherwise the pure-PyTorch +``*_pytorch`` fallbacks in this module run (also what the open-source export +ships and what the bitexact test exercises). + +Note on qk L2 norm: the fused *preprocess* kernel applies the query/key L2 norm +itself and reports ``use_qk_l2norm_in_kernel=False`` so the downstream delta-rule +kernel skips it; the PyTorch path defers the L2 norm to the delta-rule kernel and +reports ``True``. So for equal downstream results, +``l2norm(pytorch_query) == kernel_query`` (and likewise for key). +""" + +import torch +import torch.nn.functional as F + +try: + from external_accel import gated_delta_post_conv_pack as _external_accel_pack +except Exception: + _external_accel_pack = None + +try: + from external_accel import ( + gated_delta_post_conv_pack_preprocess as _external_accel_preprocess, + ) +except Exception: + _external_accel_preprocess = None + + +def gated_delta_pack_available(reference_tensor: torch.Tensor) -> bool: + """True iff the fused pack kernel can run for tensors like *reference_tensor*.""" + return _external_accel_pack is not None and reference_tensor.is_cuda + + +def gated_delta_preprocess_available(reference_tensor: torch.Tensor) -> bool: + """True iff the fused preprocess kernel can run for *reference_tensor*.""" + return _external_accel_preprocess is not None and reference_tensor.is_cuda + + +def pack_gated_delta_qkv_pytorch( + module, + mixed_qkv: torch.Tensor, + batch_size: int, + seq_len: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pure-PyTorch split of post-conv mixed_qkv into per-head q/k/v.""" + mixed_qkv = mixed_qkv.transpose(1, 2) + query, key, value = torch.split( + mixed_qkv, + [ + module.key_dim, + module.key_dim, + module.value_dim, + ], + dim=-1, + ) + + query = query.reshape(batch_size, seq_len, -1, module.head_k_dim) + key = key.reshape(batch_size, seq_len, -1, module.head_k_dim) + value = value.reshape(batch_size, seq_len, -1, module.head_v_dim) + if module.num_v_heads // module.num_k_heads > 1: + query = query.repeat_interleave(module.num_v_heads // module.num_k_heads, dim=2) + key = key.repeat_interleave(module.num_v_heads // module.num_k_heads, dim=2) + return query, key, value + + +def preprocess_gated_delta_inputs_pytorch( + module, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + batch_size: int, + seq_len: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool]: + """Pure-PyTorch pack + beta/g gates; defers qk L2 norm (returns ``True``).""" + query, key, value = pack_gated_delta_qkv_pytorch( + module, mixed_qkv, batch_size, seq_len + ) + beta = b.sigmoid() + # If the model is loaded in fp16, without the .float() here, A might be -inf + g = -module.A_log.float().exp() * F.softplus(a.float() + module.dt_bias) + return query, key, value, beta, g, True + + +def pack_gated_delta_qkv( + module, + mixed_qkv: torch.Tensor, + batch_size: int, + seq_len: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Split post-conv mixed_qkv into per-head q/k/v (fused kernel or PyTorch).""" + if _external_accel_pack is not None and mixed_qkv.is_cuda: + return _external_accel_pack( + mixed_qkv, + module.num_k_heads, + module.num_v_heads, + module.head_k_dim, + module.head_v_dim, + ) + return pack_gated_delta_qkv_pytorch(module, mixed_qkv, batch_size, seq_len) + + +def preprocess_gated_delta_inputs( + module, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + batch_size: int, + seq_len: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool]: + """Pack q/k/v and compute beta/g gates (fused kernel or PyTorch). + + Returns ``(query, key, value, beta, g, use_qk_l2norm_in_kernel)``: the fused + kernel already applies the qk L2 norm, so downstream must skip it + (``use_qk_l2norm_in_kernel=False``); the PyTorch path leaves it to the + delta-rule kernel (``True``). + """ + if _external_accel_preprocess is not None and mixed_qkv.is_cuda: + query, key, value, beta, g = _external_accel_preprocess( + mixed_qkv, + b, + a, + module.A_log.float().contiguous(), + module.dt_bias.float().contiguous(), + module.num_k_heads, + module.num_v_heads, + module.head_k_dim, + module.head_v_dim, + ) + return query, key, value, beta, g, False + + return preprocess_gated_delta_inputs_pytorch( + module, mixed_qkv, b, a, batch_size, seq_len + ) diff --git a/wall_x/model/core/ops/index.py b/wall_x/model/core/ops/index.py index 1323d492..fa5e540a 100644 --- a/wall_x/model/core/ops/index.py +++ b/wall_x/model/core/ops/index.py @@ -5,7 +5,6 @@ import torch import torch.nn.functional as F - from wall_x.model.core.ops.base import OpsProxy logger = logging.getLogger(__name__) diff --git a/wall_x/model/core/ops/moe.py b/wall_x/model/core/ops/moe.py index c50d946c..98083f2b 100644 --- a/wall_x/model/core/ops/moe.py +++ b/wall_x/model/core/ops/moe.py @@ -7,10 +7,23 @@ import logging import torch - +from wall_x.model.core.action.layout import TwoExpertLayout from wall_x.model.core.ops.base import OpsProxy logger = logging.getLogger(__name__) +_SLOW_TWO_EXPERT_FALLBACK_WARNED = set() + + +def _warn_slow_two_expert_fallback(op_name: str) -> None: + if op_name in _SLOW_TWO_EXPERT_FALLBACK_WARNED: + return + _SLOW_TWO_EXPERT_FALLBACK_WARNED.add(op_name) + logger.warning( + "%s: using PyTorch fallback with per-batch Python loops. " + "This path is much slower than the external_accel CUDA kernel; install/build " + "external_accel for production training.", + op_name, + ) class PermuteOp(OpsProxy): @@ -95,3 +108,213 @@ def _pytorch_fallback(self, permuted_tokens, sorted_indices, probs=None): permute = PermuteOp() unpermute = UnpermuteOp() + + +class TwoExpertPackOp(OpsProxy): + """Pack fixed two-expert VLA tokens without sorting. + + Input shape is [B, S, D]. Tokens in [action_starts[b], action_ends[b]) + are routed to expert 1/action; all other tokens are routed to expert 0/LLM. Output + shape is [B*S, D], ordered as [all LLM tokens][all action tokens]. + """ + + @property + def _external_accel_name(self): + return "two_expert_pack" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import two_expert_pack_kernel + + return two_expert_pack_kernel + except ImportError: + return None + except Exception as e: + logger.warning("TwoExpertPackOp: CUDA kernel load failed: %s", e) + return None + + def __call__( + self, + tokens, + layout_or_action_starts, + action_ends=None, + llm_offsets=None, + action_offsets=None, + ): + if isinstance(layout_or_action_starts, TwoExpertLayout): + layout = layout_or_action_starts + return super().__call__( + tokens, + layout.action_starts, + layout.action_ends, + layout.llm_offsets, + layout.action_offsets, + ) + if llm_offsets is None or action_offsets is None: + raise RuntimeError("two_expert_pack requires precomputed layout offsets") + return super().__call__( + tokens, + layout_or_action_starts, + action_ends, + llm_offsets, + action_offsets, + ) + + def _pytorch_fallback( + self, tokens, action_starts, action_ends, llm_offsets, action_offsets + ): + _warn_slow_two_expert_fallback("two_expert_pack") + if tokens.dim() != 3: + raise RuntimeError("two_expert_pack expects input shape [B, S, D]") + bsz, seq_length, hidden_dim = tokens.shape + action_starts = torch.as_tensor( + action_starts, dtype=torch.int64, device=tokens.device + ).reshape(-1) + action_ends = torch.as_tensor( + action_ends, dtype=torch.int64, device=tokens.device + ).reshape(-1) + if action_starts.numel() != bsz or action_ends.numel() != bsz: + raise RuntimeError("two_expert_pack expects [B] action ranges") + if llm_offsets.numel() != bsz + 1 or action_offsets.numel() != bsz + 1: + raise RuntimeError("two_expert_pack expects [B+1] layout offsets") + + llm_tokens = [] + action_tokens = [] + for batch in range(bsz): + start = int(action_starts[batch].item()) + end = int(action_ends[batch].item()) + if start < 0 or end < start or end > seq_length: + raise RuntimeError( + f"invalid action range [{start}, {end}) for seq_length={seq_length}" + ) + llm_tokens.append( + torch.cat((tokens[batch, :start], tokens[batch, end:]), dim=0) + ) + action_tokens.append(tokens[batch, start:end]) + + return torch.cat( + ( + torch.cat(llm_tokens, dim=0).reshape(-1, hidden_dim), + torch.cat(action_tokens, dim=0).reshape(-1, hidden_dim), + ), + dim=0, + ) + + +class TwoExpertUnpackOp(OpsProxy): + """Restore fixed two-expert packed tokens back to [B, S, D].""" + + @property + def _external_accel_name(self): + return "two_expert_unpack" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import two_expert_unpack_kernel + + return two_expert_unpack_kernel + except ImportError: + return None + except Exception as e: + logger.warning("TwoExpertUnpackOp: CUDA kernel load failed: %s", e) + return None + + def __call__( + self, + packed_tokens, + batch_size: int, + seq_length: int, + layout_or_action_starts, + action_ends=None, + llm_offsets=None, + action_offsets=None, + ): + if isinstance(layout_or_action_starts, TwoExpertLayout): + layout = layout_or_action_starts + return super().__call__( + packed_tokens, + batch_size, + seq_length, + layout.action_starts, + layout.action_ends, + layout.llm_offsets, + layout.action_offsets, + ) + if llm_offsets is None or action_offsets is None: + raise RuntimeError("two_expert_unpack requires precomputed layout offsets") + return super().__call__( + packed_tokens, + batch_size, + seq_length, + layout_or_action_starts, + action_ends, + llm_offsets, + action_offsets, + ) + + def _pytorch_fallback( + self, + packed_tokens, + batch_size: int, + seq_length: int, + action_starts, + action_ends, + llm_offsets, + action_offsets, + ): + _warn_slow_two_expert_fallback("two_expert_unpack") + if packed_tokens.dim() != 2: + raise RuntimeError("two_expert_unpack expects input shape [B*S, D]") + batch_size = int(batch_size) + seq_length = int(seq_length) + if packed_tokens.size(0) != batch_size * seq_length: + raise RuntimeError( + f"packed rows mismatch: got {packed_tokens.size(0)}, expected {batch_size * seq_length}" + ) + + hidden_dim = packed_tokens.size(1) + action_starts = torch.as_tensor( + action_starts, dtype=torch.int64, device=packed_tokens.device + ).reshape(-1) + action_ends = torch.as_tensor( + action_ends, dtype=torch.int64, device=packed_tokens.device + ).reshape(-1) + if action_starts.numel() != batch_size or action_ends.numel() != batch_size: + raise RuntimeError("two_expert_unpack expects [B] action ranges") + if llm_offsets.numel() != batch_size + 1 or action_offsets.numel() != ( + batch_size + 1 + ): + raise RuntimeError("two_expert_unpack expects [B+1] layout offsets") + llm_offsets = llm_offsets.to(device=packed_tokens.device, dtype=torch.int64) + action_offsets = action_offsets.to( + device=packed_tokens.device, dtype=torch.int64 + ) + total_llm = int(llm_offsets[-1].item()) + + output = torch.empty( + (batch_size, seq_length, hidden_dim), + dtype=packed_tokens.dtype, + device=packed_tokens.device, + ) + for batch in range(batch_size): + start = int(action_starts[batch].item()) + end = int(action_ends[batch].item()) + if start < 0 or end < start or end > seq_length: + raise RuntimeError( + f"invalid action range [{start}, {end}) for seq_length={seq_length}" + ) + llm_start = int(llm_offsets[batch].item()) + llm_count = int(llm_offsets[batch + 1].item()) - llm_start + action_start = total_llm + int(action_offsets[batch].item()) + llm_tokens = packed_tokens[llm_start : llm_start + llm_count] + action_tokens = packed_tokens[action_start : action_start + (end - start)] + if start > 0: + output[batch, :start, :] = llm_tokens[:start] + output[batch, start:end, :] = action_tokens + if end < seq_length: + output[batch, end:, :] = llm_tokens[start:] + return output + + +two_expert_pack = TwoExpertPackOp() +two_expert_unpack = TwoExpertUnpackOp() diff --git a/wall_x/model/core/ops/norm.py b/wall_x/model/core/ops/norm.py index f410540a..574fa43a 100644 --- a/wall_x/model/core/ops/norm.py +++ b/wall_x/model/core/ops/norm.py @@ -1,9 +1,10 @@ -"""RMS normalization operator.""" +"""Normalization operators (RMSNorm + LayerNorm).""" +import inspect import logging import torch - +import torch.nn.functional as F from wall_x.model.core.ops.base import OpsProxy logger = logging.getLogger(__name__) @@ -19,6 +20,39 @@ class RMSNormOp(OpsProxy): def _external_accel_name(self): return "rmsnorm" + def _probe_external_accel(self): + fn = super()._probe_external_accel() + if fn is None: + return None + + try: + supports_zero_centered = "zero_centered" in inspect.signature(fn).parameters + except (TypeError, ValueError): + supports_zero_centered = True + + if supports_zero_centered: + return fn + + try: + _supports_inference = "inference" in inspect.signature(fn).parameters + except (TypeError, ValueError): + _supports_inference = True + + def legacy_rmsnorm( + hidden_states, weight, eps=1e-6, inference=False, zero_centered=False + ): + if zero_centered: + return self._pytorch_fallback( + hidden_states, weight, eps, zero_centered=True + ) + # Some legacy external_accel rmsnorm kernels predate the `inference` kwarg + # (added alongside zero_centered); only pass it when accepted. + if _supports_inference: + return fn(hidden_states, weight, eps, inference=inference) + return fn(hidden_states, weight, eps) + + return legacy_rmsnorm + def _get_cuda_kernel(self): try: from wall_x.model.core.ops._cuda_wrappers import Rmsnorm @@ -30,13 +64,97 @@ def _get_cuda_kernel(self): logger.warning("RMSNormOp: CUDA kernel load failed: %s", e) return None - def _pytorch_fallback(self, hidden_states, weight, eps=1e-6): + def _pytorch_fallback(self, hidden_states, weight, eps=1e-6, zero_centered=False): """Pure PyTorch RMSNorm. eps default matches external_accel (1e-6, not PyTorch's 1e-5).""" input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + eps) + if zero_centered: + weight = weight + 1.0 return (weight * hidden_states).to(input_dtype) +class AddRMSNormOp(OpsProxy): + """Fused residual-add + RMSNorm. + + add_rmsnorm(residual, input, weight, eps, zero_centered=False) + -> (sum_output, norm_output) + where sum_output = residual + input (materialized, reused by the next + residual add) and norm_output = RMSNorm(sum_output) * weight_effective. + The external_accel backward reuses the forward's saved rstd (fast path for + hidden_size <= 4096). See docs qwen35_mfu_report (AddRMSNorm fast backward). + """ + + @property + def _external_accel_name(self): + return "add_rmsnorm" + + def _pytorch_fallback( + self, residual, input, weight, eps=1e-6, inference=False, zero_centered=False + ): + sum_output = residual + input + input_dtype = sum_output.dtype + x = sum_output.to(torch.float32) + variance = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + eps) + if zero_centered: + weight = weight + 1.0 + norm_output = (weight * x).to(input_dtype) + return sum_output, norm_output + + +class GatedRMSNormOp(OpsProxy): + """Gated RMSNorm used at the end of GatedDeltaNet: + + gated_rmsnorm(input, gate, weight, eps) + = (RMSNorm(input) * weight) * silu(gate) + + Prefers the external_accel fused kernel when the loaded build provides it + forces the PyTorch fallback below. + See docs qwen35_mfu_report (GatedDeltaNet gated RMSNorm via external_accel). + """ + + @property + def _external_accel_name(self): + return "gated_rmsnorm" + + def _pytorch_fallback(self, input, gate, weight, eps=1e-6, inference=False): + input_dtype = input.dtype + hidden_states = input.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + eps) + hidden_states = weight * hidden_states.to(input_dtype) + hidden_states = hidden_states * torch.nn.functional.silu(gate.to(torch.float32)) + return hidden_states.to(input_dtype) + + rmsnorm = RMSNormOp() +add_rmsnorm = AddRMSNormOp() +gated_rmsnorm = GatedRMSNormOp() + + +class LayerNormOp(OpsProxy): + """Layer normalization: layernorm(input, weight, bias, eps). + + external_accel signature (external_accel/kernel/layernorm.py): + layernorm(input_tensor, weight, bias, epsilon) -> Tensor + bias is required (external_accel kernel always applies affine). + + Currently unused (SigLIP path switched to flash_attn.ops.triton.layer_norm. + layer_norm_fn behind PI0Config.use_fused_layernorm in session-17 after the + external_accel kernel was shown to produce wrong output scale). Kept as a + reference for the proxy pattern; safe to remove if no other consumer + appears. + """ + + @property + def _external_accel_name(self): + return "layernorm" + + def _pytorch_fallback(self, input_tensor, weight, bias, eps): + # F.layer_norm normalizes over weight.shape (the trailing dims). + return F.layer_norm(input_tensor, weight.shape, weight, bias, eps) + + +layernorm = LayerNormOp() diff --git a/wall_x/model/core/ops/position.py b/wall_x/model/core/ops/position.py new file mode 100644 index 00000000..b06c96f0 --- /dev/null +++ b/wall_x/model/core/ops/position.py @@ -0,0 +1,96 @@ +"""Position embedding operators.""" + +import torch + +from wall_x.model.core.ops.base import OpsProxy + + +class PosEmbedInterpOp(OpsProxy): + """Fused Qwen3.5 ViT position embedding interpolation. + + Signature: + pos_embed_interp(grid_thw, pos_embed_weight, num_grid_per_side, + spatial_merge_size, total_output_tokens=None) + """ + + @property + def _external_accel_name(self): + return "pos_embed_interp" + + def _pytorch_fallback( + self, + grid_thw, + pos_embed_weight, + num_grid_per_side: int, + spatial_merge_size: int, + total_output_tokens=None, + ): + # total_output_tokens lets the external_accel backend skip a GPU->CPU sync; + # the PyTorch fallback derives sizes from grid_thw directly. + del total_output_tokens + pos_embed_weight = pos_embed_weight.float() + grid_thw_list = grid_thw.tolist() + grid_ts = [row[0] for row in grid_thw_list] + grid_hs = [row[1] for row in grid_thw_list] + grid_ws = [row[2] for row in grid_thw_list] + device = pos_embed_weight.device + + idx_list = [[] for _ in range(4)] + weight_list = [[] for _ in range(4)] + + for _t, h, w in grid_thw_list: + h_idxs = torch.linspace(0, num_grid_per_side - 1, h) + w_idxs = torch.linspace(0, num_grid_per_side - 1, w) + + h_idxs_floor = h_idxs.int() + w_idxs_floor = w_idxs.int() + h_idxs_ceil = (h_idxs.int() + 1).clip(max=num_grid_per_side - 1) + w_idxs_ceil = (w_idxs.int() + 1).clip(max=num_grid_per_side - 1) + + dh = h_idxs - h_idxs_floor + dw = w_idxs - w_idxs_floor + + base_h = h_idxs_floor * num_grid_per_side + base_h_ceil = h_idxs_ceil * num_grid_per_side + + indices = [ + (base_h[None].T + w_idxs_floor[None]).flatten(), + (base_h[None].T + w_idxs_ceil[None]).flatten(), + (base_h_ceil[None].T + w_idxs_floor[None]).flatten(), + (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(), + ] + weights = [ + ((1 - dh)[None].T * (1 - dw)[None]).flatten(), + ((1 - dh)[None].T * dw[None]).flatten(), + (dh[None].T * (1 - dw)[None]).flatten(), + (dh[None].T * dw[None]).flatten(), + ] + for i in range(4): + idx_list[i].extend(indices[i].tolist()) + weight_list[i].extend(weights[i].tolist()) + + idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=device) + weight_tensor = torch.tensor(weight_list, dtype=torch.float32, device=device) + pos_embeds = pos_embed_weight[idx_tensor] * weight_tensor[:, :, None] + patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3] + + patch_pos_embeds = patch_pos_embeds.split( + [h * w for h, w in zip(grid_hs, grid_ws)] + ) + + patch_pos_embeds_permute = [] + merge_size = spatial_merge_size + for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws): + pos_embed = pos_embed.repeat(t, 1) + pos_embed = ( + pos_embed.view( + t, h // merge_size, merge_size, w // merge_size, merge_size, -1 + ) + .permute(0, 1, 3, 2, 4, 5) + .flatten(0, 4) + ) + patch_pos_embeds_permute.append(pos_embed) + return torch.cat(patch_pos_embeds_permute) + + +pos_embed_interp = PosEmbedInterpOp() diff --git a/wall_x/model/core/ops/qkv_norm_replace.py b/wall_x/model/core/ops/qkv_norm_replace.py new file mode 100644 index 00000000..599660ad --- /dev/null +++ b/wall_x/model/core/ops/qkv_norm_replace.py @@ -0,0 +1,121 @@ +"""Fused Qwen3.5 VLA qkv norm-replace operator (kernel + PyTorch fallback). + +One op splits a fused text/action qkv tensor into q/gate/k/v, applies the +zero-centered q/k RMSNorm (raw norm weights; the `+1` is added here), and +replaces the action-token rows of each batch with the action expert's values, +emitting a FlashMask-friendly layout. + +``qkv_norm_replace`` dispatches by capability: the fused external_accel kernel when it +is available and the inputs are on CUDA, otherwise the self-contained PyTorch +fallback in this module. The fallback is the code path the open-source export +runs (it ships without external_accel) and is what the bitexact test exercises. + +Input layout (matches the kernel, per external_accel ``_check_inputs``): +``text_qkv`` / ``action_qkv`` are ``[B, S, C]`` with +``C = (num_heads*2 + num_key_value_heads*2) * head_dim`` laid out as +``[q|gate head-interleaved (H*2D)] [key (Hkv*D)] [value (Hkv*D)]``; within each +head's ``2*head_dim`` slice the first ``head_dim`` is query, the next is gate. +Action tokens are contiguous rows ``[action_starts[b]:action_ends[b])`` per +batch; ``action_offsets`` is the ``[B+1]`` cumsum mapping each batch's rows into +the flattened action buffer. + +Returns ``(query, key, value, gate)`` with shapes +``[B,S,H,D] / [B,S,Hkv,D] / [B,S,Hkv,D] / [B,S,H*D]``. +""" + +import torch + +from wall_x.model.core.ops.norm import rmsnorm as _rmsnorm + +try: + from external_accel import qwen35_vla_qkv_norm_replace_kernel as _kernel +except Exception: + _kernel = None + + +def qkv_norm_replace_available(reference_tensor: torch.Tensor) -> bool: + """True iff the fused kernel can run for tensors like *reference_tensor*.""" + return _kernel is not None and reference_tensor.is_cuda + + +def _zero_centered_rmsnorm( + x: torch.Tensor, weight: torch.Tensor, eps: float +) -> torch.Tensor: + # Same math as Qwen3_5RMSNorm: float32 upcast, (1 + weight) scale, cast back. + return _rmsnorm._pytorch_fallback(x, weight.float(), eps, zero_centered=True) + + +def qkv_norm_replace_pytorch( + text_qkv: torch.Tensor, + action_qkv: torch.Tensor, + action_starts: torch.Tensor, + action_ends: torch.Tensor, + action_offsets: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + q_action_weight: torch.Tensor, + k_action_weight: torch.Tensor, + num_heads: int, + num_key_value_heads: int, + head_dim: int, + epsilon: float, +): + """Pure-tensor equivalent of the fused kernel (see module docstring).""" + H, Hkv, D = int(num_heads), int(num_key_value_heads), int(head_dim) + q_block = H * 2 * D + k_block = Hkv * D + + def split_norm(qkv, qw, kw): + bsz, seqlen, _ = qkv.shape + qg = qkv[..., :q_block].reshape(bsz, seqlen, H, 2 * D) + query = _zero_centered_rmsnorm(qg[..., :D].contiguous(), qw, epsilon) + gate = qg[..., D:].reshape(bsz, seqlen, H * D) + key = _zero_centered_rmsnorm( + qkv[..., q_block : q_block + k_block] + .reshape(bsz, seqlen, Hkv, D) + .contiguous(), + kw, + epsilon, + ) + value = qkv[..., q_block + k_block :].reshape(bsz, seqlen, Hkv, D).contiguous() + return query, key, value, gate + + query, key, value, gate = split_norm(text_qkv, q_weight, k_weight) + B = query.shape[0] + + a_query, a_key, a_value, a_gate = split_norm( + action_qkv, q_action_weight, k_action_weight + ) + total_action = action_qkv.shape[0] * action_qkv.shape[1] + a_query = a_query.reshape(total_action, H, D) + a_key = a_key.reshape(total_action, Hkv, D) + a_value = a_value.reshape(total_action, Hkv, D) + a_gate = a_gate.reshape(total_action, H * D) + + # Replace each batch's contiguous action rows with the action expert's, + # indexing the flattened action buffer via action_offsets (matches kernel). + query, key, value, gate = query.clone(), key.clone(), value.clone(), gate.clone() + starts = action_starts.tolist() + ends = action_ends.tolist() + offsets = action_offsets.tolist() + for b in range(B): + s, e = int(starts[b]), int(ends[b]) + o0, o1 = int(offsets[b]), int(offsets[b + 1]) + if e > s: + query[b, s:e] = a_query[o0:o1] + key[b, s:e] = a_key[o0:o1] + value[b, s:e] = a_value[o0:o1] + gate[b, s:e] = a_gate[o0:o1] + return query, key, value, gate + + +def qkv_norm_replace(*args, **kwargs): + """Capability-dispatched fused qkv norm-replace. + + Uses the external_accel kernel when it is available and the fused qkv tensor is on + CUDA; otherwise runs :func:`qkv_norm_replace_pytorch`. + """ + text_qkv = args[0] if args else kwargs["text_qkv"] + if _kernel is not None and text_qkv.is_cuda: + return _kernel(*args, **kwargs) + return qkv_norm_replace_pytorch(*args, **kwargs) diff --git a/wall_x/model/core/ops/qwen35_index.py b/wall_x/model/core/ops/qwen35_index.py new file mode 100644 index 00000000..7b5c32cc --- /dev/null +++ b/wall_x/model/core/ops/qwen35_index.py @@ -0,0 +1,126 @@ +"""Qwen3.5-specific index computation operators.""" + +import itertools +from typing import Optional, Tuple + +import torch + +from wall_x.model.core.ops.base import OpsProxy + + +class Qwen35GetRopeIndexOp(OpsProxy): + """Compute Qwen3.5 3D RoPE position indices from modality-type groups. + + Signature: + qwen35_get_rope_index(input_ids, mm_token_type_ids, image_grid_thw, + video_grid_thw, attention_mask, spatial_merge_size) + -> (position_ids, mrope_position_deltas) + """ + + @property + def _external_accel_name(self): + return "qwen35_get_rope_index" + + def _pytorch_fallback( + self, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + spatial_merge_size: int = 1, + ) -> Tuple[torch.Tensor, torch.Tensor]: + position_ids = torch.zeros( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=torch.int64, + device=input_ids.device, + ) + mrope_position_deltas = [] + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + valid_mask = attention_mask[batch_idx].bool() + current_input_ids = current_input_ids[valid_mask] + input_token_type = input_token_type[valid_mask] + + input_type_group = [] + for key, group in itertools.groupby( + enumerate(input_token_type.tolist()), lambda x: x[1] + ): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device) + .view(1, -1) + .expand(3, -1) + + current_pos + ) + current_pos += text_len + else: + grid_iter = grid_iters[modality_type] + if grid_iter is None: + raise ValueError( + f"missing grid_thw for modality_type={modality_type}" + ) + grid_thw = next(grid_iter) + llm_grid_t = grid_thw[0].item() + llm_grid_h = grid_thw[1].item() // spatial_merge_size + llm_grid_w = grid_thw[2].item() // spatial_merge_size + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange( + current_pos, + current_pos + llm_grid_w, + device=input_ids.device, + ).repeat(llm_grid_h * llm_grid_t) + position_height = torch.arange( + current_pos, + current_pos + llm_grid_h, + device=input_ids.device, + ).repeat_interleave(llm_grid_w * llm_grid_t) + position_temporal = torch.full( + (image_seq_length,), + current_pos, + device=input_ids.device, + dtype=torch.long, + ) + llm_pos_ids_list.append( + torch.stack( + [position_temporal, position_height, position_width], + dim=0, + ) + ) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = ( + llm_positions.to(position_ids.device) + ) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(current_input_ids) + ) + + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=input_ids.device + ).unsqueeze(1) + return position_ids, mrope_position_deltas + + +qwen35_get_rope_index = Qwen35GetRopeIndexOp() diff --git a/wall_x/model/core/ops/rope.py b/wall_x/model/core/ops/rope.py index 09f6f42d..019483fb 100644 --- a/wall_x/model/core/ops/rope.py +++ b/wall_x/model/core/ops/rope.py @@ -4,7 +4,6 @@ from typing import List import torch - from wall_x.model.core.ops.base import OpsProxy logger = logging.getLogger(__name__) @@ -81,13 +80,95 @@ def _pytorch_fallback(self, q, k, cos, sin, interleave=False, **kwargs): k_embed = k_embed.to(k.dtype) return q_embed, k_embed + def full(self, q, k, cos, sin, interleave=False, rotary_dim=None): + """Partial-rotary RoPE that also emits the pass-through part. + + Fuses the rotary-prefix rotation and the pass-through copy into a single + external_accel kernel, avoiding the Python-side ``torch.cat([rotated, pass])``. + q/k are the FULL head tensors ([..., head_dim]); cos/sin are the compact + half-dim (rotary_dim // 2). Falls back to slice + rope + cat when the + external_accel ``.full`` kernel is unavailable. + See docs qwen35_mfu_report (RoPE full fused kernel). + """ + if self._resolved_fn is None: + self._resolve() + # external_accel rope_full fuses the rotary-prefix rotation and pass-through + # copy into one kernel. Verified numerically equal to a pure-PyTorch + # partial-rotary reference (bf16 tolerance). Used whenever the loaded + # backend provides .full. See docs qwen35_mfu_report (RoPE full). + use_external_accel_full = hasattr(self._resolved_fn, "full") + if use_external_accel_full: + try: + return self._resolved_fn.full( + q, k, cos, sin, interleave=interleave, rotary_dim=rotary_dim + ) + except Exception as exc: + logger.warning( + "RoPEOp.full: external_accel path failed (%s); using PyTorch fallback", + exc, + ) + return self._pytorch_full_fallback( + q, k, cos, sin, interleave=interleave, rotary_dim=rotary_dim + ) + + def _pytorch_full_fallback(self, q, k, cos, sin, interleave=False, rotary_dim=None): + if rotary_dim is None: + rotary_dim = cos.size(-1) * 2 + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + q_embed, k_embed = self(q_rot, k_rot, cos, sin, interleave=interleave) + q_embed = torch.cat([q_embed, q_pass], dim=-1) + k_embed = torch.cat([k_embed, k_pass], dim=-1) + return q_embed, k_embed + def pack(self, *args, **kwargs): - """Delegate to resolved backend's pack method if available.""" + """Apply RoPE to packed qkv through the selected backend. + + Callers rely on three contracts a bare ``return args`` breaks: some use + the return value directly (qwen3_5 vision), some truth-test it + (qwen2_5 vision), and the infer paths ignore the return and rely on the + in-place q/k writes. The fallback must therefore apply RoPE in place + and return the qkv tensor. + """ if self._resolved_fn is None: self._resolve() if hasattr(self._resolved_fn, "pack"): return self._resolved_fn.pack(*args, **kwargs) - return args + return self._pytorch_pack_fallback(*args, **kwargs) + + def _pytorch_pack_fallback( + self, qkv, q_num_heads, kv_num_heads, cos, sin, interleave=False, **kwargs + ): + if qkv.dim() != 2: + raise ValueError(f"Expected packed qkv to be 2D, got {qkv.dim()}D") + if qkv.stride(-1) != 1: + raise ValueError("Expected packed qkv last dimension to have stride 1") + if cos.dtype != torch.float32 or sin.dtype != torch.float32: + raise ValueError("Expected rope cos/sin to be float32") + if not cos.is_contiguous() or not sin.is_contiguous(): + raise ValueError("Expected rope cos/sin to be contiguous") + + head_dim = cos.size(-1) * 2 + q_dim = q_num_heads * head_dim + kv_dim = kv_num_heads * head_dim + if qkv.size(-1) != q_dim + 2 * kv_dim: + raise ValueError( + f"Expected packed qkv last dim {q_dim + 2 * kv_dim}, got {qkv.size(-1)}" + ) + + q = qkv[..., :q_dim].view(*qkv.shape[:-1], q_num_heads, head_dim) + k = qkv[..., q_dim : q_dim + kv_dim].view( + *qkv.shape[:-1], kv_num_heads, head_dim + ) + + if qkv.dim() == 2 and cos.dim() == 3 and cos.shape[0] == 1: + cos = cos.squeeze(0) + sin = sin.squeeze(0) + + q_embed, k_embed = self._pytorch_fallback(q, k, cos, sin, interleave=interleave) + q.copy_(q_embed) + k.copy_(k_embed) + return qkv class MRoPEOp(OpsProxy): @@ -151,12 +232,43 @@ def _pytorch_fallback( return q_embed.to(query_states.dtype), k_embed.to(key_states.dtype) def pack(self, *args, **kwargs): - """Delegate to resolved backend's pack method if available.""" + """Apply MRoPE to packed qkv through the selected backend. + + See RoPEOp.pack: the infer kvcache path ignores the return value and + relies on the in-place q/k writes, so the fallback must apply MRoPE in + place rather than return the raw args. + """ if self._resolved_fn is None: self._resolve() if hasattr(self._resolved_fn, "pack"): return self._resolved_fn.pack(*args, **kwargs) - return args + return self._pytorch_pack_fallback(*args, **kwargs) + + def _pytorch_pack_fallback( + self, qkv, q_num_heads, kv_num_heads, cos, sin, mrope_section, **kwargs + ): + if qkv.dim() != 3: + raise ValueError(f"Expected packed qkv to be 3D, got {qkv.dim()}D") + if qkv.stride(-1) != 1: + raise ValueError("Expected packed qkv last dimension to have stride 1") + + head_dim = sum(mrope_section) * 2 + q_dim = q_num_heads * head_dim + kv_dim = kv_num_heads * head_dim + if qkv.size(-1) != q_dim + 2 * kv_dim: + raise ValueError( + f"Expected packed qkv last dim {q_dim + 2 * kv_dim}, got {qkv.size(-1)}" + ) + + q = qkv[..., :q_dim].view(*qkv.shape[:-1], q_num_heads, head_dim) + k = qkv[..., q_dim : q_dim + kv_dim].view( + *qkv.shape[:-1], kv_num_heads, head_dim + ) + + q_embed, k_embed = self._pytorch_fallback(q, k, cos, sin, mrope_section) + q.copy_(q_embed) + k.copy_(k_embed) + return qkv class RotPosEmbOp(OpsProxy): @@ -180,7 +292,12 @@ def _get_cuda_kernel(self): logger.warning("RotPosEmbOp: CUDA kernel load failed: %s", e) return None - def _pytorch_fallback(self, inv_freq, grid_thw, spatial_merge_size): + def _pytorch_fallback( + self, inv_freq, grid_thw, spatial_merge_size, total_output_tokens=None + ): + # total_output_tokens lets the external_accel backend skip a GPU->CPU sync; + # the PyTorch fallback derives sizes from grid_thw directly. + del total_output_tokens if inv_freq.dtype != torch.float32: inv_freq = inv_freq.to(torch.float32) pos_ids = [] diff --git a/wall_x/model/core/vla_mixin.py b/wall_x/model/core/vla_mixin.py index b4385897..75d5bb0e 100644 --- a/wall_x/model/core/vla_mixin.py +++ b/wall_x/model/core/vla_mixin.py @@ -3,24 +3,39 @@ import torch import torch.nn as nn import torch.utils.checkpoint as cp -from peft import LoraConfig, get_peft_model -from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from torch.distributed.fsdp import MixedPrecision as MP +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + + +from peft import LoraConfig, get_peft_model + from transformers import GenerationMixin from transformers.utils import logging -from wall_x.model.core.action.moe import SparseMoeBlock, TokenTypeRouter +from wall_x.utils.constant import is_action_dataset_name + from wall_x.model.core.action.processor import ActionProcessor -from wall_x.model.core.attention.mask import ( - update_joint_attention_flash_mask, - update_joint_attention_mask_2d, +from wall_x.model.core.action.moe import ( + TokenTypeRouter, + SparseMoeBlock, +) +from wall_x.model.core.action.layout import TwoExpertLayout, prepare_two_expert_layout +from wall_x.model.core.ops.attention_mask import ( update_position_ids, + update_joint_attention_mask_2d, + update_joint_attention_flash_mask, ) -from wall_x.utils.constant import is_action_dataset_name logger = logging.get_logger(__name__) +def _adarms_horizon(config, cond): + if isinstance(cond, torch.Tensor) and cond.dim() == 3: + return cond.shape[1] + return config.action_horizon_flow + + class ActionModelMixMin: # config: Qwen2_5_VLConfig action_preprocessor: ActionProcessor @@ -45,10 +60,23 @@ def set_normalizer(self, normalizer_action, normalizer_propri): "ActionModelMixMin.set_normalizer is called but action_preprocessor is not set" ) - def _apply_mlp_moe(self, hidden_states, token_types, start_indices, end_indices): + def _apply_mlp_moe( + self, + hidden_states, + token_types, + action_starts=None, + action_ends=None, + expert_segments=None, + two_expert_layout: TwoExpertLayout = None, + ): if self.config.mlp_moe: hidden_states = self.moe( - hidden_states, token_types, start_indices, end_indices + hidden_states, + token_types, + action_starts, + action_ends, + expert_segments, + two_expert_layout, ) else: hidden_states = self.mlp(hidden_states) @@ -61,8 +89,10 @@ def _apply_norm_moe( adarms_conds, norms, # list of norm layers (expert-wise) norm, # shared norm if not norm_moe - start_indices=None, - end_indices=None, + action_starts=None, + action_ends=None, + expert_segments=None, + two_expert_layout: TwoExpertLayout = None, use_selective_recompute=False, ): """ @@ -81,14 +111,24 @@ def _apply_norm_moe( if self.config.norm_moe: # --------------------------------------------------------- - # Case 1A: mot_opt=True (segments assigned by start/end) + # Case 1A: mot_opt=True (segments assigned by fixed LLM/action layout) # --------------------------------------------------------- if self.config.mot_opt: new_hidden_states = torch.zeros_like(hidden_states) + batch_size, seq_length = token_types.shape + if two_expert_layout is None: + two_expert_layout = prepare_two_expert_layout( + token_types, + batch_size=batch_size, + seq_length=seq_length, + action_starts=action_starts, + action_ends=action_ends, + device=hidden_states.device, + ) + expert_segments = two_expert_layout.expert_segments for expert_idx, expert_norm in enumerate(norms): - start = start_indices[expert_idx] - end = end_indices[expert_idx] + start, end = expert_segments[expert_idx] if start == end: continue @@ -97,9 +137,12 @@ def _apply_norm_moe( # ====== reshape if adarms on flow expert ====== if self.config.use_adarms and expert_idx == 1: + action_horizon = _adarms_horizon( + self.config, adarms_conds[expert_idx] + ) selected = selected.view( -1, - self.config.action_horizon_flow, + action_horizon, selected.shape[-1], ) input_slice = selected[:, :, :dim_input] @@ -160,10 +203,13 @@ def norm_chunk(t_x, t_cond, expert_norm=expert_norm): selected = hidden_states[mask] # [K, D] if self.config.use_adarms and expert_idx == 1: + action_horizon = _adarms_horizon( + self.config, adarms_conds[expert_idx] + ) gate_mask = mask selected = selected.view( -1, - self.config.action_horizon_flow, + action_horizon, selected.shape[-1], ) input_slice = selected[:, :, :dim_input] @@ -231,7 +277,17 @@ def norm_chunk_shared(t_x, dummy, norm_module=norm): return hidden_states, gate, gate_mask - def _gated_residual(self, x, y, gate, start_indices=None, end_indices=None): + def _gated_residual( + self, + x, + y, + gate, + token_types=None, + action_starts=None, + action_ends=None, + expert_segments=None, + two_expert_layout: TwoExpertLayout = None, + ): """ Applies gated residual connection with optional gate parameter. @@ -251,14 +307,47 @@ def _gated_residual(self, x, y, gate, start_indices=None, end_indices=None): return x + y new_y = y.clone() - selected_y = y[start_indices[1] : end_indices[1]] - selected_y = selected_y.view( - -1, self.config.action_horizon_flow, selected_y.shape[-1] - )[:, :, : self.config.dim_inputs[1]] - selected_y = selected_y.to(torch.float32) * gate - new_y[start_indices[1] : end_indices[1], : self.config.dim_inputs[1]] = ( - selected_y.view(-1, self.config.dim_inputs[1]).to(new_y.dtype) + if self.config.mot_opt: + if two_expert_layout is None: + if token_types is None: + raise ValueError( + "token_types or two_expert_layout must be provided for " + "mot_opt gated residual" + ) + batch_size, seq_length = token_types.shape + two_expert_layout = prepare_two_expert_layout( + token_types, + batch_size=batch_size, + seq_length=seq_length, + action_starts=action_starts, + action_ends=action_ends, + device=y.device, + ) + _, action_segment = two_expert_layout.expert_segments + start, end = action_segment + selected_y = y[start:end] + else: + if token_types is None: + raise ValueError("token_types must be provided for gated residual") + action_mask = token_types == 1 + selected_y = y[action_mask] + + # ours (align_cw): dynamic action horizon — use gate's horizon dim when the + # gate is 3-D with horizon>1, else fall back to the static action_horizon_flow. + action_horizon = ( + gate.shape[1] + if isinstance(gate, torch.Tensor) and gate.dim() == 3 and gate.shape[1] > 1 + else self.config.action_horizon_flow ) + selected_y = selected_y.view(-1, action_horizon, selected_y.shape[-1])[ + :, :, : self.config.dim_inputs[1] + ] + selected_y = selected_y.to(torch.float32) * gate + selected_y = selected_y.view(-1, self.config.dim_inputs[1]).to(new_y.dtype) + if self.config.mot_opt: + new_y[start:end, : self.config.dim_inputs[1]] = selected_y + else: + new_y[action_mask, : self.config.dim_inputs[1]] = selected_y return x + new_y @@ -340,7 +429,14 @@ def _update_joint_attention_mask_2d( ) def _update_joint_attention_flash_mask( - self, attention_mask, moe_token_types, positional_masks, debug=False + self, + attention_mask, + moe_token_types, + positional_masks, + debug=False, + batch_size=None, + seq_len=None, + device=None, ): return update_joint_attention_flash_mask( attention_mask, @@ -348,6 +444,9 @@ def _update_joint_attention_flash_mask( positional_masks, causal_action_attention_mask=self.config.causal_action_attention_mask, debug=debug, + batch_size=batch_size, + seq_len=seq_len, + device=device, ) @@ -401,19 +500,19 @@ def _assign_child(root_module, dotted_name: str, new_child): cast_root_forward_inputs=False, ) - from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl import ( - Qwen2_5_VLVisionBlock, - ) from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl_act import ( Qwen2_5_VLDecoderLayer_with_MoE, ) + from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl import ( + Qwen2_5_VLVisionBlock, + ) target_classes = ( Qwen2_5_VLDecoderLayer_with_MoE, Qwen2_5_VLVisionBlock, ) - # Step 1: Find top-level ActionProcessor modules and wrap them separately with FSDP (FP32) + # Step 1️⃣: Find top-level ActionProcessor modules and wrap them separately with FSDP (FP32) for name, module in list(self.named_modules()): if isinstance(module, nn.Module) and any( k in name.lower() @@ -421,6 +520,7 @@ def _assign_child(root_module, dotted_name: str, new_child): "input_layernorm", "post_attention_layernorm", "model.norm", + "pos_embed", "action_preprocessor", ] ): @@ -461,7 +561,7 @@ def _assign_child(root_module, dotted_name: str, new_child): _assign_child(self, name, wrapped) setattr(wrapped, "_fsdp_wrapped", True) - # Step 2: Wrap the outer module with FSDP using the BF16 policy + # Step 2️⃣: Wrap the outer module with FSDP using the BF16 policy logger.info("[FSDP v1] wrapping root model with bf16 mixed precision...") self = FSDP( self, @@ -485,6 +585,7 @@ def _assign_child(root_module, dotted_name: str, new_child): "input_layernorm", "post_attention_layernorm", "model.norm", + "pos_embed", "action_preprocessor", "action_processor", ] @@ -657,7 +758,7 @@ def compute_loss( with torch.no_grad(): cross_entropy_loss.detach() - # compute action token accuracy(computed uniformly through tokenizer_mixin) + # compute action token accuracy（computed uniformly through tokenizer_mixin） # if self.tokenizer_mixin is not None and self.action_mapper is not None: # accuracy_dict = self.tokenizer_mixin.compute_accuracy( # logits, labels, self.action_mapper, self.action_token_id_set diff --git a/wall_x/model/qact/qwen2_5/__init__.py b/wall_x/model/qact/qwen2_5/__init__.py index 3040cdd9..1daa3c4d 100644 --- a/wall_x/model/qact/qwen2_5/__init__.py +++ b/wall_x/model/qact/qwen2_5/__init__.py @@ -1,2 +1,2 @@ +from .modeling_qwen2_5_vl_act import Qwen2_5_VLMoEModel, Qwen2_5_VLMoEForAction from .configuration_qwen2_5_vl import Qwen2_5_VLConfig -from .modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction, Qwen2_5_VLMoEModel diff --git a/wall_x/model/qact/qwen2_5/adapter.py b/wall_x/model/qact/qwen2_5/adapter.py index 17352f49..b2cd8d9e 100644 --- a/wall_x/model/qact/qwen2_5/adapter.py +++ b/wall_x/model/qact/qwen2_5/adapter.py @@ -1,4 +1,4 @@ -"""Qwen2.5 VLA adapter - variant-specific overrides on top of VLAdapter.""" +"""Qwen2.5 VLA adapter — variant-specific overrides on top of VLAdapter.""" from wall_x.model.registry import register_model from wall_x.trainer.adapters.vla_model_adapter import VLAdapter @@ -22,6 +22,7 @@ def config_class(cls): @classmethod def inference_model_class(cls): + """Return the model class to use for inference.""" return cls.model_class() def get_transformer_layer_cls(self): diff --git a/wall_x/model/qact/qwen2_5/configuration_qwen2_5_vl.py b/wall_x/model/qact/qwen2_5/configuration_qwen2_5_vl.py index 8956834e..ee891eb0 100644 --- a/wall_x/model/qact/qwen2_5/configuration_qwen2_5_vl.py +++ b/wall_x/model/qact/qwen2_5/configuration_qwen2_5_vl.py @@ -1,9 +1,9 @@ -# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_qwen2_5_vl.py file directly. One of our CI enforces this. -# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # coding=utf-8 # Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. # @@ -339,16 +339,21 @@ def update_model_config(self, train_config): self.action_horizon_flow = train_config["data"].get("action_horizon_flow", 32) - if train_config.get("_attn_implementation", None) is not None: - self._attn_implementation = train_config["_attn_implementation"] + attn_impl = train_config.get( + "_attn_implementation", train_config.get("attn_implementation") + ) + if attn_impl is not None: + self._attn_implementation = attn_impl - if train_config.get("attn_deterministic", None) is not None: - self.attn_deterministic = train_config["attn_deterministic"] - self.vision_config.attn_deterministic = train_config["attn_deterministic"] + # attn_deterministic derived from the single hyperparams.deterministic_algorithms + # switch (MR !556 review) — turning determinism on enables all paths at once. + _det = train_config.get("hyperparams", {}).get( + "deterministic_algorithms", False + ) + self.attn_deterministic = _det + self.vision_config.attn_deterministic = _det + if _det: logger.debug("Attention is using deterministic kernel for this run") - else: - self.attn_deterministic = True - self.vision_config.attn_deterministic = True if train_config.get("noise_scheduler", None) is not None: self.noise_scheduler = train_config["noise_scheduler"] diff --git a/wall_x/model/qact/qwen2_5/inference_mixin.py b/wall_x/model/qact/qwen2_5/inference_mixin.py index 2325a884..3bdffec6 100644 --- a/wall_x/model/qact/qwen2_5/inference_mixin.py +++ b/wall_x/model/qact/qwen2_5/inference_mixin.py @@ -1,8 +1,8 @@ import logging -from typing import List, Optional, Tuple import numpy as np import torch +from typing import Optional, List, Tuple from torchdiffeq import odeint logger = logging.getLogger(__name__) @@ -291,17 +291,12 @@ def generate_dllm_action( flow_action_mask = input_ids == self.action_token_id_set["action_token_id"] inputs_embeds[flow_action_mask] = action_embed - # Compute the token span for each expert group after permutation. - group_size = torch.zeros( - self.config.num_experts, dtype=torch.long, device="cpu" + two_expert_layout = self._prepare_two_expert_layout_for_model( + moe_token_types, + inputs_embeds, + action_starts=kwargs.get("action_starts"), + action_ends=kwargs.get("action_ends"), ) - for i in range(self.config.num_experts): - group_size[i] = (moe_token_types == i).sum() - - # Calculate start and end indices for each expert group - start_indices = torch.cumsum(group_size, dim=0) - group_size - end_indices = torch.cumsum(group_size, dim=0) - prefetch_output = self.model( input_ids=None, attention_mask=attention_mask, @@ -315,8 +310,7 @@ def generate_dllm_action( output_hidden_states=False, return_dict=True, adarms_conds=[None, adarms_cond], - start_indices=start_indices, - end_indices=end_indices, + two_expert_layout=two_expert_layout, ) hidden_states = prefetch_output.last_hidden_state @@ -335,8 +329,6 @@ def generate_dllm_action( ][ 0 ] # Attention, not tested yet with batch - start_indices[1] -= ar_prefix_length.item() - end_indices -= ar_prefix_length.item() if hasattr(prefix_kv_cache, "key_cache"): for layer_i in range(len(prefix_kv_cache.key_cache)): prefix_kv_cache.key_cache[layer_i] = prefix_kv_cache.key_cache[layer_i][ @@ -359,6 +351,9 @@ def generate_dllm_action( ar_postfix_attention_mask = attention_mask[:, ar_prefix_length:] ar_postfix_moe_token_types = moe_token_types[:, ar_prefix_length:] ar_postfix_input_ids = input_ids[:, ar_prefix_length:] + ar_postfix_two_expert_layout = self._prepare_two_expert_layout_for_model( + ar_postfix_moe_token_types, ar_postfix_inputs_embeds + ) noisy_steps_per_sample = ( torch.ceil((1.0 - times[0] - dt) * (total_ar_step + 1)).long() - 1 @@ -462,8 +457,7 @@ def step_with_kvcache( output_hidden_states=False, return_dict=True, adarms_conds=[None, adarms_cond], - start_indices=start_indices, - end_indices=end_indices, + two_expert_layout=ar_postfix_two_expert_layout, ) hidden_states = transformer_outputs.last_hidden_state @@ -561,15 +555,17 @@ def step_with_kvcache( output["predict_action"] = predict_action else: - predict_action = action_trajectory[-1] + predict_action_norm = action_trajectory[-1] + output["predict_action_norm"] = predict_action_norm predict_action = ( self.action_preprocessor.normalizer_action.unnormalize_data( - predict_action, dataset_names + predict_action_norm, dataset_names ) ) output["predict_action"] = predict_action if action_chunk is not None: + output["gt_action_norm"] = action_chunk output["gt_action"] = ( self.action_preprocessor.normalizer_action.unnormalize_data( action_chunk, dataset_names @@ -790,6 +786,7 @@ def generate_flow_action_no_cache( input_ids, action_horizon, action_dim, + init_noise: Optional[torch.Tensor] = None, num_inference_timesteps: int = 10, prefix_length: Optional[int] = None, # Prefix length. attention_mask: Optional[torch.Tensor] = None, @@ -797,8 +794,6 @@ def generate_flow_action_no_cache( past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, moe_token_types: Optional[torch.LongTensor] = None, - start_indices: Optional[torch.Tensor] = None, - end_indices: Optional[torch.Tensor] = None, positional_masks: Optional[torch.LongTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -824,9 +819,10 @@ def generate_flow_action_no_cache( batch_size = ( input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] ) - assert ( - batch_size == 1 - ), "generate_flow_action_no_cache only support batch size 1" + # batch_size > 1 is supported for batched eval where every sample in the + # batch shares an identical prompt prefix (e.g. frames from one episode): + # the body below is batch-safe (noise/noisy_action use batch_size, + # masked-scatter spans the whole batch, prefix_length is unused). output_attentions = ( output_attentions @@ -938,28 +934,23 @@ def generate_flow_action_no_cache( position_ids = position_ids.add(delta) position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) - if start_indices is None or end_indices is None: - # Compute the token span for each expert group after permutation. - group_size = torch.zeros( - self.config.num_experts, dtype=torch.long, device="cpu" - ) - for i in range(self.config.num_experts): - group_size[i] = (moe_token_types == i).sum() - - # Calculate start and end indices for each expert group - start_indices = torch.cumsum(group_size, dim=0) - group_size - end_indices = torch.cumsum(group_size, dim=0) - if action_chunk is not None: action_chunk = action_chunk.to(inputs_embeds.device).to(torch.float32) output = {} - noise = torch.randn( - size=(batch_size, action_horizon, action_dim), - dtype=torch.float32, - device=inputs_embeds.device, - ) + if init_noise is not None: + noise = init_noise.to(device=inputs_embeds.device, dtype=torch.float32) + assert noise.shape == (batch_size, action_horizon, action_dim), ( + f"init_noise shape {tuple(noise.shape)} != " + f"{(batch_size, action_horizon, action_dim)}" + ) + else: + noise = torch.randn( + size=(batch_size, action_horizon, action_dim), + dtype=torch.float32, + device=inputs_embeds.device, + ) noisy_action = noise.clone() dof_mask = dof_mask.to(inputs_embeds.device).to(torch.float32) @@ -967,6 +958,12 @@ def generate_flow_action_no_cache( num_inference_timesteps, inputs_embeds.device, torch.float32 ) flow_action_mask = input_ids == self.action_token_id_set["action_token_id"] + two_expert_layout = self._prepare_two_expert_layout_for_model( + moe_token_types, + inputs_embeds, + action_starts=kwargs.get("action_starts"), + action_ends=kwargs.get("action_ends"), + ) if not dof_mask.all(): padding_action = ( @@ -996,14 +993,13 @@ def step(timestep, noisy_action): past_key_values=None, inputs_embeds=inputs_embeds, moe_token_types=moe_token_types, - start_indices=start_indices, - end_indices=end_indices, positional_masks=positional_masks, use_cache=True, output_attentions=False, output_hidden_states=False, return_dict=True, adarms_conds=[None, adarms_cond], + two_expert_layout=two_expert_layout, ) hidden_states = model_output.last_hidden_state @@ -1011,8 +1007,14 @@ def step(timestep, noisy_action): action_pred = self.action_preprocessor.action_proj_back( action_hidden_states[:, : self.action_preprocessor.action_hidden_size] ) + # The masked gather flattens action tokens to (B*H, D); restore the + # (B, H, D) layout so the batched blends below broadcast correctly + # for batch_size > 1 (bs==1 is numerically unchanged). + action_pred = action_pred.reshape(batch_size, action_horizon, action_dim) if getattr(self.config, "use_x_pred", False): - v_t = (action_pred - noisy_action) / torch.clamp(1 - timestep, min=0.05) + v_t = (action_pred - noisy_action) / torch.clamp( + 1 - timestep.view(batch_size, 1, 1), min=0.05 + ) else: v_t = action_pred @@ -1023,9 +1025,12 @@ def step(timestep, noisy_action): action_trajectory = odeint(step, noisy_action, times, method="euler") - predict_action = action_trajectory[-1] + predict_action_norm = action_trajectory[-1] + output["predict_action_norm"] = predict_action_norm + if action_chunk is not None: + output["gt_action_norm"] = action_chunk predict_action = self.action_preprocessor.normalizer_action.unnormalize_data( - predict_action, dataset_names + predict_action_norm, dataset_names ) output["predict_action"] = predict_action # normalize action chunk to get gt_action diff --git a/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py index 9eff1feb..3cc73717 100644 --- a/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py +++ b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py @@ -1,9 +1,9 @@ -# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_qwen2_5_vl.py file directly. One of our CI enforces this. -# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # coding=utf-8 # Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. # @@ -64,22 +64,30 @@ from .configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig from wall_x.model.core.attention.selector import AttentionsSelectorMixin -from wall_x.model.core.ops import rot_pos_emb, get_window_index, m_rope +from wall_x.model.core.ops import ( + adarmsnorm, + rot_pos_emb, + get_window_index, + m_rope, + rmsnorm, + rope, + silu_and_mul, +) if is_flash_attn_2_available(): from flash_attn import flash_attn_func - from flash_attn import flash_attn_varlen_func + from flash_attn import flash_attn_varlen_qkvpacked_func from flash_attn.layers.rotary import apply_rotary_emb else: - flash_attn_varlen_func = None + flash_attn_varlen_qkvpacked_func = None apply_rotary_emb = None if is_flash_attn_2_available(): pass else: - flash_attn_varlen_func = None + flash_attn_varlen_qkvpacked_func = None try: @@ -118,6 +126,7 @@ def __init__(self, config, bias=False, use_selective_recompute=False): ) self.act_fn = ACT2FN[config.hidden_act] + self.use_silu_and_mul = config.hidden_act in ("silu", "swish") self.use_selective_recompute = use_selective_recompute def _load_from_state_dict( @@ -154,7 +163,7 @@ def _load_from_state_dict( key_dw = prefix + "down_proj.weight" if key_dw in state_dict and state_dict[key_dw].shape[1] == orig: state_dict[key_dw] = F.pad(state_dict[key_dw], (0, pad)) - # down_proj.bias: [hidden] - no change needed + # down_proj.bias: [hidden] — no change needed super()._load_from_state_dict( state_dict, @@ -175,27 +184,32 @@ def _save_to_state_dict(self, destination, prefix, keep_vars): orig = self.intermediate_size padded = self.padded_intermediate_size - # gate_up_proj.weight: [2*padded, hidden] -> [2*orig, hidden] + # gate_up_proj.weight: [2*padded, hidden] → [2*orig, hidden] key_w = prefix + "gate_up_proj.weight" if key_w in destination: w = destination[key_w] gate_w, up_w = w.split([padded, padded], dim=0) destination[key_w] = torch.cat([gate_w[:orig], up_w[:orig]], dim=0) - # gate_up_proj.bias: [2*padded] -> [2*orig] + # gate_up_proj.bias: [2*padded] → [2*orig] key_b = prefix + "gate_up_proj.bias" if key_b in destination: b = destination[key_b] gate_b, up_b = b.split([padded, padded]) destination[key_b] = torch.cat([gate_b[:orig], up_b[:orig]]) - # down_proj.weight: [hidden, padded] -> [hidden, orig] + # down_proj.weight: [hidden, padded] → [hidden, orig] key_dw = prefix + "down_proj.weight" if key_dw in destination: destination[key_dw] = destination[key_dw][:, :orig] - def _activation_chunk(self, gate_out, up_out): + def _activation_chunk(self, gate_up_out): # gate_out / up_out are precomputed GEMM results + if self.use_silu_and_mul: + return silu_and_mul(gate_up_out, inplace=True) + gate_out, up_out = gate_up_out.split( + [self.padded_intermediate_size, self.padded_intermediate_size], dim=-1 + ) act = self.act_fn(gate_out) return act * up_out @@ -204,17 +218,8 @@ def forward(self, hidden_state): # 1. Precompute GEMM without checkpointing # ----------------------------- gate_up_out = self.gate_up_proj(hidden_state) - gate_out, up_out = gate_up_out.split( - [self.padded_intermediate_size, self.padded_intermediate_size], dim=-1 - ) - if self.use_selective_recompute: - # checkpoint only activation, not GEMM - out = cp.checkpoint( - self._activation_chunk, gate_out, up_out, use_reentrant=False - ) - else: - out = self._activation_chunk(gate_out, up_out) + out = self._activation_chunk(gate_up_out) return self.down_proj(out) @@ -308,35 +313,16 @@ def _norm(self, x): return normed_inputs def forward(self, hidden_states, cond=None): - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - normed_inputs = self._norm(hidden_states) - if cond is None or self.dense is None: - # Regular RMSNorm - normed_inputs = self.weight * normed_inputs - return normed_inputs.to(input_dtype), None + return rmsnorm(hidden_states, self.weight, self.variance_epsilon), None - # Adaptive RMSNorm if cond.shape[-1] != self.cond_dim: raise ValueError( f"Expected cond dimension {self.cond_dim}, got {cond.shape[-1]}" ) - # Compute modulation parameters - cond = cond.to(dtype=self.dense.weight.dtype) modulation = self.dense(cond) - if len(hidden_states.shape) == 3: # [batch, seq, features] - modulation = modulation.unsqueeze(1) - - scale, shift, gate = torch.chunk(modulation, 3, dim=-1) - - # Apply adaptive normalization - normed_inputs = normed_inputs * (1 + scale.to(torch.float32)) + shift.to( - torch.float32 - ) - - return normed_inputs.to(input_dtype), gate.to(input_dtype) + return adarmsnorm(hidden_states, modulation, self.variance_epsilon) def extra_repr(self): repr_str = f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" @@ -362,16 +348,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x -def apply_rotary_pos_emb_flashatt( - q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor -) -> Tuple[torch.Tensor, torch.Tensor]: - cos = cos.chunk(2, dim=-1)[0].contiguous() - sin = sin.chunk(2, dim=-1)[0].contiguous() - q_embed = apply_rotary_emb(q.float(), cos.float(), sin.float()).type_as(q) - k_embed = apply_rotary_emb(k.float(), cos.float(), sin.float()).type_as(k) - return q_embed, k_embed - - class Qwen2_5_VLVisionFlashAttention2(nn.Module): def __init__(self, config: Qwen2_5_VLConfig, dim: int, num_heads: int = 16) -> None: super().__init__() @@ -389,12 +365,7 @@ def forward( position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: seq_length = hidden_states.shape[0] - q, k, v = ( - self.qkv(hidden_states) - .reshape(seq_length, 3, self.num_heads, -1) - .permute(1, 0, 2, 3) - .unbind(0) - ) + qkv = self.qkv(hidden_states) if position_embeddings is None: logger.warning_once( "The attention layers in this model are transitioning from computing the RoPE embeddings internally " @@ -402,25 +373,25 @@ def forward( "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.54 `rotary_pos_emb` will be " "removed and `position_embeddings` will be mandatory." ) - emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) - cos = emb.cos().float() - sin = emb.sin().float() + cos = rotary_pos_emb.cos().float().contiguous() + sin = rotary_pos_emb.sin().float().contiguous() else: cos, sin = position_embeddings - q, k = apply_rotary_pos_emb_flashatt(q.unsqueeze(0), k.unsqueeze(0), cos, sin) - q = q.squeeze(0) - k = k.squeeze(0) + qkv = rope.pack( + qkv, + self.num_heads, + self.num_heads, + cos.unsqueeze(0), + sin.unsqueeze(0), + ) + qkv = qkv.view(seq_length, 3, self.num_heads, -1) if max_seqlen is None: max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() - attn_output = flash_attn_varlen_func( - q, - k, - v, - cu_seqlens, + attn_output = flash_attn_varlen_qkvpacked_func( + qkv, cu_seqlens, max_seqlen, - max_seqlen, deterministic=self.deterministic, ).reshape(seq_length, -1) attn_output = self.proj(attn_output) @@ -440,7 +411,8 @@ def apply_rotary_pos_emb_vision( orig_q_dtype = q.dtype orig_k_dtype = k.dtype q, k = q.float(), k.float() - cos, sin = cos.unsqueeze(-2), sin.unsqueeze(-2) + cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2) + sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) q_embed = q_embed.to(orig_q_dtype) @@ -478,12 +450,11 @@ def forward( "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.54 `rotary_pos_emb` will be " "removed and `position_embeddings` will be mandatory." ) - emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) - cos = emb.cos().float() - sin = emb.sin().float() + cos = rotary_pos_emb.cos().float().contiguous() + sin = rotary_pos_emb.sin().float().contiguous() else: cos, sin = position_embeddings - q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) + q, k = rope(q, k, cos, sin) attention_mask = torch.full( [1, seq_length, seq_length], @@ -541,12 +512,11 @@ def forward( "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.54 `rotary_pos_emb` will be " "removed and `position_embeddings` will be mandatory." ) - emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) - cos = emb.cos().float() - sin = emb.sin().float() + cos = rotary_pos_emb.cos().float().contiguous() + sin = rotary_pos_emb.sin().float().contiguous() else: cos, sin = position_embeddings - q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) + q, k = rope(q, k, cos, sin) attention_mask = torch.zeros( [1, seq_length, seq_length], device=q.device, dtype=torch.bool @@ -854,8 +824,10 @@ def forward( ) rotary_pos_emb = rotary_pos_emb[window_index, :, :] rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) - emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) - position_embeddings = (emb.cos(), emb.sin()) + position_embeddings = ( + rotary_pos_emb.cos().float().contiguous(), + rotary_pos_emb.sin().float().contiguous(), + ) cu_seqlens = torch.repeat_interleave( grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] @@ -1003,15 +975,14 @@ def forward(self, x, position_ids): freqs = ( inv_freq_expanded.float() @ position_ids_expanded.float() ).transpose(2, 3) - emb = torch.cat((freqs, freqs), dim=-1) - cos = emb.cos() - sin = emb.sin() + cos = freqs.cos() + sin = freqs.sin() # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention cos = cos * self.attention_scaling sin = sin * self.attention_scaling - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + return cos.float().contiguous(), sin.float().contiguous() class Qwen2MLP(nn.Module): @@ -1060,24 +1031,25 @@ def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim unsqueeze_dim (`int`, *optional*, defaults to 1): The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim // 2]. Then, if q and k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ - mrope_section = mrope_section * 2 cos = torch.cat( [m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1 - ).unsqueeze(unsqueeze_dim) + ) sin = torch.cat( [m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1 - ).unsqueeze(unsqueeze_dim) + ) + cos = torch.cat((cos, cos), dim=-1).unsqueeze(unsqueeze_dim) + sin = torch.cat((sin, sin), dim=-1).unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed + return q_embed.to(q.dtype), k_embed.to(k.dtype) def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: @@ -1274,10 +1246,10 @@ def forward( cos, sin = position_embeddings query_states, key_states = m_rope( - query_states.contiguous(), - key_states.contiguous(), - cos[..., : (cos.size(-1) // 2)].contiguous().float(), - sin[..., : (sin.size(-1) // 2)].contiguous().float(), + query_states, + key_states, + cos, + sin, self.rope_scaling["mrope_section"], ) if past_key_value is not None: @@ -1301,14 +1273,13 @@ def forward( # value_states = repeat_kv(value_states, self.num_key_value_groups) dropout_rate = 0.0 if not self.training else self.attention_dropout - # In PEFT, usually we cast the layer norms in float32 for training stability reasons - # therefore the input hidden states gets silently casted in float32. Hence, we need - # cast them back in float16 just to be sure everything works as expected. + # In PEFT/FSDP mixed precision, layer norms may upcast hidden states to + # float32. Flash Attention 2 does not support fp32 inputs, so restore the + # original graceful cast path instead of hard-failing the training step. input_dtype = query_states.dtype if input_dtype == torch.float32: if torch.is_autocast_enabled(): target_dtype = torch.get_autocast_gpu_dtype() - # Handle the case where the model is quantized elif hasattr(self.config, "_pre_quantization_dtype"): target_dtype = self.config._pre_quantization_dtype else: @@ -1711,8 +1682,7 @@ def forward( cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): - Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, - with `head_dim` being the embedding dimension of each attention head. + Tuple containing cosine and sine positional embeddings of shape `(3, batch_size, seq_len, head_dim // 2)`. kwargs (`dict`, *optional*): Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code into the model diff --git a/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py index 4a0f9917..eef23c27 100644 --- a/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py +++ b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py @@ -16,6 +16,7 @@ Qwen2RMSNorm, Qwen2_5_VLPatchMerger, Qwen2_5_VLVisionBlock, + Qwen2_5_VisionPatchEmbed, Qwen2_5_VisionTransformerPretrainedModel, ) @@ -49,7 +50,12 @@ normalize_data_with_virtual_tail, unnormalize_data_with_virtual_tail, ) -from wall_x.model.core.action.moe import TokenTypeRouter, SparseMoeBlock +from wall_x.model.core.action.moe import ( + TokenTypeRouter, + SparseMoeBlock, +) +from wall_x.model.core.action.layout import TwoExpertLayout, prepare_two_expert_layout +from wall_x.utils.checkpoint import maybe_checkpoint from wall_x.model.core.attention.selector import ( ATTENTION_TYPES_WITH_2D_MASK, ATTENTION_TYPES_WITH_FLASH_MASK, @@ -57,7 +63,7 @@ from wall_x.model.core.attention.joint import JOINT_QWEN_ATTENTION_CLASSES from wall_x.model.qact.qwen2_5.inference_mixin import VLAInferenceMixin -from wall_x.model.core.ops import unpermute, permute, get_rope_index +from wall_x.model.core.ops import get_rope_index, two_expert_pack, two_expert_unpack logger = logging.get_logger(__name__) @@ -68,6 +74,11 @@ "norm", ) +_QWEN25_DMUON_SHARDED_ADAMW_NAME_PARTS = ( + "embed_tokens", + "lm_head", +) + def _is_qwen25_dmuon_target_param(name: str, param: nn.Parameter) -> bool: """Select trainable Qwen2.5 VLA matrices for DMuon.""" @@ -76,6 +87,86 @@ def _is_qwen25_dmuon_target_param(name: str, param: nn.Parameter) -> bool: return not any(part in name for part in _QWEN25_DMUON_BLOCKED_NAME_PARTS) +def _is_qwen25_dmuon_sharded_adamw_param(name: str, param: nn.Parameter) -> bool: + if not param.requires_grad or param.ndim != 2 or not name.endswith(".weight"): + return False + return any(part in name for part in _QWEN25_DMUON_SHARDED_ADAMW_NAME_PARTS) + + +def _is_qwen25_dmuon_managed_param(name: str, param: nn.Parameter) -> bool: + return _is_qwen25_dmuon_target_param( + name, param + ) or _is_qwen25_dmuon_sharded_adamw_param(name, param) + + +def _qwen25_dmuon_route_hint(name: str, param: nn.Parameter) -> Optional[str]: + if _is_qwen25_dmuon_sharded_adamw_param(name, param): + return "sharded_adamw" + if _is_qwen25_dmuon_target_param(name, param): + return "muon" + return None + + +def _is_qwen25_dmuon_trainable_param(name: str, param: nn.Parameter) -> bool: + del name + return bool(param.requires_grad) + + +def _qwen25_pure_dmuon_route_hint(name: str, param: nn.Parameter) -> Optional[str]: + if not param.requires_grad: + return None + if _is_qwen25_dmuon_target_param(name, param): + return "muon" + if _is_qwen25_dmuon_sharded_adamw_param(name, param): + return "sharded_adamw" + return "adamw" + + +def _format_qwen25_route_examples(examples: dict[str, list[str]]) -> str: + return ", ".join( + f"{route}={names}" for route, names in sorted(examples.items()) if names + ) + + +def _validate_qwen25_pure_dmuon_route_plan( + named_parameters, +) -> tuple[dict[str, int], dict[str, list[str]]]: + route_counts: dict[str, int] = {} + route_examples: dict[str, list[str]] = { + "muon": [], + "adamw": [], + "sharded_adamw": [], + } + bad_routes: list[tuple[str, str]] = [] + + for name, param in named_parameters: + if not param.requires_grad: + continue + route = _qwen25_pure_dmuon_route_hint(name, param) + route_key = str(route) + route_counts[route_key] = route_counts.get(route_key, 0) + 1 + route_examples.setdefault(route_key, []) + if len(route_examples[route_key]) < 8: + route_examples[route_key].append(name) + + if route == "sharded_adamw" and not _is_qwen25_dmuon_sharded_adamw_param( + name, param + ): + bad_routes.append((name, "unexpected sharded_adamw")) + if route == "muon" and not _is_qwen25_dmuon_target_param(name, param): + bad_routes.append((name, "unexpected muon")) + if route == "adamw" and ( + _is_qwen25_dmuon_target_param(name, param) + or _is_qwen25_dmuon_sharded_adamw_param(name, param) + ): + bad_routes.append((name, "unexpected adamw")) + + if bad_routes: + preview = bad_routes[:16] + raise RuntimeError(f"Invalid Qwen2.5 Pure DMuon route plan: {preview}") + return route_counts, route_examples + + @dataclass class Qwen2_5_VLACausalLMOutputWithPast(ModelOutput): loss: Optional[torch.FloatTensor] = None @@ -184,11 +275,11 @@ def forward( ] = None, # necessary, but kept here for BC # for vla token_types: Optional[torch.LongTensor] = None, - start_indices: Optional[torch.Tensor] = None, - end_indices: Optional[torch.Tensor] = None, - probs: Optional[torch.Tensor] = None, - row_id_map: Optional[torch.Tensor] = None, orig_shape: Optional[Tuple[int, int, int]] = None, + action_starts: Optional[torch.Tensor] = None, + action_ends: Optional[torch.Tensor] = None, + expert_segments: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None, + two_expert_layout: Optional[TwoExpertLayout] = None, adarms_conds: Optional[List[torch.Tensor]] = [None, None], **kwargs, ) -> Tuple[ @@ -209,8 +300,7 @@ def forward( cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): - Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, - with `head_dim` being the embedding dimension of each attention head. + Tuple containing cosine and sine positional embeddings of shape `(3, batch_size, seq_len, head_dim // 2)`. kwargs (`dict`, *optional*): Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code into the model @@ -223,9 +313,11 @@ def forward( adarms_conds, self.input_layernorms, self.input_layernorm, - start_indices, - end_indices, - self.use_selective_recompute, + action_starts=action_starts, + action_ends=action_ends, + expert_segments=expert_segments, + two_expert_layout=two_expert_layout, + use_selective_recompute=self.use_selective_recompute, ) # Self Attention @@ -239,11 +331,11 @@ def forward( use_cache=use_cache, cache_position=cache_position, token_types=token_types, - start_indices=start_indices, - end_indices=end_indices, - probs=probs, - row_id_map=row_id_map, orig_shape=orig_shape, + action_starts=action_starts, + action_ends=action_ends, + expert_segments=expert_segments, + two_expert_layout=two_expert_layout, position_embeddings=position_embeddings, ) else: @@ -259,7 +351,14 @@ def forward( ) hidden_states = self._gated_residual( - residual, hidden_states, gate, start_indices, end_indices + residual, + hidden_states, + gate, + token_types=token_types, + action_starts=action_starts, + action_ends=action_ends, + expert_segments=expert_segments, + two_expert_layout=two_expert_layout, ) # Fully Connected @@ -271,17 +370,31 @@ def forward( adarms_conds, self.post_attention_layernorms, self.post_attention_layernorm, - start_indices, - end_indices, - self.use_selective_recompute, + action_starts=action_starts, + action_ends=action_ends, + expert_segments=expert_segments, + two_expert_layout=two_expert_layout, + use_selective_recompute=self.use_selective_recompute, ) hidden_states = self._apply_mlp_moe( - hidden_states, token_types, start_indices, end_indices + hidden_states, + token_types, + action_starts=action_starts, + action_ends=action_ends, + expert_segments=expert_segments, + two_expert_layout=two_expert_layout, ) hidden_states = self._gated_residual( - residual, hidden_states, gate, start_indices, end_indices + residual, + hidden_states, + gate, + token_types=token_types, + action_starts=action_starts, + action_ends=action_ends, + expert_segments=expert_segments, + two_expert_layout=two_expert_layout, ) outputs = (hidden_states,) @@ -356,9 +469,37 @@ def __init__(self, config: Qwen2_5_VLConfig, use_selective_recompute=False): self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) self.gradient_checkpointing = False + # Partial recompute: only the first N decoder layers are checkpointed + # (None = all layers). Set via Qwen2_5_VLMoEForAction. + # gradient_checkpointing_enable(num_layers=...); same semantics as the + # pi0 path (gemma_pytorch / distribution_strategy._maybe_enable_grad_ckpt). + self.gradient_checkpoint_num_layers = None # Initialize weights and apply final processing self.post_init() + def _resolve_two_expert_layout( + self, + moe_token_types: torch.Tensor, + action_starts: Optional[torch.Tensor] = None, + action_ends: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, + ) -> TwoExpertLayout: + if self.config.num_experts != 2: + raise ValueError( + "mot_opt fixed scatter path requires config.num_experts == 2" + ) + if moe_token_types is None or moe_token_types.dim() != 2: + raise ValueError("mot_opt fixed scatter path expects [B, S] token types") + + return prepare_two_expert_layout( + moe_token_types, + batch_size=moe_token_types.size(0), + seq_length=moe_token_types.size(1), + action_starts=action_starts, + action_ends=action_ends, + device=device if device is not None else moe_token_types.device, + ) + def get_input_embeddings(self): return self.embed_tokens @@ -373,8 +514,9 @@ def forward( past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, moe_token_types: Optional[torch.LongTensor] = None, # new parameter - start_indices: Optional[torch.Tensor] = None, - end_indices: Optional[torch.Tensor] = None, + action_starts: Optional[torch.Tensor] = None, + action_ends: Optional[torch.Tensor] = None, + two_expert_layout: Optional[TwoExpertLayout] = None, positional_masks: Optional[ dict ] = None, # stores token position masks needed by each category @@ -408,10 +550,6 @@ def forward( ) if moe_token_types is None: raise ValueError("moe_token_types must be provided for MoE routing") - if start_indices is None or end_indices is None: - raise ValueError( - "start_indices and end_indices must be provided for MoE routing" - ) if self.gradient_checkpointing and self.training: if use_cache: @@ -473,13 +611,21 @@ def forward( # When mot_opt is enabled, permute tokens from different experts first; shape becomes [Tokens, HiddenSize] orig_shape = hidden_states.shape moe_token_types = moe_token_types.contiguous() - if self.config.mot_opt: - hidden_states = hidden_states.view(-1, hidden_states.size(-1)) - hidden_states, row_id_map = permute(hidden_states, moe_token_types.view(-1)) + if self.config.mot_opt or self.config.mlp_moe: + if two_expert_layout is None: + two_expert_layout = self._resolve_two_expert_layout( + moe_token_types, + action_starts=action_starts, + action_ends=action_ends, + device=hidden_states.device, + ) + action_starts = two_expert_layout.action_starts + action_ends = two_expert_layout.action_ends + expert_segments = two_expert_layout.expert_segments else: - row_id_map = None - # Use reshape instead of view to support non-contiguous tensors - probs = torch.ones_like(moe_token_types, dtype=torch.float32).reshape(-1, 1) + expert_segments = None + if self.config.mot_opt: + hidden_states = two_expert_pack(hidden_states, two_expert_layout) # decoder layers all_hidden_states = () if output_hidden_states else None @@ -508,52 +654,44 @@ def forward( positional_masks=positional_masks, ) - for decoder_layer in self.layers: + n_ckpt_layers = self.gradient_checkpoint_num_layers + if n_ckpt_layers is None or not 0 <= n_ckpt_layers <= len(self.layers): + n_ckpt_layers = len(self.layers) + + for layer_idx, decoder_layer in enumerate(self.layers): if output_hidden_states: assert ( self.config.mot_opt is False ), "When using mot_opt, output_hidden_states is not supported yet." all_hidden_states += (hidden_states,) - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - causal_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - # for vla - moe_token_types, - start_indices, - end_indices, - probs, - row_id_map, - orig_shape, - adarms_conds, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_value=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - # for vla - token_types=moe_token_types, - start_indices=start_indices, - end_indices=end_indices, - probs=probs, - row_id_map=row_id_map, - orig_shape=orig_shape, - adarms_conds=adarms_conds, - ) + # wall_x.utils.checkpoint（use_reentrant=False / preserve_rng_state=False + # 固化，见该模块 docstring）：enabled=False 时直接裸调，替代原 + # HF _gradient_checkpointing_func + if/else 双分支。 + layer_outputs = maybe_checkpoint( + decoder_layer.__call__, + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + # for vla + token_types=moe_token_types, + orig_shape=orig_shape, + action_starts=action_starts, + action_ends=action_ends, + expert_segments=expert_segments, + two_expert_layout=two_expert_layout, + adarms_conds=adarms_conds, + enabled=( + self.gradient_checkpointing + and self.training + and layer_idx < n_ckpt_layers + ), + ) hidden_states = layer_outputs[0] if use_cache: @@ -568,9 +706,11 @@ def forward( adarms_conds, self.norms, self.norm, - start_indices, - end_indices, - self.use_selective_recompute, + action_starts=action_starts, + action_ends=action_ends, + expert_segments=expert_segments, + two_expert_layout=two_expert_layout, + use_selective_recompute=self.use_selective_recompute, ) # add hidden states from the last decoder layer @@ -583,8 +723,12 @@ def forward( next_cache = next_decoder_cache if use_cache else None if self.config.mot_opt: - hidden_states = unpermute(hidden_states, row_id_map, probs) - hidden_states = hidden_states.view(orig_shape) + hidden_states = two_expert_unpack( + hidden_states, + orig_shape[0], + orig_shape[1], + two_expert_layout, + ) if not return_dict: return tuple( @@ -861,6 +1005,23 @@ def _build_visual(self, config, use_selective_recompute): use_selective_recompute=use_selective_recompute, ) + def gradient_checkpointing_enable( + self, gradient_checkpointing_kwargs=None, num_layers=None + ): + """Partial recompute: `num_layers` limits full-recompute to the first N + decoder layers (None = all, legacy). The `num_layers` parameter name is + what distribution_strategy._maybe_enable_grad_ckpt probes for, mirroring + the pi0 path; vision blocks are unaffected (still all checkpointed).""" + self.model.gradient_checkpoint_num_layers = num_layers + if gradient_checkpointing_kwargs is None: + # HF default is reentrant checkpoint, which rejects the extra + # kwargs the vision block forward takes (cu_seqlens/...) and is + # incompatible with FSDP2 nested fully_shard (see + # wall_x/utils/checkpoint.py docstring). Previously papered over + # by an exp-side sitecustomize shim; make False the in-repo default. + gradient_checkpointing_kwargs = {"use_reentrant": False} + super().gradient_checkpointing_enable(gradient_checkpointing_kwargs) + def _post_init_engine(self, config): """Hook called at the end of __init__. Subclasses may override to add engine-specific state.""" pass @@ -883,6 +1044,29 @@ def set_decoder(self, decoder): def get_decoder(self): return self.model + def _prepare_two_expert_layout_for_model( + self, + moe_token_types: Optional[torch.Tensor], + inputs_embeds: Optional[torch.Tensor], + action_starts: Optional[torch.Tensor] = None, + action_ends: Optional[torch.Tensor] = None, + ) -> Optional[TwoExpertLayout]: + if not (self.model.config.mot_opt or self.model.config.mlp_moe): + return None + if moe_token_types is None: + return None + device = ( + inputs_embeds.device + if inputs_embeds is not None + else moe_token_types.device + ) + return self.model._resolve_two_expert_layout( + moe_token_types, + action_starts=action_starts, + action_ends=action_ends, + device=device, + ) + def get_rope_index( self, input_ids: Optional[torch.LongTensor] = None, @@ -1116,8 +1300,6 @@ def get_attention_maps( past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, moe_token_types: Optional[torch.LongTensor] = None, - start_indices: Optional[torch.Tensor] = None, - end_indices: Optional[torch.Tensor] = None, positional_masks: Optional[torch.LongTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -1151,18 +1333,6 @@ def get_attention_maps( return_dict if return_dict is not None else self.config.use_return_dict ) - if start_indices is None or end_indices is None: - # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here - group_size = torch.zeros( - self.config.num_experts, dtype=torch.long, device="cpu" - ) - for i in range(self.config.num_experts): - group_size[i] = (moe_token_types == i).sum() - - # Calculate start and end indices for each expert group - start_indices = torch.cumsum(group_size, dim=0) - group_size - end_indices = torch.cumsum(group_size, dim=0) - # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme if position_ids is None and ( attention_mask is None or attention_mask.ndim == 2 @@ -1246,6 +1416,12 @@ def get_attention_maps( if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) + two_expert_layout = self._prepare_two_expert_layout_for_model( + moe_token_types, + inputs_embeds, + action_starts=kwargs.get("action_starts"), + action_ends=kwargs.get("action_ends"), + ) outputs = self.model( input_ids=None, position_ids=position_ids, @@ -1253,14 +1429,13 @@ def get_attention_maps( past_key_values=past_key_values, inputs_embeds=inputs_embeds, moe_token_types=moe_token_types, # pass token_types - start_indices=start_indices, - end_indices=end_indices, positional_masks=positional_masks, # pass masks for each category use_cache=False, output_attentions=True, output_hidden_states=False, return_dict=return_dict, adarms_conds=[None, adarms_cond], + two_expert_layout=two_expert_layout, # cache_position=cache_position, ) attention_maps = torch.stack(outputs.attentions, dim=0)[ @@ -1293,8 +1468,6 @@ def train_step_forward( # for vla sample_time: Optional[torch.FloatTensor] = None, moe_token_types: Optional[torch.LongTensor] = None, - start_indices: Optional[torch.Tensor] = None, - end_indices: Optional[torch.Tensor] = None, positional_masks: Optional[dict] = None, action_chunk: Optional[torch.FloatTensor] = None, proprioception: Optional[torch.FloatTensor] = None, @@ -1321,18 +1494,6 @@ def train_step_forward( return_dict if return_dict is not None else self.config.use_return_dict ) - if start_indices is None or end_indices is None: - # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here - group_size = torch.zeros( - self.config.num_experts, dtype=torch.long, device="cpu" - ) - for i in range(self.config.num_experts): - group_size[i] = (moe_token_types == i).sum() - - # Calculate start and end indices for each expert group - start_indices = torch.cumsum(group_size, dim=0) - group_size - end_indices = torch.cumsum(group_size, dim=0) - # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme if position_ids is None and ( attention_mask is None or attention_mask.ndim == 2 @@ -1421,6 +1582,12 @@ def train_step_forward( if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) + two_expert_layout = self._prepare_two_expert_layout_for_model( + moe_token_types, + inputs_embeds, + action_starts=kwargs.get("action_starts"), + action_ends=kwargs.get("action_ends"), + ) outputs = self.model( input_ids=None, position_ids=position_ids, @@ -1428,14 +1595,13 @@ def train_step_forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, moe_token_types=moe_token_types, # pass token_types - start_indices=start_indices, - end_indices=end_indices, positional_masks=positional_masks, # pass masks for each category use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, adarms_conds=[None, adarms_cond], + two_expert_layout=two_expert_layout, # cache_position=cache_position, ) @@ -1444,19 +1610,17 @@ def train_step_forward( # --- optimization: loss-token-only LM Head --- # During training, only compute lm_head for tokens that contribute to CE loss, # reducing GEMM size from [B*S, V] to [N_loss, V]. - # FSDP fp32 layernorms emit fp32 hidden_states; lm_head runs in bf16. - _lm_dtype = self.lm_head.weight.dtype _lm_loss_mask = None if labels is not None and self.training: shift_labels = labels[..., 1:].contiguous() _lm_loss_mask = shift_labels != -100 # [B, S-1] if _lm_loss_mask.any(): lm_hidden = hidden_states[:, :-1, :][_lm_loss_mask] # [N_loss, H] - logits = self.lm_head(lm_hidden.to(_lm_dtype)) # [N_loss, V] + logits = self.lm_head(lm_hidden) # [N_loss, V] else: logits = None else: - logits = self.lm_head(hidden_states.to(_lm_dtype)) + logits = self.lm_head(hidden_states) ( loss, @@ -1517,8 +1681,6 @@ def generate_text( past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, moe_token_types: Optional[torch.LongTensor] = None, # new parameter - start_indices: Optional[torch.Tensor] = None, - end_indices: Optional[torch.Tensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -1568,25 +1730,11 @@ def generate_text( else: prefix_length = input_ids.shape[1] - if start_indices is None or end_indices is None: - # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here - group_size = torch.zeros( - self.config.num_experts, dtype=torch.long, device="cpu" - ) - for i in range(self.config.num_experts): - group_size[i] = (moe_token_types == i).sum() - - # Calculate start and end indices for each expert group - start_indices = torch.cumsum(group_size, dim=0) - group_size - end_indices = torch.cumsum(group_size, dim=0) - batch = { "input_ids": input_ids, "attention_mask": attention_mask, "pixel_values": pixel_values, "moe_token_types": moe_token_types, - "start_indices": start_indices, - "end_indices": end_indices, "image_grid_thw": image_grid_thw, "dof_mask": dof_mask, "agent_pos_mask": agent_pos_mask, @@ -1778,8 +1926,6 @@ def generate_flow_action( past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, moe_token_types: Optional[torch.LongTensor] = None, - start_indices: Optional[torch.Tensor] = None, - end_indices: Optional[torch.Tensor] = None, positional_masks: Optional[torch.LongTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -1814,8 +1960,6 @@ def generate_flow_action( past_key_values=past_key_values, inputs_embeds=inputs_embeds, moe_token_types=moe_token_types, - start_indices=start_indices, - end_indices=end_indices, positional_masks=positional_masks, labels=labels, use_cache=use_cache, @@ -1865,8 +2009,6 @@ def _prepare_flow_action_inputs( past_key_values, inputs_embeds, moe_token_types, - start_indices, - end_indices, positional_masks, labels, use_cache, @@ -2023,29 +2165,6 @@ def _prepare_flow_action_inputs( ) prefix_length = int(prefix_length[0].item()) - if start_indices is None or end_indices is None: - # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here - # Cache key: shape + sum is sufficient for stable-state (moe_token_types structure doesn't change) - _moe_key = (moe_token_types.shape, int(moe_token_types.sum().item())) - _moe_cache = self._infer_stable_cache.get("moe_indices") - if _moe_cache is None or _moe_cache[0] != _moe_key: - group_size = torch.zeros( - self.config.num_experts, dtype=torch.long, device="cpu" - ) - for i in range(self.config.num_experts): - group_size[i] = (moe_token_types == i).sum() - - # Calculate start and end indices for each expert group - start_indices = torch.cumsum(group_size, dim=0) - group_size - end_indices = torch.cumsum(group_size, dim=0) - self._infer_stable_cache["moe_indices"] = ( - _moe_key, - start_indices, - end_indices, - ) - else: - _, start_indices, end_indices = _moe_cache - # Timing: action initialization if action_chunk is not None: action_chunk = action_chunk.to(inputs_embeds.device).to(torch.float32) @@ -2100,8 +2219,6 @@ def _prepare_flow_action_inputs( "position_ids": position_ids, "prefix_length": prefix_length, "flow_action_mask": flow_action_mask, - "start_indices": start_indices, - "end_indices": end_indices, "dof_mask": dof_mask, "noise": noise, "noisy_action": noisy_action, @@ -2123,6 +2240,12 @@ def _generate_flow_action_vanilla(self, ctx): # Timing: prefill forward pass all_attention_maps = [] output_attentions = True if ctx["output_attentions"] else False + two_expert_layout = self._prepare_two_expert_layout_for_model( + ctx["moe_token_types"], + ctx["inputs_embeds"], + action_starts=ctx.get("action_starts"), + action_ends=ctx.get("action_ends"), + ) prefetch_output = self.model( input_ids=None, attention_mask=ctx["attention_mask"], @@ -2130,14 +2253,13 @@ def _generate_flow_action_vanilla(self, ctx): past_key_values=None, inputs_embeds=ctx["inputs_embeds"], moe_token_types=ctx["moe_token_types"], - start_indices=ctx["start_indices"], - end_indices=ctx["end_indices"], positional_masks=ctx["positional_masks"], use_cache=True, output_attentions=output_attentions, output_hidden_states=False, return_dict=True, adarms_conds=[None, ctx["adarms_cond"]], + two_expert_layout=two_expert_layout, ) hidden_states = prefetch_output.last_hidden_state prefix_kv_cache = prefetch_output.past_key_values @@ -2166,10 +2288,11 @@ def _generate_flow_action_vanilla(self, ctx): _postfix_attention_mask, postfix_moe_token_types, postfix_input_ids, - postfix_start_indices, - postfix_end_indices, padding_mask, ) = self._prepare_flow_postfix(ctx, prefix_kv_cache) + postfix_two_expert_layout = self._prepare_two_expert_layout_for_model( + postfix_moe_token_types, postfix_inputs_embeds + ) def step_with_kvcache(timestep, noisy_action): action_mask = ( @@ -2192,13 +2315,12 @@ def step_with_kvcache(timestep, noisy_action): past_key_values=prefix_kv_cache, inputs_embeds=temp_inputs_embeds, moe_token_types=postfix_moe_token_types, - start_indices=postfix_start_indices, - end_indices=postfix_end_indices, use_cache=False, output_attentions=output_attentions, output_hidden_states=False, return_dict=True, adarms_conds=[None, adarms_cond], + two_expert_layout=postfix_two_expert_layout, ) if output_attentions: all_attention_maps.append(transformer_outputs.attentions) @@ -2252,8 +2374,6 @@ def _prepare_flow_postfix(self, ctx, prefix_kv_cache=None): postfix_moe_token_types = ctx["moe_token_types"][:, ctx["prefix_length"] :] postfix_input_ids = ctx["input_ids"][:, ctx["prefix_length"] :] - postfix_start_indices = None - postfix_end_indices = None padding_mask = None if prefix_kv_cache is not None: @@ -2288,16 +2408,6 @@ def _prepare_flow_postfix(self, ctx, prefix_kv_cache=None): layer_i ].values[:, :, : ctx["prefix_length"], :] - group_size = torch.zeros( - self.config.num_experts, dtype=torch.long, device="cpu" - ) - for i in range(self.config.num_experts): - group_size[i] = (postfix_moe_token_types == i).sum() - - # Calculate start and end indices for each expert group - postfix_start_indices = torch.cumsum(group_size, dim=0) - group_size - postfix_end_indices = torch.cumsum(group_size, dim=0) - pad_token_id = self.processor.tokenizer.pad_token_id padding_mask = ctx["input_ids"] == pad_token_id else: @@ -2334,20 +2444,22 @@ def _prepare_flow_postfix(self, ctx, prefix_kv_cache=None): _postfix_attention_mask, postfix_moe_token_types, postfix_input_ids, - postfix_start_indices, - postfix_end_indices, padding_mask, ) def _finalize_flow_action_output(self, action_trajectory, ctx): - predict_action = action_trajectory[-1] + predict_action_norm = action_trajectory[-1] + output = {"predict_action_norm": predict_action_norm} + if ctx["action_chunk"] is not None: + output["gt_action_norm"] = ctx["action_chunk"] + predict_action = unnormalize_data_with_virtual_tail( self.action_preprocessor.normalizer_action, - predict_action, + predict_action_norm, ctx["dataset_names"], ctx.get("dof_mask"), ) - output = {"predict_action": predict_action} + output["predict_action"] = predict_action if ctx["action_chunk"] is not None: output["gt_action"] = unnormalize_data_with_virtual_tail( self.action_preprocessor.normalizer_action, @@ -2377,8 +2489,6 @@ def prepare_inputs_for_generation( attention_mask=None, inputs_embeds=None, moe_token_types=None, - start_indices=None, - end_indices=None, cache_position=None, position_ids=None, use_cache=True, @@ -2480,8 +2590,6 @@ def prepare_inputs_for_generation( "position_ids": position_ids, "past_key_values": past_key_values, "moe_token_types": moe_token_types, - "start_indices": start_indices, - "end_indices": end_indices, "use_cache": use_cache, "attention_mask": attention_mask, "pixel_values": pixel_values, @@ -2867,7 +2975,7 @@ def convert_to_fused(state_dict): ) # ========================= - # Phase 4: sanity check(optional) + # Phase 4: sanity check（optional） # ========================= for k in fused.keys(): assert not any( @@ -2903,6 +3011,19 @@ def convert_to_mix_precision(self): if name in params_to_keep_float32: param.data = param.data.to(torch.float32) + def convert_to_pure_dmuon_mix_precision(self): + """Prepare live parameter dtypes before pure DMuon wrapping.""" + + params_to_keep_float32 = [] + for name, _ in self.named_parameters(): + if any(k in name for k in ["action_preprocessor", "action_processor"]): + params_to_keep_float32.append(name) + for name, param in self.named_parameters(): + if name in params_to_keep_float32: + param.data = param.data.to(torch.float32) + else: + param.data = param.data.to(torch.bfloat16) + def convert_to_fsdp( self, *, @@ -2911,6 +3032,7 @@ def convert_to_fsdp( offload_policy=None, reshard_after_forward: bool = True, use_dmuon: bool = False, + dmuon_owner_kwargs=None, ): """Wrap with FSDP2: per-decoder-layer + per-vision-block fully_shard with bf16 ``mp_policy``; layernorms and ``action_preprocessor`` leaves @@ -2921,10 +3043,10 @@ def convert_to_fsdp( ``mesh`` and ``mp_policy`` are required and produced by ``FSDPStrategy._build_fsdp2_layout``. The pre-migration FSDP1 wrap used ``ShardingStrategy.SHARD_GRAD_OP`` regardless of the yaml - config; the new path honors ``fsdp_sharding_strategy`` (default - ``full_shard`` -> ``reshard_after_forward=True``). Active runs that + config; the new path honors ``distributed.fsdp.sharding_strategy`` + (default ``full_shard`` -> ``reshard_after_forward=True``). Active runs that relied on the implicit SHARD_GRAD_OP throughput should set - ``distributed.fsdp_sharding_strategy: shard_grad_op`` explicitly. + ``distributed.fsdp.sharding_strategy: shard_grad_op`` explicitly. When ``use_dmuon=True``, DMuon owns selected trainable matrix weights before the remaining parameters are wrapped by FSDP2. @@ -2933,6 +3055,7 @@ def convert_to_fsdp( if use_dmuon: import dmuon + from wall_x.trainer.optimizer.dmuon.utils import dmuon_dedicate_mesh_args # DMuon keeps fp32 master weights for dedicated params; symmetric # FSDP2 params still use MixedPrecisionPolicy for bf16 compute. @@ -2945,31 +3068,47 @@ def convert_to_fsdp( for module in self.action_preprocessor.modules() if isinstance(module, nn.Linear) } + dmuon_sharded_adamw_module_ids = { + id(module) + for module in ( + getattr(self, "lm_head", None), + getattr(getattr(self, "model", None), "embed_tokens", None), + ) + if isinstance(module, (nn.Embedding, nn.Linear)) + } def hook_boundary(module: nn.Module) -> bool: # ActionProcessor helper methods call internal Linear layers # directly, bypassing ActionProcessor.forward(). Those internal # Linear layers need their own hooks. - if id(module) in action_preprocessor_linear_ids: + if ( + id(module) in action_preprocessor_linear_ids + or id(module) in dmuon_sharded_adamw_module_ids + ): return True return isinstance( module, ( Qwen2_5_VLDecoderLayer_with_MoE, Qwen2_5_VLVisionBlock, + Qwen2_5_VisionPatchEmbed, Qwen2_5_VLPatchMerger, type(self.action_preprocessor), ), ) + dmuon_mesh, dmuon_mesh_kwargs = dmuon_dedicate_mesh_args(mesh) assignment = dmuon.dedicate_params( self, - mesh, - predicate=_is_qwen25_dmuon_target_param, + dmuon_mesh, + predicate=_is_qwen25_dmuon_managed_param, + route_hint_fn=_qwen25_dmuon_route_hint, hook_boundary_predicate=hook_boundary, hook_boundary_strict=True, compute_dtype=torch.bfloat16, reshard_after_forward=reshard_after_forward, + **dmuon_mesh_kwargs, + **(dmuon_owner_kwargs or {}), ) logger.info( "[Qwen2.5 FSDP2 + DMuon] dedicate_params: %d params assigned", @@ -3017,17 +3156,17 @@ def is_norm_leaf(leaf_name: str, leaf: nn.Module) -> bool: fully_shard(module, **fp32_shard_kwargs) # Layernorm leaves inside decoder layers / vision blocks / model.norm - # - wrap with fp32 mp_policy BEFORE wrapping their parents so the + # — wrap with fp32 mp_policy BEFORE wrapping their parents so the # nested policy survives. Match-by-FQN mirrors the FSDP1 logic. - for module_name, module in list(self.named_modules()): - for child_name, child in list(module.named_children()): - if any( - k in child_name.lower() - for k in ("input_layernorm", "post_attention_layernorm") - ) or ("norm" in child_name.lower() and module_name.endswith("model")): - fully_shard_fp32_leaf_or_container( - f"{module_name}.{child_name}", child - ) + # for module_name, module in list(self.named_modules()): + # for child_name, child in list(module.named_children()): + # if any( + # k in child_name.lower() + # for k in ("input_layernorm", "post_attention_layernorm") + # ) or ("norm" in child_name.lower() and module_name.endswith("model")): + # fully_shard_fp32_leaf_or_container( + # f"{module_name}.{child_name}", child + # ) # action_preprocessor leaves: same fp32 treatment as layernorms. for name, module in list(self.named_modules()): @@ -3063,3 +3202,98 @@ def is_norm_leaf(leaf_name: str, leaf: nn.Module) -> bool: logger.info("[FSDP2] fully_shard root (bf16)") fully_shard(self, **shard_kwargs) return self + + def convert_to_dmuon( + self, *, mesh, reshard_after_forward: bool = False, dmuon_owner_kwargs=None + ): + """Install pure DMuon ownership without calling FSDP2 ``fully_shard``.""" + import dmuon + from wall_x.trainer.optimizer.dmuon.policy import ( + make_wallx_action_fp32_dmuon_param_policy_fn, + summarize_dmuon_param_policy, + ) + from wall_x.trainer.optimizer.dmuon.utils import dmuon_dedicate_mesh_args + + trainable_count = sum(1 for _, p in self.named_parameters() if p.requires_grad) + route_counts, route_examples = _validate_qwen25_pure_dmuon_route_plan( + self.named_parameters() + ) + logger.info( + "[Qwen2.5 Pure DMuon] route plan: counts=%s examples=%s", + route_counts, + _format_qwen25_route_examples(route_examples), + ) + + # DMuon keeps fp32 master weights for dedicated params. + for param in self.parameters(): + if param.requires_grad and param.dtype != torch.float32: + param.data = param.data.float() + + action_preprocessor_linear_ids = { + id(module) + for module in self.action_preprocessor.modules() + if isinstance(module, nn.Linear) + } + dmuon_sharded_adamw_module_ids = { + id(module) + for module in ( + getattr(self, "lm_head", None), + getattr(getattr(self, "model", None), "embed_tokens", None), + ) + if isinstance(module, (nn.Embedding, nn.Linear)) + } + + def hook_boundary(module: nn.Module) -> bool: + if ( + id(module) in action_preprocessor_linear_ids + or id(module) in dmuon_sharded_adamw_module_ids + ): + return True + return isinstance( + module, + ( + Qwen2_5_VLDecoderLayer_with_MoE, + Qwen2_5_VLVisionBlock, + Qwen2_5_VisionPatchEmbed, + Qwen2_5_VLPatchMerger, + type(self.action_preprocessor), + ), + ) + + dmuon_mesh, dmuon_mesh_kwargs = dmuon_dedicate_mesh_args(mesh) + param_policy_fn = make_wallx_action_fp32_dmuon_param_policy_fn( + _qwen25_pure_dmuon_route_hint + ) + policy_counts, policy_examples = summarize_dmuon_param_policy( + self.named_parameters(), param_policy_fn + ) + logger.info( + "[Qwen2.5 Pure DMuon] param policy: counts=%s examples=%s", + policy_counts, + _format_qwen25_route_examples(policy_examples), + ) + assignment = dmuon.dedicate_params( + self, + dmuon_mesh, + predicate=_is_qwen25_dmuon_trainable_param, + param_policy_fn=param_policy_fn, + hook_boundary_predicate=hook_boundary, + hook_boundary_strict=False, + reshard_after_forward=reshard_after_forward, + **dmuon_mesh_kwargs, + **(dmuon_owner_kwargs or {}), + ) + logger.info( + "[Qwen2.5 Pure DMuon] dedicate_params: %d/%d trainable params " + "assigned (reshard_after_forward=%s)", + len(assignment), + trainable_count, + reshard_after_forward, + ) + if len(assignment) != trainable_count: + logger.warning( + "[Qwen2.5 Pure DMuon] only %d/%d trainable params were assigned", + len(assignment), + trainable_count, + ) + return self diff --git a/wall_x/model/qact/tokenizer_mixin.py b/wall_x/model/qact/tokenizer_mixin.py index 221bcaa7..4a07f736 100644 --- a/wall_x/model/qact/tokenizer_mixin.py +++ b/wall_x/model/qact/tokenizer_mixin.py @@ -1,14 +1,22 @@ """Action tokenizer mixin for loading tokenizers and action mappings.""" -from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Tuple, Union +import fcntl +import hashlib +import json +import logging +import os -import numpy as np import torch -from transformers import AutoProcessor +import numpy as np +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, Tuple, Union +from transformers import AutoProcessor, PreTrainedTokenizerFast +from transformers.dynamic_module_utils import get_class_from_dynamic_module from wall_x.utils.constant import is_action_dataset_name +logger = logging.getLogger(__name__) + # Delay imports so missing optional packages do not fail at import time try: from spatial_tokenizer.spatial_tokenizer import SpatialActionTokenizer @@ -16,6 +24,7 @@ SpatialActionTokenizer = None + class ActionTokenizerMixin(ABC): """Base class for action tokenizers""" @@ -98,7 +107,7 @@ def build_action_mapper(self, processor) -> Optional[Dict]: Build action_mapper Args: - processor: HuggingFace processor,used to convert tokens to IDs + processor: HuggingFace processor，used to convert tokens to IDs Returns: action_mapper dictionary; format depends on tokenizer type @@ -309,11 +318,72 @@ def __init__(self): super().__init__() self._tokenizer_type = "fast" + @staticmethod + def _load_processor(action_tokenizer_path: str): + try: + return AutoProcessor.from_pretrained( + action_tokenizer_path, trust_remote_code=True + ) + except ValueError as exc: + # Transformers 5.x loads non-primary ProcessorMixin attributes from + # a named subfolder. Older fast-tokenizer assets keep the backing + # tokenizer beside processor_config.json, so load that flat layout + # explicitly while leaving all other AutoProcessor failures intact. + required_files = ("processor_config.json", "tokenizer.json") + if "Couldn't instantiate the backend tokenizer" not in str(exc) or not all( + os.path.isfile(os.path.join(action_tokenizer_path, name)) + for name in required_files + ): + raise + + logger.debug( + "AutoProcessor failed; falling back to legacy flat layout: %s", exc + ) + + with open( + os.path.join(action_tokenizer_path, "processor_config.json"), "r" + ) as f: + processor_config = json.load(f) + class_reference = processor_config.get("auto_map", {}).get("AutoProcessor") + if not class_reference: + raise ValueError( + f"processor_config.json in {action_tokenizer_path} " + "does not define auto_map.AutoProcessor" + ) from exc + + processor_class = get_class_from_dynamic_module( + class_reference, action_tokenizer_path, local_files_only=True + ) + bpe_tokenizer = PreTrainedTokenizerFast.from_pretrained( + action_tokenizer_path, local_files_only=True + ) + processor_kwargs = { + key: value + for key, value in processor_config.items() + if key not in {"auto_map", "processor_class"} + } + return processor_class(bpe_tokenizer, **processor_kwargs) + def load_tokenizer(self, config: dict, normalizer, device: str = "cpu") -> Any: """Load the fast tokenizer""" - self._tokenizer = AutoProcessor.from_pretrained( - config["action_tokenizer_path"], trust_remote_code=True - ) + action_tokenizer_path = config["action_tokenizer_path"] + lock_key = hashlib.sha1( + os.path.abspath(action_tokenizer_path).encode("utf-8") + ).hexdigest()[:16] + lock_path = f"/tmp/wallx_autoprocessor_{lock_key}.lock" + + # Serialize concurrent AutoProcessor(trust_remote_code=True) loads across + # ranks — they race on HF's remote-code module cache. fcntl.flock is a + # kernel-held advisory lock that auto-releases when the fd is closed or the + # process dies, so there is no stale-lock timeout / TOCTOU cleanup to get + # wrong; waiters queue in the kernel instead of polling. + lock_fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o644) + try: + fcntl.flock(lock_fd, fcntl.LOCK_EX) + self._tokenizer = self._load_processor(action_tokenizer_path) + finally: + fcntl.flock(lock_fd, fcntl.LOCK_UN) + os.close(lock_fd) self.action_normalizer = normalizer return self._tokenizer @@ -931,6 +1001,7 @@ def update_placeholder_mask(self, processor, prefix_length, input_ids): return {"ar_action_mask": ar_action_mask} + # ============================================================ # Factory function # ============================================================ diff --git a/wall_x/model/registry.py b/wall_x/model/registry.py index 250d8e81..ab7820ad 100644 --- a/wall_x/model/registry.py +++ b/wall_x/model/registry.py @@ -4,7 +4,7 @@ from typing import Dict, Type -# Global registry mapping model_type string -> adapter class +# Global registry mapping model_type string → adapter class _MODEL_REGISTRY: Dict[str, Type] = {} diff --git a/wall_x/trainer/adapters/__init__.py b/wall_x/trainer/adapters/__init__.py index af6752d3..b9646e01 100644 --- a/wall_x/trainer/adapters/__init__.py +++ b/wall_x/trainer/adapters/__init__.py @@ -5,6 +5,7 @@ from wall_x.trainer.adapters.base_adapter import ModelAdapter from wall_x.trainer.adapters.vla_model_adapter import VLAdapter + _ADAPTER_SPECS = { "qwen2_5": ("wall_x.model.qact.qwen2_5.adapter", "Qwen2_5Adapter"), } @@ -100,14 +101,6 @@ def format_adapter_error(model_type: str) -> str: return msg -def resolve_adapter(model_type: str) -> type[ModelAdapter]: - """Return a registered adapter class, loading lazily when needed.""" - adapter_cls = ADAPTER_REGISTRY.get(model_type) - if adapter_cls is None: - raise ValueError(format_adapter_error(model_type)) - return adapter_cls - - def __getattr__(name: str): model_type = _ADAPTER_CLASS_EXPORTS.get(name) if model_type is not None: @@ -124,5 +117,4 @@ def __getattr__(name: str): "ADAPTER_REGISTRY", "ADAPTER_IMPORT_ERRORS", "format_adapter_error", - "resolve_adapter", ] diff --git a/wall_x/trainer/adapters/base_adapter.py b/wall_x/trainer/adapters/base_adapter.py index 99f08032..007efb98 100644 --- a/wall_x/trainer/adapters/base_adapter.py +++ b/wall_x/trainer/adapters/base_adapter.py @@ -1,3 +1,4 @@ +import dataclasses import logging import os from abc import ABC, abstractmethod @@ -12,6 +13,33 @@ logger = logging.getLogger(__name__) +def _distributed_union_keys(local_keys): + keys = sorted(str(key) for key in local_keys) + if not (dist.is_available() and dist.is_initialized()): + return keys + world_size = dist.get_world_size() + if world_size <= 1: + return keys + gathered = [None for _ in range(world_size)] + dist.all_gather_object(gathered, keys) + union = set() + for rank_keys in gathered: + if isinstance(rank_keys, (list, tuple)): + union.update(str(key) for key in rank_keys) + return sorted(union) + + +def _channel_stats_device(*mappings): + for mapping in mappings: + if isinstance(mapping, dict): + for value in mapping.values(): + if torch.is_tensor(value): + return value.device + if torch.cuda.is_available(): + return torch.device("cuda", torch.cuda.current_device()) + return torch.device("cpu") + + def load_trainer_data_config(cfg): """Return optional backend-specific trainer data config. @@ -100,6 +128,17 @@ def get_transformer_layer_cls(self): """Return layer classes for FSDP transformer_auto_wrap_policy.""" ... + # ---- sequence parallelism ---- + def sp_sharded_param_prefixes(self) -> tuple[str, ...]: + """FQN prefixes of params trained on the SP sequence shard. + + These params see only this rank's sequence chunk under sequence + parallelism, so their grads are partial and get summed across the SP + group (see reduce_sp_partial_grads); params on the gathered full + sequence are excluded. Default: none (model has no SP-sharded params). + """ + return () + # ---- dataset ---- @abstractmethod def load_dataset(self, data_config, processor, rank, world_size, **kwargs): @@ -181,16 +220,16 @@ def collect_output_stats( per-dataset channel losses, and action accuracy metrics from the model outputs and merges them into *step_stats* in-place. - This is an optional hook - subclasses override it when their + This is an optional hook — subclasses override it when their output format carries extra fields. The default implementation is a no-op. Args: outputs: raw model outputs (dict or object). step_stats: dict being built by the trainer; mutated in-place. - reduce_tensor_fn: callable(tensor, average=True) -> scalar tensor, + reduce_tensor_fn: callable(tensor, average=True) → scalar tensor, wraps dist.all_reduce for the current parallelism. - true_gather_fn: callable(tensor) -> scalar | None, + true_gather_fn: callable(tensor) → scalar | None, gathers a value to rank-0 (returns None on non-main ranks or when input is None). tokenizer_mixin: optional ActionTokenizerMixin used to retrieve @@ -204,10 +243,10 @@ def collect_param_norms( """Compute per-group parameter L2 norms and merge into step_stats. Groups parameters by name keyword: - - "visual" -> visual_param_norm - - "action" -> action_expert_param_norm - - (others) -> org_vlm_param_norm - - all -> total_param_norm + - "visual" → visual_param_norm + - "action" → action_expert_param_norm + - (others) → org_vlm_param_norm + - all → total_param_norm When parameters are sharded across ranks (FSDP FULL_SHARD / HYBRID_SHARD), the local squared-sum is partial and must be @@ -215,13 +254,13 @@ def collect_param_norms( When parameters are replicated (DDP, or FSDP with SHARD_GRAD_OP / NO_SHARD), all_reduce would inflate the result by world_size. - This is an optional hook - the default implementation is a no-op. + This is an optional hook — the default implementation is a no-op. Args: model: the (possibly FSDP/DDP-wrapped) model. step_stats: dict being built by the trainer; mutated in-place. device: torch.device for tensor allocation. - reduce_tensor_fn: callable(tensor, average=False) -> tensor, + reduce_tensor_fn: callable(tensor, average=False) → tensor, wraps dist.all_reduce for the current parallelism. params_sharded: True when parameters are actually sharded across ranks (FSDP FULL_SHARD / HYBRID_SHARD). @@ -286,7 +325,8 @@ def convert_to_mix_precision_hint( model, *, device, - use_fsdp: bool, + use_fsdp: bool | None = None, + backend: str | None = None, log_fn=None, ): """Prepare *model* dtype + placement before FSDP/DDP wrap. @@ -294,15 +334,33 @@ def convert_to_mix_precision_hint( FSDP path: cast to fp32 on device; MixedPrecision policy handles the bf16 compute cast at forward time. DDP path: delegate to the model's ``convert_to_mix_precision`` method (manual bf16 conversion, since - DDP has no equivalent of MixedPrecision), then move to device. + DDP has no equivalent of MixedPrecision), then move to device. Pure + DMuon has its own forward dtype policy and must not rely on FSDP2 + MixedPrecision hooks. Adapters that handle placement themselves may override this hook. """ _log = log_fn or (lambda _msg: None) - if use_fsdp: + if backend is None: + backend = "fsdp2" if use_fsdp else "ddp" + + if backend == "fsdp2": model.to(dtype=torch.float32) _log("FSDP mode: all params fp32, MixedPrecision handles bf16 compute") model.to(device) + elif backend == "dmuon": + if hasattr(model, "convert_to_pure_dmuon_mix_precision"): + model.convert_to_pure_dmuon_mix_precision() + _log( + "Pure DMuon mode: model applied explicit live forward dtype policy" + ) + else: + model.to(dtype=torch.float32) + _log( + "Pure DMuon mode: all params fp32; no model-specific live dtype " + "policy is defined" + ) + model.to(device) else: model.convert_to_mix_precision() _log("DDP mode: params converted to bf16 (no MixedPrecision policy)") @@ -423,7 +481,7 @@ def predict( """Run an inference loop for the given ``prediction_type``. Subclasses implement the dispatch table (flow_action / ar_action / - dllm_action / text / ...). Default raises - only VLAdapter has + dllm_action / text / ...). Default raises — only VLAdapter has meaningful predict paths today. """ raise NotImplementedError( @@ -459,36 +517,80 @@ def collect_channel_loss_stats(outputs, prefix="", tokenizer_mixin=None): model output dict lacks ``channel_loss_dict``. """ stats = {} - channel_loss_dict = ( - outputs.get("channel_loss_dict") - if isinstance(outputs, dict) - else getattr(outputs, "channel_loss_dict", None) - ) - if channel_loss_dict is None: + + def output_field(name): + return ( + outputs.get(name) + if isinstance(outputs, dict) + else getattr(outputs, name, None) + ) + + channel_loss_dict = output_field("channel_loss_dict") or {} + channel_loss_count_dict = output_field("channel_loss_count_dict") or {} + if not channel_loss_dict and not (dist.is_available() and dist.is_initialized()): return stats - channel_loss_count_dict = ( - outputs.get("channel_loss_count_dict") - if isinstance(outputs, dict) - else getattr(outputs, "channel_loss_count_dict", None) - ) - for dataset_name_i in channel_loss_dict: - count_tensor = channel_loss_count_dict[dataset_name_i].clone() - loss_tensor = channel_loss_dict[dataset_name_i].detach().clone() - dist.all_reduce(count_tensor, op=dist.ReduceOp.SUM) - dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM) + + local_channel_keys = sorted(channel_loss_dict.keys()) + if dist.is_available() and dist.is_initialized(): + gathered_channel_keys = [None] * dist.get_world_size() + dist.all_gather_object(gathered_channel_keys, local_channel_keys) + channel_keys = sorted( + {key for rank_keys in gathered_channel_keys for key in (rank_keys or [])} + ) + else: + channel_keys = local_channel_keys + + metric_ref = None + for value in list(channel_loss_dict.values()) + list( + channel_loss_count_dict.values() + ): + if torch.is_tensor(value): + metric_ref = value + break + if metric_ref is None: + for name in ("loss", "total_loss", "flow_loss", "cross_entropy_loss"): + value = output_field(name) + if torch.is_tensor(value): + metric_ref = value + break + + def zero_metric_tensor(): + if metric_ref is None: + raise RuntimeError( + "Cannot reduce channel metrics without a tensor device reference " + "on this rank" + ) + return torch.zeros((), device=metric_ref.device, dtype=torch.float32) + + def channel_metric_tensor(key, source_dict): + value = source_dict.get(key) + if value is None: + return zero_metric_tensor() + return value.detach().clone().to(dtype=torch.float32) + + def reduce_metric_tensor(tensor): + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + return tensor + + for dataset_name_i in channel_keys: + count_tensor = channel_metric_tensor(dataset_name_i, channel_loss_count_dict) + loss_tensor = channel_metric_tensor(dataset_name_i, channel_loss_dict) + reduce_metric_tensor(count_tensor) + reduce_metric_tensor(loss_tensor) count_sum = count_tensor.item() if count_sum >= 0.5: stats[f"{prefix}channel_loss_{dataset_name_i}"] = ( loss_tensor.item() / count_sum ) - if "action_accuracy" in channel_loss_dict and tokenizer_mixin is not None: - acc_tensor = channel_loss_dict["action_accuracy"].detach().clone() - dist.all_reduce(acc_tensor, op=dist.ReduceOp.SUM) + if "action_accuracy" in channel_keys and tokenizer_mixin is not None: + acc_tensor = channel_metric_tensor("action_accuracy", channel_loss_dict) + reduce_metric_tensor(acc_tensor) world_size = dist.get_world_size() if dist.is_initialized() else 1 stats[f"{prefix}action_accuracy"] = acc_tensor.item() / world_size for key in tokenizer_mixin.get_accuracy_keys(): - if key != "action_accuracy" and key in channel_loss_dict: - rvq_acc_tensor = channel_loss_dict[key].detach().clone() - dist.all_reduce(rvq_acc_tensor, op=dist.ReduceOp.SUM) + if key != "action_accuracy" and key in channel_keys: + rvq_acc_tensor = channel_metric_tensor(key, channel_loss_dict) + reduce_metric_tensor(rvq_acc_tensor) stats[f"{prefix}{key}"] = rvq_acc_tensor.item() / world_size return stats diff --git a/wall_x/trainer/adapters/vla_model_adapter.py b/wall_x/trainer/adapters/vla_model_adapter.py index 955c9dce..4d8c0c4f 100644 --- a/wall_x/trainer/adapters/vla_model_adapter.py +++ b/wall_x/trainer/adapters/vla_model_adapter.py @@ -5,13 +5,42 @@ from tqdm import tqdm from wall_x.trainer.adapters.base_adapter import ModelAdapter -from wall_x.trainer.optimizer.dmuon import is_dmuon_model from wall_x.trainer.trainer_utils import ( compute_action_metrics, - load_qwen_pretrain_weight, load_wallx_processors_from_cfg, + load_qwen_pretrain_weight, save_text_results_to_file, ) +from wall_x.trainer.optimizer.dmuon import is_dmuon_model + + +_MISSING = object() + + +# Approximate peak per-GPU bf16 (dense, non-sparse) throughput in TFLOP/s, keyed +# by a substring of ``torch.cuda.get_device_name()``. Used ONLY as the MFU +# denominator. These are the vendor spec sheet numbers; real achievable peak is +# lower, so reported MFU is a conservative lower bound. Override via +# ``debug.mfu_peak_tflops`` for cards/precisions not covered here. +_GPU_BF16_TFLOPS = { + "A100": 312.0, + "A800": 312.0, + "H200": 989.0, + "H800": 989.0, + "H100": 989.0, # SXM; PCIe variant is special-cased in _resolve_peak_tflops + "H20": 148.0, + "A40": 149.0, + "A30": 165.0, + "A10": 125.0, + "A6000": 155.0, + "6000 Ada": 182.0, + "L40S": 362.0, + "L40": 181.0, + "L20": 119.5, + "L4": 121.0, + "4090": 165.0, + "V100": 125.0, # fp16 (Volta has no native bf16); MFU is approximate +} class VLAdapter(ModelAdapter): @@ -38,6 +67,8 @@ def __init__(self, *, cfg=None, logger=None, model_type=None): self._dmuon_module = None self._dmuon_import_checked = False self._dmuon_named_param_cache = None + self._mfu_missing_info_logged = False + self._mfu_compute_error_logged = False def _get_dmuon_module(self): if not self._dmuon_import_checked: @@ -114,10 +145,17 @@ def _build_processor_dict(self) -> dict: flat["batch_size_per_gpu"] = self.cfg.hyperparams.batch_size_per_gpu return flat - def load_processor(self, action_statistic_dof): + def load_processor(self, action_statistic_dof, ar_normalizer=None): + # ``ar_normalizer`` lets openloop scaling-eval inject the AR-mode + # normalizer for instruction-config override; falls back to the + # adapter's own ``normalizer_action`` (training path) when unset. processors_dict = load_wallx_processors_from_cfg( self.cfg, - normalizer=getattr(self, "normalizer_action", None), + normalizer=( + ar_normalizer + if ar_normalizer is not None + else getattr(self, "normalizer_action", None) + ), action_statistic_dof=action_statistic_dof, ) self.logger.info( @@ -156,12 +194,17 @@ def build_model_config(self): def create_model(self, processor, tokenizer_mixin, model_config): ModelClass, _ = self._get_model_and_config_class() use_selective_recompute = self.cfg.distributed.use_selective_recompute - return ModelClass( + model = ModelClass( model_config, processor, tokenizer_mixin=tokenizer_mixin, use_selective_recompute=use_selective_recompute, ) + if self.cfg.distributed.use_gradient_checkpointing: + model.gradient_checkpointing_enable() + if self.logger: + self.logger.info("Enabled gradient checkpointing for VLA model") + return model def load_weights(self, model, normalizer_action, normalizer_propri, **kwargs): import copy @@ -202,7 +245,7 @@ def _make_inputs_require_grad(module, input, output): def get_transformer_layer_cls(self): """Return the FSDP transformer-wrap layer classes for this variant. - Subclass override - see the per-variant adapter for the actual + Subclass override — see the per-variant adapter for the actual layer classes. """ raise NotImplementedError( @@ -224,17 +267,32 @@ def load_dataset(self, data_config, processor, rank, world_size, **kwargs): resume_state = None resume_batches = kwargs.get("resume_batches", 0) indices = kwargs.get("resume_indices") + sampler_state_path = kwargs.get("sampler_state_path") + dataset_state = kwargs.get("dataset_state") if indices is None and self.cfg.checkpoint.resume_from: checkpoint_path = self.cfg.checkpoint.resume_from if os.path.isfile(checkpoint_path): checkpoint_path = os.path.dirname(checkpoint_path) if os.path.isdir(checkpoint_path): indices = self._load_episode_indices(checkpoint_path, rank) - - if resume_batches or indices or kwargs.get("episode_container_checkpoint_path"): + candidate = os.path.join( + checkpoint_path, f"sampler_state_rank_{rank}.pt" + ) + if os.path.exists(candidate): + sampler_state_path = candidate + + if ( + resume_batches + or indices + or sampler_state_path + or dataset_state + or kwargs.get("episode_container_checkpoint_path") + ): resume_state = { "resume_batches": resume_batches, "indices": indices, + "sampler_state_path": sampler_state_path, + "dataset_state": dataset_state, "episode_container_checkpoint_path": kwargs.get( "episode_container_checkpoint_path" ), @@ -245,6 +303,8 @@ def load_dataset(self, data_config, processor, rank, world_size, **kwargs): world_size=world_size, processor=processor, tokenizer_mixin=kwargs.get("tokenizer_mixin"), + train_action_tokenizer=kwargs.get("train_action_tokenizer"), + val_action_tokenizer=kwargs.get("val_action_tokenizer"), normalizer_action=copy.deepcopy(kwargs.get("normalizer_action")), normalizer_propri=copy.deepcopy(kwargs.get("normalizer_propri")), model_config=kwargs.get("model_config"), @@ -300,35 +360,130 @@ def collect_output_stats( if log_flow_loss is not None: step_stats["flow_loss"] = log_flow_loss - channel_loss_dict = self.get_output_field(outputs, "channel_loss_dict") - channel_loss_count_dict = self.get_output_field( - outputs, "channel_loss_count_dict" + channel_loss_dict = self.get_output_field(outputs, "channel_loss_dict") or {} + channel_loss_count_dict = ( + self.get_output_field(outputs, "channel_loss_count_dict") or {} ) - if channel_loss_dict is not None: - for dataset_name_i in channel_loss_dict: - count_tensor = channel_loss_count_dict[dataset_name_i].clone() - loss_tensor = channel_loss_dict[dataset_name_i].detach().clone() + local_channel_keys = sorted(channel_loss_dict.keys()) + dist_ready = dist.is_available() and dist.is_initialized() + if dist_ready: + metric_ref = None + for value in list(channel_loss_dict.values()) + list( + channel_loss_count_dict.values() + ): + if torch.is_tensor(value): + metric_ref = value + break + if metric_ref is None: + metric_ref = flow_loss if torch.is_tensor(flow_loss) else None + if metric_ref is None: + metric_ref = ( + cross_entropy_loss if torch.is_tensor(cross_entropy_loss) else None + ) + + metric_device = ( + metric_ref.device + if metric_ref is not None + else ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else torch.device("cpu") + ) + ) + has_local_channel = torch.tensor( + [len(local_channel_keys)], + device=metric_device, + dtype=torch.int64, + ) + dist.all_reduce(has_local_channel, op=dist.ReduceOp.SUM) + has_any_channel = has_local_channel.item() > 0 + else: + metric_ref = None + has_any_channel = bool(local_channel_keys) + + if has_any_channel: + if dist_ready: + gathered_channel_keys = [None] * dist.get_world_size() + dist.all_gather_object(gathered_channel_keys, local_channel_keys) + channel_keys = sorted( + { + key + for rank_keys in gathered_channel_keys + for key in (rank_keys or []) + } + ) + else: + channel_keys = local_channel_keys + + if metric_ref is None: + for value in list(channel_loss_dict.values()) + list( + channel_loss_count_dict.values() + ): + if torch.is_tensor(value): + metric_ref = value + break + if metric_ref is None: + metric_ref = flow_loss if torch.is_tensor(flow_loss) else None + if metric_ref is None: + metric_ref = ( + cross_entropy_loss if torch.is_tensor(cross_entropy_loss) else None + ) + if metric_ref is None: + raise RuntimeError( + "Cannot reduce channel metrics without a tensor device reference " + "on this rank" + ) - dist.all_reduce(count_tensor, op=dist.ReduceOp.SUM) - dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM) + device = metric_ref.device + loss_count = len(channel_keys) + loss_count_pack = torch.zeros( + 2, loss_count, device=device, dtype=torch.float32 + ) + for idx, key in enumerate(channel_keys): + loss_value = channel_loss_dict.get(key) + if loss_value is not None: + loss_count_pack[0, idx] = loss_value.detach().to( + dtype=torch.float32 + ) + count_value = channel_loss_count_dict.get(key) + if count_value is not None: + loss_count_pack[1, idx] = count_value.detach().to( + dtype=torch.float32 + ) - cout_sum = count_tensor.item() + if dist_ready: + dist.all_reduce(loss_count_pack, op=dist.ReduceOp.SUM) + + loss_vec = loss_count_pack[0] + count_vec = loss_count_pack[1] + for idx, dataset_name_i in enumerate(channel_keys): + cout_sum = count_vec[idx].item() if cout_sum >= 0.5: step_stats[f"channel_loss_{dataset_name_i}"] = ( - loss_tensor.item() / cout_sum + loss_vec[idx].item() / cout_sum ) - if "action_accuracy" in channel_loss_dict and tokenizer_mixin is not None: - acc_tensor = channel_loss_dict["action_accuracy"].detach().clone() - dist.all_reduce(acc_tensor, op=dist.ReduceOp.SUM) - world_size = dist.get_world_size() if dist.is_initialized() else 1 - step_stats["action_accuracy"] = acc_tensor.item() / world_size - - for key in tokenizer_mixin.get_accuracy_keys(): - if key != "action_accuracy" and key in channel_loss_dict: - rvq_acc_tensor = channel_loss_dict[key].detach().clone() - dist.all_reduce(rvq_acc_tensor, op=dist.ReduceOp.SUM) - step_stats[key] = rvq_acc_tensor.item() / world_size + if "action_accuracy" in channel_keys and tokenizer_mixin is not None: + accuracy_keys = [ + key + for key in tokenizer_mixin.get_accuracy_keys() + if key in channel_keys + ] + if accuracy_keys: + acc_vec = torch.zeros( + len(accuracy_keys), device=device, dtype=torch.float32 + ) + for idx, key in enumerate(accuracy_keys): + acc_value = channel_loss_dict.get(key) + if acc_value is not None: + acc_vec[idx] = acc_value.detach().to(dtype=torch.float32) + if dist_ready: + dist.all_reduce(acc_vec, op=dist.ReduceOp.SUM) + world_size = dist.get_world_size() + else: + world_size = 1 + for idx, key in enumerate(accuracy_keys): + step_stats[key] = acc_vec[idx].item() / world_size def collect_param_norms( self, model, step_stats, device, reduce_tensor_fn, params_sharded=False @@ -336,10 +491,10 @@ def collect_param_norms( """Compute per-group parameter L2 norms for VL models. Groups: - "visual" -> visual_param_norm - "action" -> action_expert_param_norm - others -> org_vlm_param_norm - all -> total_param_norm + "visual" → visual_param_norm + "action" → action_expert_param_norm + others → org_vlm_param_norm + all → total_param_norm """ total_sq = torch.tensor(0.0, dtype=torch.float32, device=device) visual_sq = torch.tensor(0.0, dtype=torch.float32, device=device) @@ -462,6 +617,44 @@ def _extract_batch_info(batch): return info + @staticmethod + def _cfg_get(cfg, key, default=None): + if cfg is None: + return default + if isinstance(cfg, dict): + return cfg.get(key, default) + return getattr(cfg, key, default) + + @classmethod + def _cfg_first(cls, configs, key, default=None): + for cfg in configs: + value = cls._cfg_get(cfg, key, _MISSING) + if value is not _MISSING and value is not None: + return value + return default + + def _log_mfu_once(self, attr, message, *args, exc_info=False): + if getattr(self, attr, False): + return + setattr(self, attr, True) + if self.logger is not None: + self.logger.warning(message, *args, exc_info=exc_info) + + @staticmethod + def _get_visual_module(model): + visual = getattr(model, "visual", None) + if visual is not None: + return visual + inner_model = getattr(model, "model", None) + return getattr(inner_model, "visual", None) + + @classmethod + def _flops_config_sources(cls, model_config): + text_config = cls._cfg_get(model_config, "text_config") + vla_config = cls._cfg_get(model_config, "vla_config") + nested_vla_config = cls._cfg_get(text_config, "vla_config") + return (model_config, text_config, vla_config, nested_vla_config) + def _compute_detailed_flops(self, model, config): """Compute detailed per-module FLOPs for one training step. @@ -479,6 +672,8 @@ def _compute_detailed_flops(self, model, config): if model_config is None: return None + config_sources = self._flops_config_sources(model_config) + B = batch_info["batch_size"] S = batch_info["seq_length"] num_lang = batch_info["num_lang_tokens"] @@ -489,94 +684,270 @@ def _compute_detailed_flops(self, model, config): Nv = batch_info["vision_seq_length"] num_loss_tokens = batch_info["num_loss_tokens"] - H = model_config.hidden_size - num_layers = model_config.num_hidden_layers - num_heads = model_config.num_attention_heads - num_kv = model_config.num_key_value_heads - vocab_size = getattr(model_config, "padded_vocab_size", model_config.vocab_size) + H = self._cfg_first(config_sources, "hidden_size") + num_layers = self._cfg_first(config_sources, "num_hidden_layers") + num_heads = self._cfg_first(config_sources, "num_attention_heads") + num_kv = self._cfg_first(config_sources, "num_key_value_heads", num_heads) + vocab_size = self._cfg_first( + config_sources, + "padded_vocab_size", + self._cfg_first(config_sources, "vocab_size"), + ) + missing = [ + name + for name, value in ( + ("hidden_size", H), + ("num_hidden_layers", num_layers), + ("num_attention_heads", num_heads), + ("num_key_value_heads", num_kv), + ("vocab_size", vocab_size), + ) + if value is None + ] + if missing: + self._log_mfu_once( + "_mfu_missing_info_logged", + "[MFU] Skip MFU calculation: missing model config fields: %s", + ", ".join(missing), + ) + return None - use_mot = getattr(model_config, "attention_moe", False) - use_moe_mlp = getattr(model_config, "mlp_moe", False) - dim_inputs = getattr(model_config, "dim_inputs", (H, H)) + use_mot = self._cfg_first( + config_sources, + "attention_moe", + self._cfg_first(config_sources, "use_mot", False), + ) + use_moe_mlp = self._cfg_first(config_sources, "mlp_moe", False) + dim_inputs = self._cfg_first(config_sources, "dim_inputs", (H, H)) dim_lang, dim_act = dim_inputs grad_accum = self.cfg.hyperparams.gradient_accumulation_steps + opt_cfg = self.cfg.hyperparams.optimizer + action_expert_only = self._cfg_get(opt_cfg, "train_action_expert_only", False) vit_frozen = True - if hasattr(model, "visual"): - for p in model.visual.parameters(): + visual_module = self._get_visual_module(model) + if visual_module is not None: + for p in visual_module.parameters(): if p.requires_grad: vit_frozen = False break vit_mult = 1.0 if vit_frozen else 3.0 vit_fwd_flops = 0 - if hasattr(model_config, "vision_config") and Nv > 0: - vcfg = model_config.vision_config - Hv = vcfg.hidden_size - Iv = vcfg.intermediate_size - out_hidden = vcfg.out_hidden_size - depth_v = vcfg.depth - - F_qkv_v = 6 * Nv * Hv * Hv - F_o_v = 2 * Nv * Hv * Hv - F_mlp_v = 6 * Nv * Hv * Iv + 2 * Nv * Iv - F_linear_per_layer = F_qkv_v + F_o_v + F_mlp_v - - fullatt_set = set(getattr(vcfg, "fullatt_block_indexes", [])) - num_images = max(Nv // 768, 1) - si = Nv // num_images - sum_si_sq = num_images * si * si - sum_wi_sq = sum_si_sq // 16 - - for i in range(depth_v): - if i in fullatt_set: - vit_fwd_flops += F_linear_per_layer + 4 * Hv * sum_si_sq - else: - vit_fwd_flops += F_linear_per_layer + 4 * Hv * sum_wi_sq - - merge_unit = getattr(vcfg, "spatial_merge_size", 2) ** 2 - merger_hidden = Hv * merge_unit - Nv_merged = Nv // merge_unit - vit_fwd_flops += ( - 2 * Nv_merged * merger_hidden * merger_hidden - + 2 * Nv_merged * merger_hidden * out_hidden - ) + vcfg = self._cfg_get(model_config, "vision_config") + if vcfg is not None and Nv > 0: + Hv = self._cfg_get(vcfg, "hidden_size") + Iv = self._cfg_get(vcfg, "intermediate_size") + out_hidden = self._cfg_get(vcfg, "out_hidden_size") + depth_v = self._cfg_get(vcfg, "depth") + vision_missing = [ + name + for name, value in ( + ("vision_config.hidden_size", Hv), + ("vision_config.intermediate_size", Iv), + ("vision_config.out_hidden_size", out_hidden), + ("vision_config.depth", depth_v), + ) + if value is None + ] + if vision_missing: + self._log_mfu_once( + "_mfu_missing_info_logged", + "[MFU] Skip vision FLOPs: missing model config fields: %s", + ", ".join(vision_missing), + ) + else: + F_qkv_v = 6 * Nv * Hv * Hv + F_o_v = 2 * Nv * Hv * Hv + F_mlp_v = 6 * Nv * Hv * Iv + 2 * Nv * Iv + F_linear_per_layer = F_qkv_v + F_o_v + F_mlp_v + + fullatt_set = set(self._cfg_get(vcfg, "fullatt_block_indexes", [])) + num_images = max(Nv // 768, 1) + si = Nv // num_images + sum_si_sq = num_images * si * si + sum_wi_sq = sum_si_sq // 16 + + for i in range(depth_v): + if i in fullatt_set: + vit_fwd_flops += F_linear_per_layer + 4 * Hv * sum_si_sq + else: + vit_fwd_flops += F_linear_per_layer + 4 * Hv * sum_wi_sq + + merge_unit = self._cfg_get(vcfg, "spatial_merge_size", 2) ** 2 + merger_hidden = Hv * merge_unit + Nv_merged = Nv // merge_unit + vit_fwd_flops += ( + 2 * Nv_merged * merger_hidden * merger_hidden + + 2 * Nv_merged * merger_hidden * out_hidden + ) - kv_ratio = num_kv / num_heads - F_matmul_fwd = 4 * B * (S**2) * H + text_config = self._cfg_get(model_config, "text_config", model_config) + head_dim = self._cfg_get(text_config, "head_dim", H // num_heads) + layer_types = self._cfg_get(text_config, "layer_types") + D_qo = num_heads * head_dim + D_kv = num_kv * head_dim + + def _swiglu(N, d_in, d_ffn): + return 6 * N * d_in * d_ffn + N * d_ffn + + if layer_types is not None: + vla_cfg = self._cfg_get(model_config, "vla_config") + vla_dim_inputs = self._cfg_get(vla_cfg, "dim_inputs") + if vla_dim_inputs is not None and len(vla_dim_inputs) >= 2: + dim_act = vla_dim_inputs[1] + + ffn_base = self._cfg_get(text_config, "intermediate_size") + ffn_act = self._cfg_get(vla_cfg, "action_mlp_intermediate_size", ffn_base) + if ffn_base is None or ffn_act is None: + self._log_mfu_once( + "_mfu_missing_info_logged", + "[MFU] Skip MFU calculation: missing Qwen3.5 MLP hidden size", + ) + return None - if not use_mot: - F_attn_fwd = ( - 2 * N_total * H * H - + 4 * N_total * H * H * kv_ratio - + 2 * N_total * H * H - + F_matmul_fwd + nkh_lin = self._cfg_get(text_config, "linear_num_key_heads", num_heads) + nvh_lin = self._cfg_get(text_config, "linear_num_value_heads", num_heads) + kd_lin = self._cfg_get(text_config, "linear_key_head_dim", head_dim) + vd_lin = self._cfg_get(text_config, "linear_value_head_dim", head_dim) + key_dim_lin = nkh_lin * kd_lin + val_dim_lin = nvh_lin * vd_lin + + # Qwen3.5 has a base branch for all tokens and an action branch for + # action tokens. The attention matmul/recurrent body is shared. + FA_base_proj = ( + 2 * N_total * H * (D_qo * 2) + + 2 * N_total * H * D_kv + + 2 * N_total * H * D_kv + + 2 * N_total * D_qo * H ) - else: - F_attn_fwd = ( - N_lang * dim_lang * H * (4 + 4 * kv_ratio) - + N_act * dim_act * H * (4 + 4 * kv_ratio) - + F_matmul_fwd + FA_act_proj = ( + 2 * N_act * dim_act * (D_qo * 2) + + 2 * N_act * dim_act * D_kv + + 2 * N_act * dim_act * D_kv + + 2 * N_act * D_qo * dim_act ) + FA_matmul = 4 * B * num_heads * (S**2) * head_dim - if not use_moe_mlp: - ffn_hidden = model_config.intermediate_size - F_mlp_fwd = 6 * N_total * H * ffn_hidden + 2 * N_total * ffn_hidden - else: - hid_lang = model_config.experts[0]["intermediate_size"] - hid_act = model_config.experts[1]["intermediate_size"] - F_mlp_fwd = (6 * N_lang * dim_lang * hid_lang + 2 * N_lang * hid_lang) + ( - 6 * N_act * dim_act * hid_act + 2 * N_act * hid_act + LA_base_proj = ( + 2 * N_total * H * (key_dim_lin * 2 + val_dim_lin) + + 2 * N_total * H * val_dim_lin + + 2 * N_total * H * nvh_lin * 2 + + 2 * N_total * val_dim_lin * H ) + LA_act_proj = ( + 2 * N_act * dim_act * (key_dim_lin * 2 + val_dim_lin) + + 2 * N_act * dim_act * val_dim_lin + + 2 * N_act * dim_act * nvh_lin * 2 + + 2 * N_act * val_dim_lin * dim_act + ) + conv_kernel = self._cfg_get(text_config, "linear_conv_kernel_dim", 4) + LA_conv = 2 * N_total * (key_dim_lin * 2 + val_dim_lin) * conv_kernel + LA_recur = 4 * N_total * nvh_lin * kd_lin * vd_lin + + MLP_base = _swiglu(N_total, H, ffn_base) + MLP_act = ( + _swiglu(N_act, dim_act, ffn_act) + if use_moe_mlp + else _swiglu(N_act, H, ffn_base) + ) + + base_fwd = 0 + act_fwd = 0 + shared_matmul_fwd = 0 + for layer_type in layer_types: + if layer_type == "full_attention": + base_fwd += FA_base_proj + MLP_base + act_fwd += FA_act_proj + MLP_act + shared_matmul_fwd += FA_matmul + else: + base_fwd += LA_base_proj + MLP_base + act_fwd += LA_act_proj + MLP_act + shared_matmul_fwd += LA_conv + LA_recur + + decoder_fwd_flops = base_fwd + act_fwd + shared_matmul_fwd + else: + if not use_mot: + F_proj_fwd = ( + 2 * N_total * H * D_qo + + 2 * N_total * H * D_kv + + 2 * N_total * H * D_kv + + 2 * N_total * D_qo * H + ) + else: + F_proj_fwd = N_lang * dim_lang * ( + 2 * D_qo + 2 * D_kv + 2 * D_kv + 2 * D_qo + ) + N_act * dim_act * (2 * D_qo + 2 * D_kv + 2 * D_kv + 2 * D_qo) + # QK^T + attn·V over the full S x S map (2 matmuls, 2 FLOPs each). + # We count the full S^2 (Megatron/nanoGPT convention) rather than the + # causal half; for these VLA sequences attention is ~2% of step FLOPs + # so the choice shifts MFU by <1%. Kept as the standard "model FLOPs". + F_matmul_fwd = 4 * B * num_heads * (S**2) * head_dim + + ffn_hidden = self._cfg_first(config_sources, "intermediate_size") + if not use_moe_mlp: + if ffn_hidden is None: + self._log_mfu_once( + "_mfu_missing_info_logged", + "[MFU] Skip MFU calculation: missing model config field: intermediate_size", + ) + return None + F_mlp_fwd = _swiglu(N_total, H, ffn_hidden) + else: + experts = self._cfg_first(config_sources, "experts") + if experts and len(experts) >= 2: + hid_lang = self._cfg_get(experts[0], "intermediate_size") + hid_act = self._cfg_get(experts[1], "intermediate_size") + else: + hid_lang = ffn_hidden + hid_act = self._cfg_first( + config_sources, "action_mlp_intermediate_size" + ) + if hid_lang is None or hid_act is None: + self._log_mfu_once( + "_mfu_missing_info_logged", + "[MFU] Skip MFU calculation: missing MoE MLP hidden size", + ) + return None + F_mlp_fwd = _swiglu(N_lang, dim_lang, hid_lang) + _swiglu( + N_act, dim_act, hid_act + ) - decoder_fwd_flops = num_layers * (F_attn_fwd + F_mlp_fwd) + F_attn_fwd = F_proj_fwd + F_matmul_fwd + decoder_fwd_flops = num_layers * (F_attn_fwd + F_mlp_fwd) + base_fwd = decoder_fwd_flops + act_fwd = 0 + shared_matmul_fwd = 0 N_lm = num_loss_tokens if num_loss_tokens is not None else N_total lm_head_fwd_flops = 2 * N_lm * H * vocab_size + if action_expert_only and layer_types is not None: + decoder_total_flops = ( + 2.0 * base_fwd + 3.0 * shared_matmul_fwd + 3.0 * act_fwd + ) + lm_head_mult = 2.0 + elif action_expert_only: + trainable_ratio = 0.0 + if hasattr(model, "parameters"): + total_params = sum(p.numel() for p in model.parameters()) + trainable_params = sum( + p.numel() for p in model.parameters() if p.requires_grad + ) + trainable_ratio = ( + trainable_params / total_params if total_params > 0 else 0.0 + ) + decoder_total_flops = (2.0 + trainable_ratio) * decoder_fwd_flops + lm_head_mult = 2.0 + else: + decoder_total_flops = 3.0 * decoder_fwd_flops + lm_head_mult = 3.0 + total_flops = ( - vit_mult * vit_fwd_flops + 3.0 * decoder_fwd_flops + 3.0 * lm_head_fwd_flops + vit_mult * vit_fwd_flops + + decoder_total_flops + + lm_head_mult * lm_head_fwd_flops ) * grad_accum return { @@ -584,21 +955,66 @@ def _compute_detailed_flops(self, model, config): "vit_fwd_flops": vit_fwd_flops, "decoder_fwd_flops": decoder_fwd_flops, "lm_head_fwd_flops": lm_head_fwd_flops, + "decoder_base_fwd": base_fwd, + "decoder_act_fwd": act_fwd, + "decoder_matmul_fwd": shared_matmul_fwd, "vit_mult": vit_mult, + "lm_head_mult": lm_head_mult, } + @staticmethod + def _resolve_peak_tflops(override): + """Return (peak_bf16_tflops, source_description) for the MFU denominator. + + Priority: explicit ``debug.mfu_peak_tflops`` override > lookup by + ``torch.cuda.get_device_name()`` > a loud fallback. Never guesses + silently — an unlisted GPU returns a fallback tagged UNKNOWN so the log + makes clear the MFU may be off until the override is set. + """ + if override is not None and override > 0: + return float(override), f"config override ({float(override):g} TFLOP/s)" + + name = "" + try: + if torch.cuda.is_available(): + name = torch.cuda.get_device_name() + except Exception: + name = "" + + if name: + # H100 PCIe (756) has materially lower peak than SXM (989); the SXM + # table entry would otherwise over-credit it, deflating MFU. + if "H100" in name and "PCIe" in name: + return 756.0, f"{name} bf16 dense" + for key, val in _GPU_BF16_TFLOPS.items(): + if key in name: + return val, f"{name} bf16 dense" + + return 312.0, ( + f"UNKNOWN GPU {name!r}; assuming 312 TFLOP/s (A100-class) — set " + "debug.mfu_peak_tflops for an accurate MFU on this hardware" + ) + def compute_mfu(self, model, step_time_seconds): """Compute detailed Model FLOPs Utilization for VL transformer training.""" if step_time_seconds <= 0: return None try: - config = self._build_processor_dict() - flops_info = self._compute_detailed_flops(model, config) + flops_info = self._compute_detailed_flops(model, config=None) if flops_info is None: return None - gpu_peak_tflops = 312.0 # TODO: move to DebugConfig if needed - peak_flops = gpu_peak_tflops * 1e12 + override = self._cfg_get( + getattr(self.cfg, "debug", None), "mfu_peak_tflops", None + ) + peak_tflops, peak_source = self._resolve_peak_tflops(override) + self._log_mfu_once( + "_mfu_peak_logged", + "[MFU] denominator = %.1f TFLOP/s (%s)", + peak_tflops, + peak_source, + ) + peak_flops = peak_tflops * 1e12 # Data-parallel: each GPU independently processes its own micro-batch, # so per-GPU MFU = per_gpu_flops / (step_time * per_gpu_peak). @@ -610,7 +1026,13 @@ def compute_mfu(self, model, step_time_seconds): "mfu": mfu, "flops_per_step_T": total_flops / 1e12, } - except Exception: + except Exception as exc: + self._log_mfu_once( + "_mfu_compute_error_logged", + "[MFU] Failed to compute MFU; future MFU errors are suppressed: %s", + exc, + exc_info=True, + ) return None # ---- optional hooks ---- diff --git a/wall_x/trainer/fsdp_trainer/__init__.py b/wall_x/trainer/fsdp_trainer/__init__.py index 891c5865..4bd20fc4 100644 --- a/wall_x/trainer/fsdp_trainer/__init__.py +++ b/wall_x/trainer/fsdp_trainer/__init__.py @@ -3,18 +3,20 @@ from wall_x.trainer.fsdp_trainer.base_trainer import ( BaseDistributedTrainer, DistributedTrainer, - all_gather, - all_reduce, - barrier, + setup_distributed, cleanup_distributed, - get_local_rank, get_rank, get_world_size, + get_local_rank, is_main_process, - setup_distributed, + barrier, + all_reduce, + all_gather, ) + from wall_x.trainer.fsdp_trainer.fsdp_trainer import FSDPTrainer + __all__ = [ # Base classes and utilities "BaseDistributedTrainer", diff --git a/wall_x/trainer/fsdp_trainer/base_trainer.py b/wall_x/trainer/fsdp_trainer/base_trainer.py index 2ad2ea4f..d6e09eb3 100644 --- a/wall_x/trainer/fsdp_trainer/base_trainer.py +++ b/wall_x/trainer/fsdp_trainer/base_trainer.py @@ -3,32 +3,33 @@ No accelerate dependency - can be launched directly with torchrun """ -import contextlib +import os import gc +import torch import logging -import os +import contextlib +import torch.distributed as dist from abc import ABC, abstractmethod from datetime import timedelta -from typing import Any, Optional +from typing import Optional, Any -import torch -import torch.distributed as dist - -from wall_x.config.hyperparams_config import AdamWConfig from wall_x.config.train_config import TrainConfig -from wall_x.trainer.trainer_utils import seed_all +from wall_x.config.hyperparams_config import AdamWConfig from wall_x.utils.logger import DistributedLogger +from wall_x.trainer.trainer_utils import seed_all from wall_x.utils.timers import Timers def setup_distributed(): """Initialize distributed training environment. - NCCL watchdog default is 10 min - some dataset backends (e.g. lerobot's + NCCL watchdog default is 10 min — some dataset backends (e.g. lerobot's per-rank sequential LeRobotDataset indexing) can exceed that during the first build, causing barrier #2 to time out on the waiting rank. Allow the timeout to be raised via ``WALLX_DIST_TIMEOUT_MINUTES`` (default 30). """ + # Determinism (incl. the cublas workspace pin) is config-driven and applied + # in seed_all() from hyperparams.deterministic_algorithms — nothing to set here. if not dist.is_initialized(): timeout_min = int(os.environ.get("WALLX_DIST_TIMEOUT_MINUTES", "30")) dist.init_process_group( @@ -166,7 +167,7 @@ def __init__( ): self.cfg = train_config - # wandb Run object (metric recorder) - None on non-main ranks and + # wandb Run object (metric recorder) — None on non-main ranks and # when use_wandb is disabled. Distinct from self.logger below. self.wandb_run = wandb_run @@ -191,7 +192,12 @@ def __init__( self.seed = self.cfg.hyperparams.seed self.logger.info(f"seed {self.seed}") - seed_all(self.seed) + seed_all( + self.seed, + deterministic_algorithms=getattr( + self.cfg.hyperparams, "deterministic_algorithms", False + ), + ) self.start_epoch = 0 self.global_step = 0 self.micro_step = 0 @@ -217,7 +223,7 @@ def __init__( # Performance options self.nvtx = self.cfg.debug.nvtx - self.timers = Timers(log_level=0, log_option="minmax") + self.timers = Timers(log_level=-1, log_option="minmax") # Training hyperparameters self.max_grad_norm = opt.max_grad_norm @@ -257,16 +263,23 @@ def fit(self): """Main training loop""" barrier() + enable_validation = self.cfg.logging.enable_validation + # Optional: validate before training - if self.cfg.checkpoint.validate_first: + if enable_validation and self.cfg.checkpoint.validate_first: self.val_loop() barrier() if self.nvtx: torch.cuda.cudart().cudaProfilerStart() - # Profiler setup - enable_profiling = self.cfg.debug.profile + # Profiler setup. With profile_only_rank0 (default), other ranks fall + # through to a nullcontext — saves N-1 redundant trace dumps that would + # otherwise show the same per-step structure (the interesting + # cross-rank info is in collective kernels, not in the CPU annotations). + enable_profiling = self.cfg.debug.profile and ( + not self.cfg.debug.profile_only_rank0 or self.rank == 0 + ) profiler = ( self._setup_profiler() if enable_profiling else contextlib.nullcontext() ) @@ -284,8 +297,9 @@ def fit(self): self.save_checkpoint(epoch) # Validation - self.val_loop() - barrier() + if enable_validation: + self.val_loop() + barrier() gc.collect() torch.cuda.empty_cache() @@ -295,23 +309,31 @@ def fit(self): def _setup_profiler(self): """Setup PyTorch profiler""" - return torch.profiler.profile( - activities=[ + if os.environ.get("WALLX_PROFILE_CUDA_ONLY", "0") == "1": + activities = [torch.profiler.ProfilerActivity.CUDA] + else: + activities = [ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, - ], + ] + + return torch.profiler.profile( + activities=activities, schedule=torch.profiler.schedule( wait=self.cfg.debug.profile_wait_iters, warmup=self.cfg.debug.profile_warmup_iters, active=self.cfg.debug.profile_active_iters, + repeat=self.cfg.debug.profile_repeat, ), on_trace_ready=torch.profiler.tensorboard_trace_handler( self.cfg.debug.profile_save_path, worker_name=f"worker{self.rank}", + use_gzip=self.cfg.debug.profile_use_gzip, ), - record_shapes=True, - profile_memory=True, - with_stack=True, + record_shapes=self.cfg.debug.profile_record_shapes, + profile_memory=self.cfg.debug.profile_memory, + with_stack=self.cfg.debug.profile_with_stack, + with_modules=self.cfg.debug.profile_with_modules, ) def gather_tensor(self, tensor: torch.Tensor) -> torch.Tensor: @@ -374,10 +396,6 @@ def training_log( if not is_main_process(): return - # Prefer the cross-rank reduced loss (train_loss) from - # _current_step_stats so the console line matches what wall_x_2509's - # fsdp_trainer prints and what wandb sees. Falls back to the raw - # rank-0 local tensor when stats haven't been populated yet. _smoothed = getattr(self, "_current_step_stats", None) or {} _loss_to_print = _smoothed.get("train_loss", float(loss)) @@ -406,7 +424,7 @@ def true_gather(self, value): # CRITICAL: must clone() here, not just detach(). # detach() only disconnects the computation graph but still shares the - # same underlying storage. dist.all_reduce() is **in-place** - it + # same underlying storage. dist.all_reduce() is **in-place** — it # overwrites that storage with the sum across all ranks. If the caller # passes a tensor that is also referenced elsewhere (e.g. outputs["loss"] # and outputs["flow_loss"] pointing to the same scalar_loss), the diff --git a/wall_x/trainer/fsdp_trainer/checkpoint_io.py b/wall_x/trainer/fsdp_trainer/checkpoint_io.py index 77755ac8..43bbe72c 100644 --- a/wall_x/trainer/fsdp_trainer/checkpoint_io.py +++ b/wall_x/trainer/fsdp_trainer/checkpoint_io.py @@ -26,10 +26,23 @@ from wall_x.trainer.optimizer.dmuon import is_dmuon_model +_module_logger = logging.getLogger(__name__) + + def _noop_log(_msg: str, **_kw) -> None: pass +def _log(log_fn: Callable, message: str, *, level: int = logging.INFO) -> None: + if level >= logging.WARNING: + # log_fn callables (rank0 info/print wrappers) flatten the record level; + # serious messages (cold-started optimizer, dropped state) must also go + # out as real WARNING records so level-based filters/alerting see them. + _module_logger.log(level, message) + prefix = "WARNING: " if level >= logging.WARNING else "" + log_fn(f"{prefix}{message}") + + def _dict_section(config: dict, key: str) -> dict: section = config.get(key, {}) return section if isinstance(section, dict) else {} @@ -65,7 +78,7 @@ def _detect_legacy_fsdp1_format(checkpoint_path: str) -> bool: Legacy: ``model.safetensors`` is present (rank-0 full state), but optimizer state lives in per-rank ``optimizer_rank{N}.pt`` files - instead of a single ``optimizer.pt``. After the FSDP1 -> FSDP2 + instead of a single ``optimizer.pt``. After the FSDP1 → FSDP2 migration we cannot reshard those flat_param-keyed optim files into the new DTensor layout, so the legacy loader cold-starts the optimizer and warns. @@ -130,6 +143,7 @@ def save_checkpoint( grad_scaler=grad_scaler, log_fn=log_fn, ) + _save_rank_rng_state(ckpt_path, rank, write_legacy=is_main) if is_dmuon_model(model): _save_dmuon_state_dict( @@ -191,24 +205,18 @@ def _save_training_checkpoint_metadata( # Written unconditionally (previously sharded-mode only). torch.save({"world_size": world_size}, os.path.join(ckpt_path, "world_size.pth")) - # RNG state (rank-0 snapshot; every rank restores the same state on - # resume, matching the existing seed_all(seed) convention where every - # rank is seeded identically). - rng_state = { - "torch": torch.get_rng_state(), - "numpy": np.random.get_state(), - "python": random.getstate(), - } - if torch.cuda.is_available(): - rng_state["cuda"] = torch.cuda.get_rng_state() - torch.save(rng_state, os.path.join(ckpt_path, "rng_state.pt")) - # GradScaler state (fp16 AMP only). if grad_scaler is not None: torch.save(grad_scaler.state_dict(), os.path.join(ckpt_path, "grad_scaler.pt")) with open(os.path.join(ckpt_path, "config.yml"), "w") as f: - yaml.dump(config, f, default_flow_style=False, allow_unicode=True) + # safe_dump (no !!python/* tags) so the saved config loads with a plain + # yaml.safe_load; sanitize first to coerce tuples/objects to safe types. + from wall_x.config.loader import _sanitize_for_yaml + + yaml.safe_dump( + _sanitize_for_yaml(config), f, default_flow_style=False, allow_unicode=True + ) model_cfg = _dict_section(config, "model") data_cfg = _dict_section(config, "data") @@ -219,8 +227,9 @@ def _save_training_checkpoint_metadata( # Backward compatibility: fall back to pretrained_qwen_vl_path. processor_dir = config.get("pretrained_qwen_vl_path") if processor_dir is not None: - log_fn( - "WARNING: 'pretrained_qwen_vl_path' is deprecated for processor " + _log( + log_fn, + "'pretrained_qwen_vl_path' is deprecated for processor " "file copying, please use 'processor_path' instead.", level=logging.WARNING, ) @@ -266,6 +275,24 @@ def _save_training_checkpoint_metadata( ) +def _collect_rng_state() -> dict: + rng_state = { + "torch": torch.get_rng_state(), + "numpy": np.random.get_state(), + "python": random.getstate(), + } + if torch.cuda.is_available(): + rng_state["cuda"] = torch.cuda.get_rng_state() + return rng_state + + +def _save_rank_rng_state(ckpt_path: str, rank: int, *, write_legacy: bool) -> None: + rng_state = _collect_rng_state() + torch.save(rng_state, os.path.join(ckpt_path, f"rng_state_rank_{rank}.pt")) + if write_legacy: + torch.save(rng_state, os.path.join(ckpt_path, "rng_state.pt")) + + def _save_fsdp2_full_state_dict( *, ckpt_path: str, @@ -289,6 +316,7 @@ def _save_fsdp2_full_state_dict( if is_main: model_sd = _filter_frozen_prefixes(model_sd, frozen_prefixes, log_fn) + model_sd = _strip_gemm_padding(model_sd, model, log_fn) model_sd_out = _make_contiguous_and_clone_shared(model_sd) save_file(model_sd_out, os.path.join(ckpt_path, "model.safetensors")) torch.save(optim_sd, os.path.join(ckpt_path, "optimizer.pt")) @@ -314,18 +342,19 @@ def _save_dmuon_state_dict( import dmuon model_sd = dmuon.get_model_state_dict(model, cpu_offload=True, rank0_only=True) + optim_sd = dmuon.get_optimizer_state_dict( + model, optimizer, cpu_offload=True, rank0_only=True + ) + if is_main: model_sd = _filter_frozen_prefixes(model_sd, frozen_prefixes, log_fn) + model_sd = _strip_gemm_padding(model_sd, model, log_fn) model_sd = _make_contiguous_and_clone_shared(model_sd) save_file(model_sd, os.path.join(ckpt_path, "model.safetensors")) log_fn("[Checkpoint] Saved DMuon model state dict (full, rank0)") - - optim_sd = dmuon.get_optimizer_state_dict( - model, optimizer, cpu_offload=True, rank0_only=True - ) - if is_main: torch.save(optim_sd, os.path.join(ckpt_path, "optimizer.pt")) torch.save(lr_scheduler.state_dict(), os.path.join(ckpt_path, "scheduler.pt")) + log_fn("[Checkpoint] Saved DMuon optimizer and scheduler state dict (rank0)") del model_sd, optim_sd gc.collect() @@ -343,10 +372,12 @@ def _save_ddp_state_dict( ) -> None: if not is_main: return - model_state = ( - model.module.state_dict() if isinstance(model, DDP) else model.state_dict() - ) + base_model = model.module if isinstance(model, DDP) else model + model_state = base_model.state_dict() model_state = _filter_frozen_prefixes(model_state, frozen_prefixes, log_fn) + # DDP's state_dict() already triggers Qwen2_5_VLMLP._save_to_state_dict, so + # this is a no-op here (shape-guarded/idempotent); kept for parity/safety. + model_state = _strip_gemm_padding(model_state, base_model, log_fn) model_state_contiguous = _make_contiguous_and_clone_shared(model_state) save_file(model_state_contiguous, os.path.join(ckpt_path, "model.safetensors")) torch.save(optimizer.state_dict(), os.path.join(ckpt_path, "optimizer.pt")) @@ -357,6 +388,18 @@ def _save_dataset_state( *, ckpt_path: str, dataset, rank: int, log_fn: Callable ) -> None: """Save per-rank dataset resume state if the dataset supports it.""" + if hasattr(dataset, "dataset_state_dict"): + state = dataset.dataset_state_dict() + if state: + state_path = os.path.join(ckpt_path, f"dataset_state_rank_{rank}.pt") + torch.save(state, state_path) + log_fn(f"[Checkpoint] Saved dataset state: {state_path}") + return + + if hasattr(dataset, "save_checkpoint_state"): + dataset.save_checkpoint_state(ckpt_path, rank) + return + if not hasattr(dataset, "save_episode_containers"): return ec_path = os.path.join(ckpt_path, f"episode_containers_rank_{rank}.pkl") @@ -372,6 +415,36 @@ def _save_dataset_state( break +def load_dataset_state( + *, + checkpoint_path: str, + rank: int, + current_world_size: int, + log_fn: Optional[Callable] = None, +) -> Optional[dict]: + """Load per-rank dataset resume state if present and compatible.""" + log_fn = log_fn or _noop_log + state_path = os.path.join(checkpoint_path, f"dataset_state_rank_{rank}.pt") + if not os.path.exists(state_path): + return None + + state = torch.load(state_path, map_location="cpu", weights_only=False) + ckpt_world_size = state.get("world_size") + if ckpt_world_size is not None and int(ckpt_world_size) != int(current_world_size): + _log( + log_fn, + "[Checkpoint] Dataset state world_size mismatch: " + f"ckpt_world_size={ckpt_world_size}, current={current_world_size}. " + "Skipping exact dataset resume because sampler consumed cursors " + "are per-rank.", + level=logging.WARNING, + ) + return None + + log_fn(f"[Checkpoint] Loaded dataset state: {state_path}") + return state + + def _filter_frozen_prefixes( state_dict: dict, frozen_prefixes: Optional[Tuple[str, ...]], @@ -410,6 +483,81 @@ def _make_contiguous_and_clone_shared(state_dict: dict) -> dict: return out +def _strip_gemm_padding( + state_dict: dict, model: torch.nn.Module, log_fn: Callable = _noop_log +) -> dict: + """Slice GEMM-alignment padding out of a gathered full state dict in place. + + Modules like ``Qwen2_5_VLMLP`` pad ``intermediate_size`` to a multiple of + 128 (e.g. 3420 -> 3456) for cuBLAS efficiency and rely on their own + ``_save_to_state_dict`` hook to strip the padding on save. That hook only + fires for ``nn.Module.state_dict()`` (the DDP path); the FSDP2 + (``get_state_dict``) and DMuon (``get_model_state_dict``) paths gather + DTensors directly and bypass it, so the padded rows leak into the saved + checkpoint. This restores parity by stripping post-gather. + + The strip spec is derived from the model itself: any module exposing the + padding contract (``padded_intermediate_size`` != ``intermediate_size``). + This targets only the padded class and naturally excludes the non-padded + LLM ``Qwen2MLP`` and Qwen3.5 ``Qwen3_5VisionMLP``, which share the + ``gate_up_proj``/``down_proj`` names but carry no padding. The slice is + shape-guarded (only acts when a dim still equals the padded size), so it is + idempotent and safe to run even after the DDP hook already stripped. + """ + if not state_dict: + return state_dict + + def _resolve_key(suffix: str) -> Optional[str]: + if suffix in state_dict: + return suffix + # Tolerate a wrapper prefix on the gathered FQN (defensive; FSDP2/DMuon + # produce clean keys, so this rarely fires). + matches = [k for k in state_dict if k.endswith("." + suffix)] + return matches[0] if len(matches) == 1 else None + + stripped = 0 + for module_name, module in model.named_modules(): + orig = getattr(module, "intermediate_size", None) + padded = getattr(module, "padded_intermediate_size", None) + if orig is None or padded is None or orig == padded: + continue + + # gate_up_proj is a fused [gate; up] concat along dim 0; split, drop each + # half's padding rows, then re-concat. weight: [2*padded, h] -> [2*orig, h]; + # bias: [2*padded] -> [2*orig]. + for sub in ("gate_up_proj.weight", "gate_up_proj.bias"): + key = _resolve_key(f"{module_name}.{sub}") + if key is None: + continue + t = state_dict[key] + if t.shape[0] == 2 * padded: + gate, up = t.split([padded, padded], dim=0) + state_dict[key] = torch.cat([gate[:orig], up[:orig]], dim=0) + stripped += 1 + elif t.shape[0] != 2 * orig: + log_fn( + f"[Checkpoint] WARNING: unexpected shape for '{key}': " + f"{tuple(t.shape)} (orig={orig}, padded={padded}); not stripped." + ) + + # down_proj.weight: [h, padded] -> [h, orig]. + key = _resolve_key(f"{module_name}.down_proj.weight") + if key is not None: + t = state_dict[key] + if t.shape[1] == padded: + state_dict[key] = t[:, :orig] + stripped += 1 + elif t.shape[1] != orig: + log_fn( + f"[Checkpoint] WARNING: unexpected shape for '{key}': " + f"{tuple(t.shape)} (orig={orig}, padded={padded}); not stripped." + ) + + if stripped: + log_fn(f"[Checkpoint] Stripped GEMM padding from {stripped} tensor(s).") + return state_dict + + def _world_size_for_metadata() -> int: if dist.is_initialized(): return dist.get_world_size() @@ -420,45 +568,6 @@ def _world_size_for_metadata() -> int: # Load # ---------------------------------------------------------------------- -_EMBED_WEIGHT_KEYS = ( - "model.embed_tokens.weight", - "model.language_model.embed_tokens.weight", -) - - -def _maybe_resize_token_embeddings_for_load( - model: torch.nn.Module, - state_dict: dict, - log_fn: Optional[Callable] = None, -) -> None: - """Resize model embeddings when checkpoint vocab size differs from the model.""" - log_fn = log_fn or _noop_log - for key in _EMBED_WEIGHT_KEYS: - if key not in state_dict: - continue - ckpt_vocab = state_dict[key].shape[0] - embed = ( - model.get_input_embeddings() - if hasattr(model, "get_input_embeddings") - else None - ) - if embed is None: - return - cur_vocab = embed.weight.shape[0] - if cur_vocab == ckpt_vocab: - return - log_fn( - f"resize_token_embeddings from {cur_vocab} to {ckpt_vocab} " - f"to match checkpoint ({key})" - ) - if hasattr(model, "resize_token_embeddings"): - model.resize_token_embeddings(ckpt_vocab) - elif hasattr(model, "model") and hasattr( - model.model, "resize_token_embeddings" - ): - model.model.resize_token_embeddings(ckpt_vocab) - return - def load_weights( *, @@ -491,7 +600,6 @@ def load_weights( log_fn("The weights is fused, skipping conversion.") filtered_state_dict = _drop_checkpoint_normalizer_state(state_dict, log_fn) - _maybe_resize_token_embeddings_for_load(model, filtered_state_dict, log_fn=log_fn) if resume_config.get("try_harder", False): log_fn("### try harder to squeeze checkpoint weights into new model ###") @@ -512,17 +620,49 @@ def reshape_compatible_state_dict( """Pad / slice action-preprocessor weights to match target model shape.""" log_fn = log_fn or _noop_log out = {} + passed_through_unknown = 0 for name, param in state_dict.items(): if name not in model_sd: - log_fn(f"Not used parameter: {name}") + # Pass unknown keys through instead of dropping them: modules with + # name-changing _load_from_state_dict hooks (e.g. qwen3_5's fused + # qkv_proj/gate_up_proj/in_proj_zba consuming legacy split + # q_proj/gate_proj/... keys) only see the checkpoint tensor if it + # survives to load_state_dict. Dropping here would leave those + # fused projections at random init while strict=False training + # continues silently. Truly-unused keys are still tolerated by + # strict=False downstream. + log_fn( + f"Key not in model state_dict (passing through for " + f"_load_from_state_dict hooks / strict=False): {name}" + ) + out[name] = param + passed_through_unknown += 1 continue - if "action_preprocessor" in name or "action_processor" in name: + # action_preprocessor/action_processor: action-dim changes (pad only). + # embed_tokens/lm_head: vocab grow/shrink (append/remove-at-end), so the + # same first-N slice handles both. These are bare nn.Embedding/nn.Linear + # with no _load_from_state_dict hook, so a mismatch here would otherwise + # crash load_state_dict(strict=False) (size mismatch still raises). + if ( + "action_preprocessor" in name + or "action_processor" in name + or "embed_tokens" in name + or "lm_head" in name + ): if param.size() == model_sd[name].size(): out[name] = param continue size_0 = param.size() size_1 = model_sd[name].size() - if any(old_dim > new_dim for old_dim, new_dim in zip(size_0, size_1)): + # Action dims are concrete joints/channels and are NOT prefix-aligned, + # so a larger action checkpoint must be padded explicitly, not sliced + # (blind slicing would mismatch joints). Vocab is append/remove-at-end, + # so a larger vocab checkpoint is safely clipped to the model's first + # rows by the slice logic below. + is_action = "action_preprocessor" in name or "action_processor" in name + if is_action and any( + old_dim > new_dim for old_dim, new_dim in zip(size_0, size_1) + ): raise ValueError( f"Shape mismatch for '{name}': checkpoint shape {tuple(size_0)} is " f"larger than model shape {tuple(size_1)} in at least one dimension. " @@ -531,25 +671,54 @@ def reshape_compatible_state_dict( "in the dataset processor so checkpoint actions match the new action " "size." ) - out[name] = model_sd[name].clone() + # Keep model init for the new dims of: normalizer buffers (delta must + # stay 1.0, else div-by-zero) and vocab embeddings (proper resize init + # — a zero embedding row is not meaningful). Zero-init projection + # weights to avoid polluting existing hidden features with noise. + keep_model_init = ( + "normalizer" in name or "embed_tokens" in name or "lm_head" in name + ) + if keep_model_init: + out[name] = model_sd[name].clone() + else: + out[name] = torch.zeros_like(model_sd[name]) slices = [ slice(0, min(old_dim, new_dim)) for old_dim, new_dim in zip(size_0, size_1) ] out[name][slices] = param[slices] - log_fn( - f"Not match key: {name}, checkpoint shape: {tuple(size_0)}, " - f"model shape: {tuple(size_1)}. Filled checkpoint weights into the first " - f"{[s.stop for s in slices]} dims, remaining dims keep model init weights." - ) - else: - if param.size() == model_sd[name].size(): - out[name] = param + clipped = any(old_dim > new_dim for old_dim, new_dim in zip(size_0, size_1)) + init_desc = "model init" if keep_model_init else "zero-initialized" + if clipped: + log_fn( + f"Not match key: {name}, checkpoint {tuple(size_0)} -> model " + f"{tuple(size_1)}. Clipped checkpoint to first " + f"{[s.stop for s in slices]} dims (rows beyond dropped)." + ) else: log_fn( - f"Skipping '{name}': checkpoint shape {tuple(param.size())} " - f"!= model shape {tuple(model_sd[name].size())}" + f"Not match key: {name}, checkpoint {tuple(size_0)} -> model " + f"{tuple(size_1)}. Filled first {[s.stop for s in slices]} dims, " + f"remaining dims {init_desc}." + ) + else: + # Always pass through, even on shape mismatch. + # Modules with a custom _load_from_state_dict hook (e.g. + # Qwen2_5_VLMLP pads ViT MLP 3420->3456 for cuBLAS alignment) + # must receive the original checkpoint tensor to do their own + # adaptation. Skipping here would bypass that hook entirely. + out[name] = param + if param.size() != model_sd[name].size(): + log_fn( + f"Shape mismatch '{name}': checkpoint {tuple(param.size())} " + f"vs model {tuple(model_sd[name].size())}. " + f"Passing through to _load_from_state_dict hook." ) + if passed_through_unknown: + log_fn( + f"Passed through {passed_through_unknown} key(s) not in model " + f"state_dict (for _load_from_state_dict hooks / strict=False)" + ) return out @@ -675,7 +844,8 @@ def _read_ckpt_world_size(checkpoint_path: str, log_fn: Callable): try: return int(torch.load(ws_path, map_location="cpu")["world_size"]) except (KeyError, ValueError, RuntimeError, TypeError) as e: - log_fn( + _log( + log_fn, f"world_size.pth unreadable ({e}); assuming same ws.", level=logging.WARNING, ) @@ -691,6 +861,174 @@ def _strip_fused_flags(osd: dict) -> None: pg.pop("foreach", None) +def _distributed_int_sum_batch(*values: int) -> list: + """Sum several ints across ranks in ONE all_reduce (fewer collective rounds).""" + if not dist.is_initialized(): + return [int(v) for v in values] + device = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else torch.device("cpu") + ) + tensor = torch.tensor([int(v) for v in values], dtype=torch.int64, device=device) + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + return tensor.tolist() + + +def _format_name_preview(names, limit: int = 8) -> str: + names = sorted(str(name) for name in names) + if not names: + return "[]" + suffix = "" if len(names) <= limit else f", ... (+{len(names) - limit})" + return "[" + ", ".join(names[:limit]) + suffix + "]" + + +def _log_dmuon_optimizer_resume_diagnostics( + *, + model, + optimizer, + optim_sd: dict, + log_fn: Callable, +) -> None: + """Audit DMuon optimizer restore coverage without changing optimizer state.""" + try: + from dmuon.checkpoint import ( + _compute_dedicated_fqns, + _owns_dedicated_adamw_optimizer_state, + _owns_dedicated_optimizer_state, + ) + except Exception as exc: + _log( + log_fn, + f"[Checkpoint][DMuon] Resume diagnostics unavailable: {exc!r}", + level=logging.WARNING, + ) + return + + rank = dist.get_rank() if dist.is_initialized() else 0 + is_main = rank == 0 + + dp_fqns = _compute_dedicated_fqns(model) + dedicated_fqn_set = set(dp_fqns.values()) + dp_to_muon_group = getattr(optimizer, "_dp_to_muon_group_idx", {}) or {} + dp_to_adamw_group = getattr(optimizer, "_dp_to_adamw_group_idx", {}) or {} + + current_muon_fqns = { + fqn for dp, fqn in dp_fqns.items() if id(dp) in dp_to_muon_group + } + current_dedicated_adamw_fqns = { + fqn for dp, fqn in dp_fqns.items() if id(dp) in dp_to_adamw_group + } + current_fsdp_fqns = { + name + for name, _param in model.named_parameters() + if name not in dedicated_fqn_set + } + + saved_dedicated_fqns = set((optim_sd.get("dedicated") or {}).keys()) + saved_dedicated_adamw_fqns = set((optim_sd.get("dedicated_adamw") or {}).keys()) + saved_fsdp_fqns = set((optim_sd.get("fsdp") or {}).keys()) + + fqn_to_dp = {fqn: dp for dp, fqn in dp_fqns.items()} + local_muon_expected = 0 + local_muon_loaded = 0 + local_muon_missing = [] + for fqn in current_muon_fqns & saved_dedicated_fqns: + dp = fqn_to_dp[fqn] + if not _owns_dedicated_optimizer_state(dp): + continue + local_muon_expected += 1 + state = optimizer.state.get(id(dp), {}) + if "momentum_buffer" in state: + local_muon_loaded += 1 + else: + local_muon_missing.append(fqn) + + local_dedicated_adamw_expected = 0 + local_dedicated_adamw_loaded = 0 + local_dedicated_adamw_missing = [] + for fqn in current_dedicated_adamw_fqns & saved_dedicated_adamw_fqns: + dp = fqn_to_dp[fqn] + if not _owns_dedicated_adamw_optimizer_state(dp): + continue + local_dedicated_adamw_expected += 1 + state = optimizer.state.get(id(dp), {}) + if state: + local_dedicated_adamw_loaded += 1 + else: + local_dedicated_adamw_missing.append(fqn) + + fsdp_param_by_fqn = { + name: param + for name, param in model.named_parameters() + if name not in dedicated_fqn_set + } + local_fsdp_expected = len(current_fsdp_fqns & saved_fsdp_fqns) + local_fsdp_loaded = 0 + local_fsdp_missing = [] + for fqn in current_fsdp_fqns & saved_fsdp_fqns: + param = fsdp_param_by_fqn[fqn] + state = optimizer.state.get(param, {}) + if state: + local_fsdp_loaded += 1 + else: + local_fsdp_missing.append(fqn) + + # One batched all_reduce instead of six separate collectives. + ( + global_muon_expected, + global_muon_loaded, + global_dedicated_adamw_expected, + global_dedicated_adamw_loaded, + global_fsdp_expected, + global_fsdp_loaded, + ) = _distributed_int_sum_batch( + local_muon_expected, + local_muon_loaded, + local_dedicated_adamw_expected, + local_dedicated_adamw_loaded, + local_fsdp_expected, + local_fsdp_loaded, + ) + + if is_main: + log_fn( + "[Checkpoint][DMuon] optimizer resume coverage: " + f"muon_saved={len(saved_dedicated_fqns)} " + f"muon_current={len(current_muon_fqns)} " + f"muon_owner_loaded={global_muon_loaded}/{global_muon_expected}; " + f"dedicated_adamw_saved={len(saved_dedicated_adamw_fqns)} " + f"dedicated_adamw_current={len(current_dedicated_adamw_fqns)} " + f"dedicated_adamw_owner_loaded={global_dedicated_adamw_loaded}/" + f"{global_dedicated_adamw_expected}; " + f"fsdp_saved={len(saved_fsdp_fqns)} " + f"fsdp_current={len(current_fsdp_fqns)} " + f"fsdp_local_loaded_sum={global_fsdp_loaded}/{global_fsdp_expected}" + ) + log_fn( + "[Checkpoint][DMuon] optimizer resume FQN diff: " + f"muon_missing_current={_format_name_preview(current_muon_fqns - saved_dedicated_fqns)} " + f"muon_unexpected_saved={_format_name_preview(saved_dedicated_fqns - current_muon_fqns)} " + f"dedicated_adamw_missing_current=" + f"{_format_name_preview(current_dedicated_adamw_fqns - saved_dedicated_adamw_fqns)} " + f"dedicated_adamw_unexpected_saved=" + f"{_format_name_preview(saved_dedicated_adamw_fqns - current_dedicated_adamw_fqns)} " + f"fsdp_missing_current={_format_name_preview(current_fsdp_fqns - saved_fsdp_fqns)} " + f"fsdp_unexpected_saved={_format_name_preview(saved_fsdp_fqns - current_fsdp_fqns)}" + ) + + if local_muon_missing or local_dedicated_adamw_missing or local_fsdp_missing: + _log( + log_fn, + "[Checkpoint][DMuon] optimizer resume local missing state: " + f"rank={rank} " + f"muon={_format_name_preview(local_muon_missing)} " + f"dedicated_adamw={_format_name_preview(local_dedicated_adamw_missing)} " + f"fsdp={_format_name_preview(local_fsdp_missing)}", + level=logging.WARNING, + ) + + def _drop_checkpoint_normalizer_state(state_dict: dict, log_fn: Callable) -> dict: """Drop saved normalizer buffers so current-run stats stay authoritative.""" normalizer_state_prefixes = ( @@ -728,7 +1066,7 @@ def _load_fsdp2_or_dmuon_full( DMuon path: every rank MUST load the full state dict from disk. ``dmuon.set_model_state_dict`` / ``set_optimizer_state_dict`` iterate - ``fqn_to_dp`` and skip any FQN not present in the provided dict - so + ``fqn_to_dp`` and skip any FQN not present in the provided dict — so if only rank 0 has the data, new owners on other ranks (after a cross-ws resume) silently miss their assigned params and ``_owned_data`` stays at fresh-init (random weights), producing a @@ -788,7 +1126,17 @@ def _load_fsdp2_or_dmuon_full( optim_sd = torch.load(optim_path, map_location="cpu", weights_only=False) _strip_fused_flags(optim_sd) elif is_main: - optim_sd = torch.load(optim_path, map_location="cpu", weights_only=False) + # FSDP2 optimizer restore reshards the full optimizer state via the + # distributed backend. Loading rank-0 state on CUDA avoids a CPU + # DTensor/backend path that fails under NCCL-only process groups. + map_location = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else "cpu" + ) + optim_sd = torch.load( + optim_path, map_location=map_location, weights_only=False + ) _strip_fused_flags(optim_sd) else: optim_sd = {} @@ -799,10 +1147,16 @@ def _load_fsdp2_or_dmuon_full( dmuon.set_optimizer_state_dict(model, optimizer, optim_sd) log_fn("[Checkpoint] DMuon optimizer state loaded.") + _log_dmuon_optimizer_resume_diagnostics( + model=model, + optimizer=optimizer, + optim_sd=optim_sd, + log_fn=log_fn, + ) else: options = StateDictOptions( full_state_dict=True, - cpu_offload=True, + cpu_offload=False, broadcast_from_rank0=True, ) set_optimizer_state_dict( @@ -816,8 +1170,14 @@ def _load_fsdp2_or_dmuon_full( "(reshard via broadcast)." ) except (ValueError, RuntimeError, TypeError, KeyError) as e: + # CUDA OOM is not a layout incompatibility: rank 0 materializes the + # full unsharded optimizer state on GPU here, and swallowing an OOM + # would silently cold-start momentum mid-run. Fail loudly instead. + if isinstance(e, RuntimeError) and "out of memory" in str(e).lower(): + raise # Most likely a legacy FSDP1 consolidated OSD (flat_param-specific). - log_fn( + _log( + log_fn, f"[Checkpoint] Optimizer state incompatible with current " f"layout ({e!r}); cold-starting optimizer. Model weights are " f"loaded; momentum resets to zero.", @@ -826,7 +1186,8 @@ def _load_fsdp2_or_dmuon_full( del optim_sd gc.collect() else: - log_fn( + _log( + log_fn, "[Checkpoint] optimizer.pt not found; cold-starting optimizer.", level=logging.WARNING, ) @@ -839,7 +1200,8 @@ def _load_fsdp2_or_dmuon_full( ) log_fn("[Checkpoint] Scheduler state loaded.") except (ValueError, RuntimeError, TypeError, KeyError) as e: - log_fn( + _log( + log_fn, f"Scheduler load failed ({e}); keeping fresh state.", level=logging.WARNING, ) @@ -863,7 +1225,8 @@ def _resume_ddp_optimizer_scheduler( optimizer.load_state_dict(optim_sd) log_fn("[Checkpoint] Optimizer state loaded (single file).") except (ValueError, RuntimeError) as e: - log_fn( + _log( + log_fn, f"Failed to load optimizer state dict, " f"optimizer will be re-initialized. Error: {e}", level=logging.WARNING, @@ -877,7 +1240,8 @@ def _resume_ddp_optimizer_scheduler( ) log_fn(f"[Checkpoint] Scheduler state loaded ({sched_path}).") except (ValueError, RuntimeError, KeyError, TypeError) as e: - log_fn( + _log( + log_fn, f"Scheduler load failed ({e}); keeping fresh state.", level=logging.WARNING, ) @@ -912,10 +1276,11 @@ def _load_legacy_fsdp1_full( model_class=model_class, log_fn=log_fn, ) - log_fn( + _log( + log_fn, f"Legacy FSDP1 ckpt detected at {checkpoint_path}: model weights " f"loaded, but per-rank optimizer files cannot reshard into FSDP2. " - f"Optimizer state DROPPED - it will warm up from scratch. For long " + f"Optimizer state DROPPED — it will warm up from scratch. For long " f"resume runs prefer a fresh FSDP2 ckpt; for short runs the loss " f"bump is usually negligible.", level=logging.WARNING, @@ -923,10 +1288,14 @@ def _load_legacy_fsdp1_full( def _resume_rng(checkpoint_path: str, log_fn: Callable) -> None: - path = os.path.join(checkpoint_path, "rng_state.pt") + rank = dist.get_rank() if dist.is_initialized() else 0 + rank_path = os.path.join(checkpoint_path, f"rng_state_rank_{rank}.pt") + legacy_path = os.path.join(checkpoint_path, "rng_state.pt") + path = rank_path if os.path.exists(rank_path) else legacy_path if not os.path.exists(path): - log_fn( - "rng_state.pt not found; RNG stays at seed_all() state.", + _log( + log_fn, + f"{os.path.basename(path)} not found; RNG stays at seed_all() state.", level=logging.WARNING, ) return @@ -940,9 +1309,10 @@ def _resume_rng(checkpoint_path: str, log_fn: Callable) -> None: np.random.set_state(sd["numpy"]) if "python" in sd: random.setstate(sd["python"]) - log_fn("[Checkpoint] RNG state restored.") + log_fn(f"[Checkpoint] RNG state restored from {os.path.basename(path)}.") except (RuntimeError, ValueError, TypeError, KeyError) as e: - log_fn( + _log( + log_fn, f"RNG load failed ({e}); keeping current RNG state.", level=logging.WARNING, ) @@ -951,7 +1321,8 @@ def _resume_rng(checkpoint_path: str, log_fn: Callable) -> None: def _resume_grad_scaler(checkpoint_path: str, grad_scaler, log_fn: Callable) -> None: path = os.path.join(checkpoint_path, "grad_scaler.pt") if not os.path.exists(path): - log_fn( + _log( + log_fn, "grad_scaler.pt not found; GradScaler stays fresh.", level=logging.WARNING, ) @@ -962,7 +1333,8 @@ def _resume_grad_scaler(checkpoint_path: str, grad_scaler, log_fn: Callable) -> ) log_fn("[Checkpoint] GradScaler state restored.") except (RuntimeError, ValueError, TypeError, KeyError) as e: - log_fn( + _log( + log_fn, f"GradScaler load failed ({e}); keeping fresh state.", level=logging.WARNING, ) diff --git a/wall_x/trainer/fsdp_trainer/distribution_strategy.py b/wall_x/trainer/fsdp_trainer/distribution_strategy.py index 65ad4cd3..0d11370c 100644 --- a/wall_x/trainer/fsdp_trainer/distribution_strategy.py +++ b/wall_x/trainer/fsdp_trainer/distribution_strategy.py @@ -2,6 +2,7 @@ from __future__ import annotations +import inspect import logging import os from abc import ABC, abstractmethod @@ -18,6 +19,8 @@ ) from torch.nn.parallel import DistributedDataParallel as DDP +from wall_x.distributed import get_sp_world_size + _logger = logging.getLogger(__name__) @@ -34,14 +37,26 @@ class FSDP2Layout(NamedTuple): reshard_after_forward: bool +class DMuonLayout(NamedTuple): + """Pure DMuon wrap configuration produced by ``DMuonStrategy``.""" + + mesh: DeviceMesh + reshard_after_forward: bool + + class DistributionStrategy(ABC): - """Common interface for FSDP2 / DDP wrapping + gradient operations.""" + """Common interface for FSDP2 / pure DMuon / DDP gradient operations.""" _cfg: dict last_grad_clip_stats: Optional[dict] = None def _maybe_enable_grad_ckpt(self, model: torch.nn.Module) -> None: - """Call ``gradient_checkpointing_enable`` if the model exposes it.""" + """Call ``gradient_checkpointing_enable`` if the model exposes it. + + Passes ``num_layers=cfg["gradient_checkpoint_num_layers"]`` when the + model's enable signature accepts it (partial-ckpt support); otherwise + falls back to a no-arg call (legacy models = ckpt all layers). + """ if not self._cfg.get("use_gradient_checkpointing", False): return method = getattr(model, "gradient_checkpointing_enable", None) @@ -56,10 +71,24 @@ def _maybe_enable_grad_ckpt(self, model: torch.nn.Module) -> None: type(model).__name__, ) return - method() + n_layers = self._cfg.get("gradient_checkpoint_num_layers", None) + ckpt_vision = self._cfg.get("gradient_checkpoint_vision_tower", True) + sig_params = inspect.signature(method).parameters + kwargs = {} + if "num_layers" in sig_params: + kwargs["num_layers"] = n_layers + if "ckpt_vision_tower" in sig_params: + kwargs["ckpt_vision_tower"] = ckpt_vision + method(**kwargs) @abstractmethod - def wrap(self, model: torch.nn.Module) -> torch.nn.Module: + def wrap( + self, + model: torch.nn.Module, + *, + use_dedicated_params: bool = False, + dmuon_owner_kwargs: Optional[dict] = None, + ) -> torch.nn.Module: """Wrap *model* with the active distribution strategy.""" @abstractmethod @@ -86,9 +115,8 @@ def no_sync(self, model: torch.nn.Module): def params_sharded(self) -> bool: """Whether parameters are actually sharded across ranks. - FSDP2 always shards (every ``fully_shard`` unit is sharded across - the mesh). DDP never shards. Used by adapters to decide whether - per-rank norm contributions need an all-reduce. + Used by adapters to decide whether per-rank norm contributions need + an all-reduce. """ @@ -126,7 +154,7 @@ def _build_fsdp2_layout(self) -> FSDP2Layout: groups, a trainer-level global process group, the matching mixed-precision policy, offload policy, and reshard_after_forward flag. """ - name = self._cfg.get("fsdp_sharding_strategy", "full_shard") + name = self._cfg.get("sharding_strategy", "full_shard") world_size = dist.get_world_size() if dist.is_initialized() else 1 dp_pg: Optional[dist.ProcessGroup] = None shard_pg: Optional[dist.ProcessGroup] = None @@ -134,18 +162,18 @@ def _build_fsdp2_layout(self) -> FSDP2Layout: # 2D mesh for HSDP variants; 1D for full_shard / shard_grad_op. if name in ("hybrid_shard", "_hybrid_shard_zero2"): - replicate_size = self._cfg.get("fsdp_hsdp_replicate_size") or int( + replicate_size = self._cfg.get("hsdp_replicate_size") or int( os.environ.get("LOCAL_WORLD_SIZE", "0") ) if replicate_size <= 0: # Fallback: assume one replica per node, equal-sized shard groups. - # Caller should set fsdp_hsdp_replicate_size explicitly when this + # Caller should set hsdp_replicate_size explicitly when this # default is wrong (e.g. uneven node sizes). replicate_size = max(1, world_size // 8) if world_size >= 8 else 1 if world_size % replicate_size != 0: raise ValueError( f"world_size ({world_size}) not divisible by HSDP replicate " - f"size ({replicate_size}); set fsdp_hsdp_replicate_size." + f"size ({replicate_size}); set hsdp_replicate_size." ) shard_size = world_size // replicate_size mesh = init_device_mesh( @@ -156,8 +184,28 @@ def _build_fsdp2_layout(self) -> FSDP2Layout: replicate_pg = mesh.get_group("replicate") shard_pg = mesh.get_group("shard") else: - mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",)) - dp_pg = mesh.get_group("dp") + sp_size = get_sp_world_size() + if sp_size > 1: + if world_size % sp_size != 0: + raise ValueError( + f"world_size ({world_size}) not divisible by " + f"sp_size ({sp_size})." + ) + # Row-major (dp, sp): the "sp" submesh is consecutive ranks and + # the "dp" submesh is strided, matching parallel_state's SP + # (consecutive) and DP (strided) groups. FSDP shards/replicates + # over DP only; the SP dimension is handled by the model's + # activation all-to-all and the trainer's scoped SP grad-sum. + full = init_device_mesh( + "cuda", + (world_size // sp_size, sp_size), + mesh_dim_names=("dp", "sp"), + ) + mesh = full["dp"] + dp_pg = mesh.get_group() + else: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",)) + dp_pg = mesh.get_group("dp") trainer_pg = self._build_trainer_process_group(world_size) @@ -168,7 +216,7 @@ def _build_fsdp2_layout(self) -> FSDP2Layout: offload_policy: Optional[OffloadPolicy] = ( OffloadPolicy(pin_memory=True) - if self._cfg.get("fsdp_cpu_offload", False) + if self._cfg.get("cpu_offload", False) else None ) @@ -177,7 +225,7 @@ def _build_fsdp2_layout(self) -> FSDP2Layout: dtype = torch.bfloat16 if self._cfg.get("bf16", True) else torch.float16 reduce_dtype = ( torch.float32 - if self._cfg.get("fsdp_reduce_dtype", "bf16") == "fp32" + if self._cfg.get("reduce_dtype", "bf16") == "fp32" else dtype ) mp_policy = MixedPrecisionPolicy( @@ -209,6 +257,7 @@ def _wrap_fsdp( model: torch.nn.Module, *, use_dmuon: bool = False, + dmuon_owner_kwargs: Optional[dict] = None, ) -> torch.nn.Module: if not hasattr(model, "convert_to_fsdp"): raise NotImplementedError( @@ -217,18 +266,41 @@ def _wrap_fsdp( self._maybe_enable_grad_ckpt(model) layout = self._build_fsdp2_layout() self._layout = layout - wrapped = model.convert_to_fsdp( + convert_kwargs = dict( mesh=layout.mesh, mp_policy=layout.mp_policy, offload_policy=layout.offload_policy, reshard_after_forward=layout.reshard_after_forward, use_dmuon=use_dmuon, + dmuon_owner_kwargs=dmuon_owner_kwargs, ) + # 若 model 接 ``replicate_non_layer_modules`` kwarg（目前只有 pi0）， + # 透传 cfg 字段；其它 model 路径不动签名 → 跳过传递避免 TypeError。 + convert_sig_params = inspect.signature(model.convert_to_fsdp).parameters + if "replicate_non_layer_modules" in convert_sig_params: + convert_kwargs["replicate_non_layer_modules"] = self._cfg.get( + "replicate_non_layer_modules", False + ) + if "fsdp_forward_prefetch_joint0_with_vision" in convert_sig_params: + convert_kwargs["fsdp_forward_prefetch_joint0_with_vision"] = self._cfg.get( + "fsdp_forward_prefetch_joint0_with_vision", True + ) + wrapped = model.convert_to_fsdp(**convert_kwargs) torch.cuda.empty_cache() return wrapped - def wrap(self, model: torch.nn.Module) -> torch.nn.Module: - return self._wrap_fsdp(model, use_dmuon=False) + def wrap( + self, + model: torch.nn.Module, + *, + use_dedicated_params: bool = False, + dmuon_owner_kwargs: Optional[dict] = None, + ) -> torch.nn.Module: + return self._wrap_fsdp( + model, + use_dmuon=use_dedicated_params, + dmuon_owner_kwargs=dmuon_owner_kwargs, + ) def clip_grad_norm( self, @@ -239,7 +311,7 @@ def clip_grad_norm( ) -> torch.Tensor: del optimizer self.last_grad_clip_stats = None - # Separate DTensor (FSDP2) from regular tensor params - PyTorch's + # Separate DTensor (FSDP2) from regular tensor params — PyTorch's # foreach_mul_ in clip_grad_norm_ can't mix them. try: from torch.distributed.tensor import DTensor @@ -267,6 +339,25 @@ def clip_grad_norm( @contextmanager def no_sync(self, model: torch.nn.Module): + if hasattr(model, "_dedicated_comm_ctx"): + try: + import dmuon + except ImportError: + _logger.warning( + "DMuon model detected but dmuon.no_sync is unavailable; " + "falling back to FSDP2 gradient-sync gating only." + ) + else: + dmuon_no_sync = getattr(dmuon, "no_sync", None) + if dmuon_no_sync is not None: + with dmuon_no_sync(model): + yield + return + _logger.warning( + "DMuon model detected but installed dmuon has no no_sync(); " + "falling back to FSDP2 gradient-sync gating only." + ) + toggled = hasattr(model, "set_requires_gradient_sync") if toggled: model.set_requires_gradient_sync(False) @@ -281,46 +372,9 @@ def params_sharded(self) -> bool: return True -class DMuonFSDPStrategy(FSDPStrategy): - """FSDP2 plus the DMuon dedicated-parameter runtime.""" - - @staticmethod - def _runtime(): - return import_module("dmuon") - - def wrap(self, model: torch.nn.Module) -> torch.nn.Module: - return self._wrap_fsdp(model, use_dmuon=True) - - def clip_grad_norm( - self, - model: torch.nn.Module, - max_norm: float, - *, - optimizer=None, - ) -> torch.Tensor: - total_norm = super().clip_grad_norm(model, max_norm, optimizer=optimizer) - self.last_grad_clip_stats = None - if optimizer is None: - return total_norm - - clip_stats = self._runtime().clip_grad_norm_(optimizer, max_norm) - self.last_grad_clip_stats = clip_stats.as_dict() - device = next(iter(p for p in model.parameters())).device - total_norm_t = ( - total_norm.to(device) - if torch.is_tensor(total_norm) - else torch.tensor(float(total_norm), device=device) - ) - dedicated_norm_t = torch.tensor(float(clip_stats.total_norm), device=device) - return (total_norm_t.pow(2) + dedicated_norm_t.pow(2)).sqrt() - - @contextmanager - def no_sync(self, model: torch.nn.Module): - with self._runtime().no_sync(model): - yield - - class DDPStrategy(DistributionStrategy): + """PyTorch DistributedDataParallel backend for debug/small-model runs.""" + def __init__( self, config: dict, @@ -336,12 +390,17 @@ def __init__( self.bucket_cap_mb = config.get("bucket_cap_mb", 25) self.last_grad_clip_stats = None - def wrap(self, model: torch.nn.Module) -> torch.nn.Module: - # Adapter.convert_to_mix_precision_hint already placed the model on - # self._device; this .to() is a defensive no-op (same as the original - # wrap_ddp_model). - model = model.to(self._device) + def wrap( + self, + model: torch.nn.Module, + *, + use_dedicated_params: bool = False, + dmuon_owner_kwargs: Optional[dict] = None, + ) -> torch.nn.Module: + if use_dedicated_params or dmuon_owner_kwargs: + raise ValueError("DDPStrategy does not support DMuon dedicated params") self._maybe_enable_grad_ckpt(model) + model = model.to(self._device) return DDP( model, device_ids=[self._local_rank], @@ -372,29 +431,441 @@ def params_sharded(self) -> bool: return False +class DMuonStrategy(DistributionStrategy): + """Pure DMuon dedicated-parameter runtime, without FSDP2 wrapping.""" + + def __init__(self, config: dict): + self._cfg = config + self._layout: Optional[DMuonLayout] = None + self.last_grad_clip_stats = None + + @staticmethod + def _runtime(): + return import_module("dmuon") + + @staticmethod + def _debug_sharded_adamw_steps() -> int: + try: + return max( + 0, int(os.environ.get("WALLX_DMUON_SHARDED_ADAMW_DEBUG_STEPS", "0")) + ) + except ValueError: + return 0 + + @staticmethod + def _debug_sharded_adamw_log_limit() -> int: + try: + return max( + 0, + int(os.environ.get("WALLX_DMUON_SHARDED_ADAMW_DEBUG_LOG_LIMIT", "24")), + ) + except ValueError: + return 24 + + @staticmethod + def _tensor_finite_stats(tensor: torch.Tensor) -> dict: + detached = tensor.detach() + finite_mask = torch.isfinite(detached) + all_finite = bool(finite_mask.all().item()) + stats = { + "finite": all_finite, + "numel": int(detached.numel()), + "dtype": str(detached.dtype), + "shape": tuple(int(dim) for dim in detached.shape), + "nan": 0, + "inf": 0, + "amax": float("nan"), + } + if all_finite: + if detached.numel() > 0: + stats["amax"] = float(detached.float().abs().max().cpu().item()) + else: + stats["amax"] = 0.0 + return stats + stats["nan"] = int(torch.isnan(detached).sum().cpu().item()) + stats["inf"] = int(torch.isinf(detached).sum().cpu().item()) + return stats + + @staticmethod + def _format_finite_stats(stats: dict) -> str: + return ( + f"finite={stats['finite']} shape={stats['shape']} " + f"dtype={stats['dtype']} numel={stats['numel']} " + f"nan={stats['nan']} inf={stats['inf']} amax={stats['amax']}" + ) + + @staticmethod + def _maybe_install_sharded_adamw_debug_hooks(model: torch.nn.Module) -> None: + debug_steps = DMuonStrategy._debug_sharded_adamw_steps() + if debug_steps <= 0: + return + comm_ctx = getattr(model, "_dedicated_comm_ctx", None) + if comm_ctx is None: + return + installed = 0 + for state in getattr(comm_ctx, "all_states", []): + group = getattr(state, "group", None) + if group is None: + continue + group_name = str( + getattr(group, "_debug_name", None) or f"group_{id(group):x}" + ) + for dparam in getattr(group, "params", []): + if getattr(dparam, "_dmuon_route", None) != "sharded_adamw": + continue + param = getattr(dparam, "_unsharded_param", None) + if param is None: + continue + if getattr(dparam, "_wallx_sharded_adamw_debug_hook", None) is not None: + continue + name = f"{group_name}.{getattr(dparam, 'param_name', '')}" + dparam._wallx_sharded_adamw_debug_name = name + dparam._wallx_sharded_adamw_raw_calls = 0 + dparam._wallx_sharded_adamw_raw_stats = None + + def _make_hook(dp): + def _hook(grad): + calls = int(getattr(dp, "_wallx_sharded_adamw_raw_calls", 0)) + if calls >= debug_steps: + return grad + stats = DMuonStrategy._tensor_finite_stats(grad) + dp._wallx_sharded_adamw_raw_calls = calls + 1 + dp._wallx_sharded_adamw_raw_stats = stats + if not stats["finite"]: + rank = dist.get_rank() if dist.is_initialized() else 0 + print( + "[WallX DMuon ShardedAdamW Debug] " + f"rank={rank} phase=raw_grad call={calls} " + f"name={getattr(dp, '_wallx_sharded_adamw_debug_name', '')} " + f"{DMuonStrategy._format_finite_stats(stats)}", + flush=True, + ) + return grad + + return _hook + + dparam._wallx_sharded_adamw_debug_hook = param.register_hook( + _make_hook(dparam) + ) + installed += 1 + rank = dist.get_rank() if dist.is_initialized() else 0 + print( + "[WallX DMuon ShardedAdamW Debug] " + f"rank={rank} installed_hooks={installed} debug_steps={debug_steps}", + flush=True, + ) + + @property + def trainer_process_group(self) -> Optional[dist.ProcessGroup]: + return None + + def _build_dmuon_layout(self) -> DMuonLayout: + name = self._cfg.get("sharding_strategy", "shard_grad_op") + world_size = dist.get_world_size() if dist.is_initialized() else 1 + + if name in ("hybrid_shard", "_hybrid_shard_zero2"): + replicate_size = self._cfg.get("hsdp_replicate_size") or int( + os.environ.get("LOCAL_WORLD_SIZE", "0") + ) + if replicate_size <= 0: + replicate_size = max(1, world_size // 8) if world_size >= 8 else 1 + if world_size % replicate_size != 0: + raise ValueError( + f"world_size ({world_size}) not divisible by DMuon HSDP " + f"replicate size ({replicate_size}); set " + "hsdp_replicate_size." + ) + mesh = init_device_mesh( + "cuda", + (replicate_size, world_size // replicate_size), + mesh_dim_names=("replicate", "shard"), + ) + else: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",)) + + reshard_after_forward = name not in ( + "shard_grad_op", + "_hybrid_shard_zero2", + ) + return DMuonLayout( + mesh=mesh, + reshard_after_forward=reshard_after_forward, + ) + + def wrap( + self, + model: torch.nn.Module, + *, + use_dedicated_params: bool = False, + dmuon_owner_kwargs: Optional[dict] = None, + ) -> torch.nn.Module: + del use_dedicated_params + if not hasattr(model, "convert_to_dmuon"): + raise NotImplementedError( + f"Model {model.__class__.__name__} has no convert_to_dmuon method" + ) + self._maybe_enable_grad_ckpt(model) + layout = self._build_dmuon_layout() + self._layout = layout + wrapped = model.convert_to_dmuon( + mesh=layout.mesh, + reshard_after_forward=layout.reshard_after_forward, + dmuon_owner_kwargs=dmuon_owner_kwargs, + ) + self._maybe_install_sharded_adamw_debug_hooks(wrapped) + torch.cuda.empty_cache() + return wrapped + + @staticmethod + def _iter_dedicated_adamw_grads(optimizer): + seen = set() + for dedicated_params in getattr(optimizer, "_adamw_group_dps", {}).values(): + for dparam in dedicated_params: + dparam_id = id(dparam) + if dparam_id in seen: + continue + seen.add(dparam_id) + if getattr(dparam, "_dmuon_route", None) == "sharded_adamw": + grad = getattr(dparam, "_sharded_adamw_grad", None) + else: + grad = getattr(dparam, "_reduced_grad", None) + if grad is not None: + yield dparam, grad + + @staticmethod + def _maybe_log_sharded_adamw_reduce_debug(optimizer, entries) -> None: + debug_steps = DMuonStrategy._debug_sharded_adamw_steps() + if debug_steps <= 0: + return + step = int(getattr(optimizer, "_wallx_sharded_adamw_debug_step", 0)) + if step >= debug_steps: + return + log_limit = DMuonStrategy._debug_sharded_adamw_log_limit() + rank = dist.get_rank() if dist.is_initialized() else 0 + total = 0 + raw_nonfinite = 0 + shard_nonfinite = 0 + logged = 0 + for dparam, grad in entries: + if getattr(dparam, "_dmuon_route", None) != "sharded_adamw": + continue + total += 1 + shard_stats = DMuonStrategy._tensor_finite_stats(grad) + reduce_input = getattr(dparam, "_sharded_adamw_reduce_input", None) + reduce_input_stats = ( + None + if reduce_input is None + else DMuonStrategy._tensor_finite_stats(reduce_input) + ) + raw_stats = getattr(dparam, "_wallx_sharded_adamw_raw_stats", None) + raw_finite = None if raw_stats is None else bool(raw_stats["finite"]) + if raw_stats is not None and not raw_stats["finite"]: + raw_nonfinite += 1 + if not shard_stats["finite"]: + shard_nonfinite += 1 + should_log = ( + not shard_stats["finite"] or raw_finite is False or logged < log_limit + ) + if should_log: + logged += 1 + name = getattr( + dparam, + "_wallx_sharded_adamw_debug_name", + getattr(dparam, "param_name", ""), + ) + raw_msg = ( + "raw_grad=missing" + if raw_stats is None + else "raw_grad=" + DMuonStrategy._format_finite_stats(raw_stats) + ) + reduce_input_msg = ( + "reduce_input=missing" + if reduce_input_stats is None + else "reduce_input=" + + DMuonStrategy._format_finite_stats(reduce_input_stats) + ) + print( + "[WallX DMuon ShardedAdamW Debug] " + f"rank={rank} step={step} phase=post_reduce " + f"name={name} {raw_msg} " + f"{reduce_input_msg} " + f"shard_grad={DMuonStrategy._format_finite_stats(shard_stats)}", + flush=True, + ) + print( + "[WallX DMuon ShardedAdamW Debug] " + f"rank={rank} step={step} phase=post_reduce_summary " + f"total={total} raw_nonfinite={raw_nonfinite} " + f"shard_nonfinite={shard_nonfinite}", + flush=True, + ) + optimizer._wallx_sharded_adamw_debug_step = step + 1 + + @staticmethod + @torch.no_grad() + def _clip_dedicated_adamw_grad_norm( + optimizer, + max_norm: float, + *, + device: torch.device, + ) -> dict: + entries = list(DMuonStrategy._iter_dedicated_adamw_grads(optimizer)) + DMuonStrategy._maybe_log_sharded_adamw_reduce_debug(optimizer, entries) + grads = [grad for _dparam, grad in entries] + + local_sq = torch.zeros((), device=device, dtype=torch.float32) + for grad in grads: + local_sq = local_sq + torch.sum(grad.detach().float().pow(2)).to(device) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(local_sq, op=dist.ReduceOp.SUM) + total_norm = torch.sqrt(local_sq) + clip_coef = float(max_norm) / (total_norm + 1e-6) + clip_coef_clamped = torch.clamp(clip_coef, max=1.0) + if torch.isnan(clip_coef_clamped): + clip_coef_clamped = torch.ones_like(clip_coef_clamped) + for grad in grads: + grad.mul_(clip_coef_clamped.to(device=grad.device)) + + raw_coef = float(clip_coef.detach().cpu().item()) + total_norm_value = float(total_norm.detach().cpu().item()) + finite = torch.isfinite(total_norm).item() + return { + "total_norm": total_norm_value, + "clip_coef": min(raw_coef, 1.0) if finite else raw_coef, + "clipped": bool(finite and raw_coef < 1.0), + "param_count": len(grads), + "found_inf": not bool(finite), + } + + def clip_grad_norm( + self, + model: torch.nn.Module, + max_norm: float, + *, + optimizer=None, + ) -> torch.Tensor: + self.last_grad_clip_stats = None + try: + device = next(iter(p for p in model.parameters())).device + except StopIteration: + device = torch.device("cuda", torch.cuda.current_device()) + + regular_params = [ + p + for p in model.parameters() + if p.grad is not None and not hasattr(p, "_dedicated_owner_rank") + ] + if regular_params: + total_norm = torch.nn.utils.clip_grad_norm_(regular_params, max_norm) + total_norm_t = total_norm.to(device) + else: + total_norm_t = torch.tensor(0.0, device=device) + + if optimizer is None: + return total_norm_t + + clip_stats = self._runtime().clip_grad_norm_(optimizer, max_norm) + adamw_clip_stats = self._clip_dedicated_adamw_grad_norm( + optimizer, + max_norm, + device=device, + ) + self.last_grad_clip_stats = clip_stats.as_dict() + self.last_grad_clip_stats.update( + { + "adamw_total_norm": adamw_clip_stats["total_norm"], + "adamw_clip_coef": adamw_clip_stats["clip_coef"], + "adamw_clipped": adamw_clip_stats["clipped"], + "adamw_param_count": adamw_clip_stats["param_count"], + "adamw_found_inf": adamw_clip_stats["found_inf"], + } + ) + dedicated_norm_t = torch.tensor(float(clip_stats.total_norm), device=device) + dedicated_adamw_norm_t = torch.tensor( + float(adamw_clip_stats["total_norm"]), device=device + ) + return ( + total_norm_t.pow(2) + + dedicated_norm_t.pow(2) + + dedicated_adamw_norm_t.pow(2) + ).sqrt() + + @contextmanager + def no_sync(self, model: torch.nn.Module): + with self._runtime().no_sync(model): + yield + + @property + def params_sharded(self) -> bool: + return True + + def build_strategy( config: dict, *, device: torch.device, local_rank: int, - use_dmuon: bool = False, + use_dmuon: bool | None = None, + optimizer_type: str | None = None, ) -> DistributionStrategy: - """Pick FSDP vs DDP based on ``use_fsdp`` flag. - - ``use_gradient_checkpointing`` is read from the config dict by - each strategy's ``_maybe_enable_grad_ckpt``; not a separate kwarg here. - """ - if config.get("use_fsdp", False): + """Pick the public distributed backend: FSDP2, pure DMuon, or DDP.""" + + backend = str(config.get("backend", "fsdp2")) + if use_dmuon is None: + use_dmuon = optimizer_type == "dmuon" + if optimizer_type == "naive_fsdp2_muon" and backend != "fsdp2": + raise ValueError( + "optimizer_type='naive_fsdp2_muon' requires distributed.backend='fsdp2'." + ) + common = { + "use_gradient_checkpointing": config.get("use_gradient_checkpointing", False), + "use_mixed_precision": config.get("use_mixed_precision", True), + "bf16": config.get("bf16", True), + } + # Forward top-level distributed keys the strategies consume directly. The + # current schema nests sharding/replicate under ``fsdp:``/``dmuon:``, but + # these model-level knobs live at distributed top level. Without forwarding + # them, _maybe_enable_grad_ckpt / convert_to_fsdp silently fall back to + # defaults (checkpoint ALL layers + vision tower, joint0 prefetch on) and + # override the adapter's config-correct setup — a large, uniform MFU + # regression on pi0. ``if k in config`` keeps new-schema configs (keys under + # the active section) untouched, since the section overlay below wins. + for _k in ( + "gradient_checkpoint_num_layers", + "gradient_checkpoint_vision_tower", + "replicate_non_layer_modules", + "fsdp_forward_prefetch_joint0_with_vision", + ): + if _k in config: + common[_k] = config[_k] + if backend == "fsdp2": if use_dmuon: - return DMuonFSDPStrategy(config) - return FSDPStrategy(config) - if use_dmuon: - raise NotImplementedError( - "DMuon dedicated-parameter training requires " - "distributed.use_fsdp=true in this trainer." + raise ValueError( + "FSDP2 + DMuon dedicated params is no longer supported. " + "Use distributed.backend=dmuon for pure DMuon." + ) + return FSDPStrategy({**common, **dict(config.get("fsdp", {}))}) + if backend == "dmuon": + if optimizer_type is not None and optimizer_type != "dmuon": + raise ValueError( + "distributed.backend=dmuon requires optimizer_type='dmuon'." + ) + if not use_dmuon: + raise ValueError( + "distributed.backend=dmuon requires optimizer_type='dmuon'." + ) + return DMuonStrategy({**common, **dict(config.get("dmuon", {}))}) + if backend == "ddp": + if use_dmuon or optimizer_type == "dmuon": + raise ValueError( + "distributed.backend=ddp is not compatible with optimizer_type='dmuon'." + ) + return DDPStrategy( + {**common, **dict(config.get("ddp", {}))}, + device=device, + local_rank=local_rank, ) - return DDPStrategy( - config, - device=device, - local_rank=local_rank, + raise ValueError( + f"Unknown distributed.backend={backend!r}; expected 'fsdp2', " + "'dmuon', or 'ddp'." ) diff --git a/wall_x/trainer/fsdp_trainer/fsdp_trainer.py b/wall_x/trainer/fsdp_trainer/fsdp_trainer.py index 19ffd773..08da7766 100644 --- a/wall_x/trainer/fsdp_trainer/fsdp_trainer.py +++ b/wall_x/trainer/fsdp_trainer/fsdp_trainer.py @@ -1,30 +1,39 @@ """FSDP trainer implementation.""" -import contextlib import dataclasses +import copy import gc -import logging +import json import os -import time -from typing import Any, Optional, Tuple - import torch +import time +import logging +import contextlib from torch.nn.parallel import DistributedDataParallel as DDP +from wall_x.distributed import ( + get_sp_world_size, + get_sp_group, + reduce_sp_partial_grads, +) +from typing import Optional, Any, Tuple +from wall_x.config.loader import dump_config +from wall_x.trainer.optimizer import get_optimizer +from wall_x.utils.profile_zones import zone, Z -from wall_x.model.core.action.normalizer import create_normalizers -from wall_x.trainer.adapters import ADAPTER_REGISTRY, format_adapter_error -from wall_x.trainer.fsdp_trainer import checkpoint_io as _ckpt_io from wall_x.trainer.fsdp_trainer.base_trainer import ( DistributedTrainer, - barrier, is_main_process, + barrier, ) +from wall_x.trainer.adapters import ADAPTER_REGISTRY, format_adapter_error +from wall_x.trainer.fsdp_trainer import checkpoint_io as _ckpt_io from wall_x.trainer.fsdp_trainer.distribution_strategy import build_strategy -from wall_x.trainer.fsdp_trainer.metrics_logger import MetricsLogger -from wall_x.trainer.optimizer import get_optimizer -from wall_x.trainer.scheduler.scheduler import get_scheduler +from wall_x.trainer.metrics import MetricsLogger from wall_x.trainer.utils import move_batch_to_device +from wall_x.model.core.action.normalizer import create_normalizers +from wall_x.trainer.scheduler.scheduler import get_scheduler + class FSDPTrainer(DistributedTrainer): """ @@ -54,6 +63,26 @@ def __init__( self.model_type = train_config.model_type super().__init__(train_config, wandb_run) + # CUDA memory snapshot — start recording BEFORE adapter / strategy / model + # construction so the dumped pickle covers: + # * adapter init (processor, tokenizer) + # * strategy build (FSDP2 mesh / process group) + # * model load + ckpt weight load + # * FSDP/DDP wrap (param sharding + DTensor wiring) + # * optimizer + scheduler creation (AdamW state ≈ 8B × #params, large) + # * dataset / dataloader build + # * the first few training steps + # Rank 0 only — multi-rank snapshots would 8× the .pickle size for no extra + # signal (the per-rank pattern is identical). + if self.cfg.debug.enable_memory_snapshot and self.rank == 0: + torch.cuda.memory._record_memory_history( + max_entries=self.cfg.debug.memory_snapshot_max_entries + ) + self.log( + f"[MemSnap] recording started in __init__ (max_entries=" + f"{self.cfg.debug.memory_snapshot_max_entries})" + ) + if self.model_type not in ADAPTER_REGISTRY: raise ValueError(format_adapter_error(self.model_type)) adapter_cls = ADAPTER_REGISTRY[self.model_type] @@ -67,20 +96,41 @@ def __init__( dataclasses.asdict(self.cfg.distributed), device=self.device, local_rank=self.local_rank, - use_dmuon=self.cfg.hyperparams.optimizer.optimizer_type == "dmuon", + optimizer_type=self.cfg.hyperparams.optimizer.optimizer_type, ) # Model config self.action_dim = self.cfg.action_dim self.use_selective_recompute = self.cfg.distributed.use_selective_recompute self.show_time_details = self.cfg.debug.show_time_details + if self.show_time_details: + self.timers._log_level = max(getattr(self.timers, "_log_level", -1), 0) + self.log( + "[TimeDetailsTimersEnabled] timers log_level set to 0", + level=logging.INFO, + ) # num_training_steps: read from any scheduler that exposes it # (CosineSchedulerConfig currently; future schedulers may add it). - # Used in train_loop to trigger loss_guard_should_stop - works with + # Used in train_loop to trigger loss_guard_should_stop — works with # constant scheduler too once the field is set. sched = self.cfg.hyperparams.scheduler self.num_training_steps = int(getattr(sched, "num_training_steps", 0) or 0) + + # ── loss_guard: optional per-step train_loss dump ───────────── + # Set ``LOSS_GUARD_PATH`` env to opt in; trainer appends + # ``{step, train_loss}`` each optimizer step (main rank only) and + # writes a JSON at end-of-training or when ``num_training_steps`` is + # reached. Used for cross-code-version loss-curve alignment (e.g. + # acc_pi/wall-x baseline ↔ online_rl/wall-x verify). + # Mirror of acc_pi fsdp_trainer.py:130-138, 1094-1106, 1268-1281. + self.loss_guard_path = os.environ.get("LOSS_GUARD_PATH") + self.loss_guard_records: list = [] + self.loss_guard_meta = { + "id": os.environ.get("LOSS_GUARD_ID"), + "name": os.environ.get("LOSS_GUARD_NAME"), + } + self.loss_guard_should_stop = False self.metrics_logger = MetricsLogger( wandb_run=self.wandb_run, log_interval=self.log_interval, @@ -91,38 +141,88 @@ def __init__( self.base_l1_loss = None self.base_l1_loss_detail = {} + startup_t0 = time.time() + stage_t0 = time.time() self.load_normalizer() + self.log("[StartupTiming] normalizer_s={:.3f}".format(time.time() - stage_t0)) + + stage_t0 = time.time() self.load_processor() + self.log("[StartupTiming] processor_s={:.3f}".format(time.time() - stage_t0)) + + stage_t0 = time.time() self.load_model() + self.log("[StartupTiming] model_s={:.3f}".format(time.time() - stage_t0)) # DDP needs to see correct requires_grad at wrap time, so freeze first. + stage_t0 = time.time() self._freeze_params_if_needed(self.model) + self.log( + "[StartupTiming] freeze_params_s={:.3f}".format(time.time() - stage_t0) + ) # Frozen-submodule prefixes for checkpoint filtering. Must be computed # before FSDP wrapping, while child modules still have stable names. + stage_t0 = time.time() self._frozen_prefixes = self._compute_frozen_prefixes() # Original shapes captured before FSDP flattening. self.model._orig_param_shapes = { name: p.shape for name, p in self.model.named_parameters() } + self.log( + "[StartupTiming] param_metadata_s={:.3f}".format(time.time() - stage_t0) + ) if self._resume_from_single_file(): + stage_t0 = time.time() self.load_state_dict( self.model, - {"ckpt": self.cfg.checkpoint.resume_from}, + { + "ckpt": self.cfg.checkpoint.resume_from, + "try_harder": self.cfg.checkpoint.try_harder, + }, + ) + self.log( + "[StartupTiming] single_file_resume_s={:.3f} path={}".format( + time.time() - stage_t0, + self.cfg.checkpoint.resume_from, + ) ) + stage_t0 = time.time() self._wrap_model(self.model) + self.log("[StartupTiming] wrap_model_s={:.3f}".format(time.time() - stage_t0)) + + stage_t0 = time.time() self._create_optimizer() + self.log("[StartupTiming] optimizer_s={:.3f}".format(time.time() - stage_t0)) + + stage_t0 = time.time() self._create_scheduler() + self.log("[StartupTiming] scheduler_s={:.3f}".format(time.time() - stage_t0)) if self._resume_from_training_checkpoint(): + stage_t0 = time.time() self.resume_from_checkpoint() + self.log( + "[StartupTiming] training_checkpoint_resume_s={:.3f}".format( + time.time() - stage_t0 + ) + ) + stage_t0 = time.time() self.load_dataset() + self.log("[StartupTiming] dataset_s={:.3f}".format(time.time() - stage_t0)) + stage_t0 = time.time() self.adapter.init_validation(self.normalizer_action, self.normalizer_propri) + self.log( + "[StartupTiming] validation_init_s={:.3f} total_s={:.3f}".format( + time.time() - stage_t0, + time.time() - startup_t0, + ) + ) def load_normalizer(self): norm_cfg = self.cfg.data.normalizer_config or {} @@ -144,12 +244,49 @@ def load_normalizer(self): ) return + # x2robot_v2 stats are nested (dataset_type/end_effector_type) and are the + # single source of truth: wall-x's flat Normalizer cannot parse them and + # silently collapses to identity. Build the model-side normalizer straight + # from v2's resolver so the saved .pth carries the real statistics and is + # bit-exact with what the v2 dataset pipeline applied during training. + if self.cfg.data.dataset_type == "x2robot_v2" and custom_stats_path: + from wall_x.model.core.action.normalizer import ( + V2_DEFAULT_DELTA_KEY, + V2_DEFAULT_MIN_KEY, + create_normalizers_from_v2_stats, + ) + + bridge_cfg = { + "dof_config": self.cfg.task.dof_config, + "agent_pos_config": self.cfg.task.agent_pos_config, + "stats_file": custom_stats_path, + "min_key": norm_cfg.get("min_key", V2_DEFAULT_MIN_KEY), + "delta_key": norm_cfg.get("delta_key", V2_DEFAULT_DELTA_KEY), + } + for k in ("type", "default_range"): + v = norm_cfg.get(k, None) + if v is not None: + bridge_cfg[k] = v + ( + self.normalizer_action, + self.normalizer_propri, + self._action_statistic_dof, + ) = create_normalizers_from_v2_stats(bridge_cfg) + self.logger.info( + "Built model normalizers from internal_dataset_backend_v2 stats (bridge); " + "dataset_type=x2robot_v2, stats=%s", + custom_stats_path, + ) + return + merged = { "dof_config": self.cfg.task.dof_config, "agent_pos_config": self.cfg.task.agent_pos_config, - "customized_action_statistic_dof": custom_stats_path, - "min_key": norm_cfg.get("min_key", "min"), - "delta_key": norm_cfg.get("delta_key", "delta"), + "normalizer_config": { + "customized_action_statistic_dof": custom_stats_path, + "min_key": norm_cfg.get("min_key", "min"), + "delta_key": norm_cfg.get("delta_key", "delta"), + }, } self.normalizer_action, self.normalizer_propri, self._action_statistic_dof = ( create_normalizers(merged) @@ -179,19 +316,61 @@ def clip_grad_norm(self, max_norm: float) -> torch.Tensor: ) return total_norm + def _reduce_sp_partial_grads(self) -> None: + """Sum the SP-sharded params' partial grads across the SP group. + + Runs once per optimizer step, after FSDP's DP reduce-scatter and before + grad-norm/clip. The mechanism lives in ``wall_x.distributed`` + (``reduce_sp_partial_grads``); *which* params are SP-sharded is declared + by the model adapter (``sp_sharded_param_prefixes``). Here the trainer + only owns the *when* and a one-time sanity guard against a prefix set + that matches nothing (SP enabled on a model that doesn't support it, or + a drifted module FQN layout). + """ + if get_sp_world_size() <= 1: + return + prefixes = self.adapter.sp_sharded_param_prefixes() + matched = reduce_sp_partial_grads( + self.model.named_parameters(), + lambda name: name.startswith(prefixes), + sp_group=get_sp_group(), + ) + if not getattr(self, "_sp_grad_checked", False): + self._sp_grad_checked = True + if matched == 0: + raise RuntimeError( + "sequence parallel is enabled (sp_size>1) but no params " + f"matched the adapter's SP-sharded prefixes {prefixes!r}; " + "either the model does not support SP or its module FQN " + "layout changed." + ) + def load_model(self): """Load and prepare model, optimizer, and scheduler""" + model_t0 = time.time() self.model_config = self.adapter.build_model_config() + self.log("[StartupTiming] model_config_s={:.3f}".format(time.time() - model_t0)) + create_t0 = time.time() self.model = self.adapter.create_model( self.processor, self.tokenizer_mixin, self.model_config ) + self.log( + "[StartupTiming] model_create_s={:.3f}".format(time.time() - create_t0) + ) type(self.adapter).log_attention_implementation(self.logger, self.model) + weights_t0 = time.time() self.adapter.load_weights( self.model, self.normalizer_action, self.normalizer_propri, processor=self.processor, ) + self.log( + "[StartupTiming] adapter_load_weights_s={:.3f} model_total_s={:.3f}".format( + time.time() - weights_t0, + time.time() - model_t0, + ) + ) def load_processor(self): """Load processor and tokenizers""" @@ -262,9 +441,27 @@ def _wrap_model(self, model: torch.nn.Module): model, device=self.device, use_fsdp=self.cfg.distributed.use_fsdp, + backend=self.cfg.distributed.backend, log_fn=self.log, ) - self.model = self.strategy.wrap(model) + from wall_x.trainer.optimizer.dmuon.utils import dmuon_owner_assignment_kwargs + + dmuon_owner_kwargs = dmuon_owner_assignment_kwargs( + self.cfg.hyperparams.optimizer + ) + if dmuon_owner_kwargs: + self.log( + "[DMuon] owner assignment config: " + f"owner_strategy={dmuon_owner_kwargs['owner_strategy']}, " + f"owner_cost_model={dmuon_owner_kwargs['owner_cost_model']}, " + f"owner_pack_small_params={dmuon_owner_kwargs['pack_small_params']}, " + f"hsdp_column_balance={dmuon_owner_kwargs['hsdp_column_balance']}" + ) + self.model = self.strategy.wrap( + model, + use_dedicated_params=self.cfg.hyperparams.optimizer.use_dedicated_params, + dmuon_owner_kwargs=dmuon_owner_kwargs, + ) trainer_pg = getattr(self.strategy, "trainer_process_group", None) if trainer_pg is not None: @@ -290,10 +487,15 @@ def _create_optimizer(self): opt_cfg, self.adapter.default_action_lr_keywords ) if lr_group_configs: - if opt_cfg.optimizer_type not in ("adamw", "muon", "dmuon"): + if opt_cfg.optimizer_type not in ( + "adamw", + "muon", + "dmuon", + "naive_fsdp2_muon", + ): raise ValueError( "optimizer.lr_groups are only supported with adamw, muon, " - "or dmuon" + "dmuon, or naive_fsdp2_muon" ) base_group_name = ( "base_lr_group" if uses_legacy_action_lr_groups(opt_cfg) else "base" @@ -317,7 +519,7 @@ def _create_optimizer(self): "opt_cfg": opt_cfg, "param_groups": param_groups, } - if opt_cfg.optimizer_type == "dmuon": + if opt_cfg.optimizer_type in ("dmuon", "naive_fsdp2_muon"): optimizer_kwargs["log_fn"] = self.log self.optimizer = get_optimizer( @@ -337,8 +539,10 @@ def _create_scheduler(self): lr = self.cfg.hyperparams.optimizer.learning_rate if isinstance(sched, ConstantSchedulerConfig): - self.lr_scheduler = torch.optim.lr_scheduler.ConstantLR( - self.optimizer, factor=1.0, total_iters=0 + self.lr_scheduler = get_scheduler( + optimizer=self.optimizer, + lr_scheduler_type="constant", + num_warmup_steps=sched.num_warmup_steps, ) elif isinstance(sched, CosineSchedulerConfig): if sched.num_training_steps <= 0: @@ -362,40 +566,84 @@ def load_dataset(self): """Load training and validation datasets""" torch.cuda.empty_cache() barrier() + self._resume_checkpoint_path = None + self._resume_loaded_global_step = None + self._resume_next_global_step = self.global_step + self._resume_consumed_samples = None + self._resume_iter_offset = 0 + self._display_train_iter_total = None # Dispatch to adapter with unified signature (supports model-specific kwargs) extra_dataset_kwargs = dict( model_config=self.model_config, tokenizer_mixin=self.tokenizer_mixin, + train_action_tokenizer=self.train_action_tokenizer, + val_action_tokenizer=self.val_action_tokenizer, normalizer_action=self.normalizer_action, normalizer_propri=self.normalizer_propri, ) resume_path = self.cfg.checkpoint.resume_from if resume_path is not None and os.path.isdir(resume_path): + resume_load_start = time.time() + self._resume_checkpoint_path = resume_path # Step and epoch restore are generic. Backend-owned episode # metadata is optional and only loaded when the adapter exposes # the hook. is_incomplete_epoch = False + sampler_state_path = os.path.join( + resume_path, f"sampler_state_rank_{self.rank}.pt" + ) + if os.path.exists(sampler_state_path): + is_incomplete_epoch = True + extra_dataset_kwargs["sampler_state_path"] = sampler_state_path + + dataset_state = _ckpt_io.load_dataset_state( + checkpoint_path=resume_path, + rank=self.rank, + current_world_size=self.world_size, + log_fn=self.log, + ) + if dataset_state is not None: + is_incomplete_epoch = True + extra_dataset_kwargs["dataset_state"] = dataset_state + self._resume_consumed_samples = self._extract_dataset_consumed_samples( + dataset_state + ) + if hasattr(self.adapter, "_load_episode_indices"): resume_indices = self.adapter._load_episode_indices( resume_path, self.rank ) if resume_indices is not None: - is_incomplete_epoch = resume_indices["is_incomplete_epoch"] + is_incomplete_epoch = bool( + is_incomplete_epoch or resume_indices["is_incomplete_epoch"] + ) extra_dataset_kwargs["resume_indices"] = resume_indices train_state = self.adapter.load_step_and_epoch( resume_path, is_incomplete_epoch ) self.global_step = train_state["global_step"] + self._resume_loaded_global_step = self.global_step + self._resume_next_global_step = self.global_step + 1 self.start_epoch = train_state["start_epoch"] self.log( - f"global_step: {self.global_step}, start_epoch: " - f"{self.start_epoch} from {resume_path}" + "[Resume] checkpoint={} loaded_global_step={} " + "next_global_step={} start_epoch={} dataset_state_loaded={} " + "consumed_samples_per_rank={} load_state_s={:.1f}".format( + resume_path, + self._resume_loaded_global_step, + self._resume_next_global_step, + self.start_epoch, + dataset_state is not None, + self._resume_consumed_samples, + time.time() - resume_load_start, + ) ) # global_step records completed steps, so resume from the next step # to avoid triggering _should_save_checkpoint() immediately. self.global_step += 1 + dataset_load_start = time.time() ( self.dataset, self.train_dataloader, @@ -407,6 +655,74 @@ def load_dataset(self): self.world_size, **extra_dataset_kwargs, ) + self._prepare_train_iter_display(len(self.train_dataloader)) + if self._resume_checkpoint_path is not None: + self.log( + "[Resume] dataloader_remaining_iters={} iter_offset={} " + "display_total_iters={} dataloader_batch_size={} " + "dataset_build_s={:.1f}".format( + len(self.train_dataloader), + self._resume_iter_offset, + self._display_train_iter_total, + self._infer_dataloader_batch_size(), + time.time() - dataset_load_start, + ) + ) + + def _maybe_dump_oom_memsnap(self, exc): + """Rescue dump of the CUDA allocator history when OOM hits before the + normal step-3 snapshot has fired. + + Only fires when global_step < memory_snapshot_until_step, because: + - After until_step, the normal dump already wrote a baseline pickle + and ``_record_memory_history(enabled=None)`` stopped recording — + a late dump here would be empty / stale. + - And by then the allocator-history ring buffer is so long the + memory_viz tool can't open the pickle anyway. + + Best-effort: any failure during dump is logged but swallowed, so we + still re-raise the original OOM at the caller. + """ + if not ( + self.cfg.debug.enable_memory_snapshot + and self.rank == 0 + and self.global_step < self.cfg.debug.memory_snapshot_until_step + ): + return + try: + os.makedirs(self.cfg.debug.memory_snapshot_dir, exist_ok=True) + snap_path = os.path.join( + self.cfg.debug.memory_snapshot_dir, + f"memory_snapshot_oom_step{self.global_step}.pickle", + ) + torch.cuda.memory._dump_snapshot(snap_path) + torch.cuda.memory._record_memory_history(enabled=None) + self.log( + f"[MemSnap-OOM] dumped to {snap_path} after " + f"{type(exc).__name__} at step {self.global_step}; " + f"recording stopped" + ) + except Exception as dump_err: + self.log( + f"[MemSnap-OOM] dump failed at step {self.global_step}: " + f"{type(dump_err).__name__}: {dump_err}", + level=logging.WARNING, + ) + + @contextlib.contextmanager + def _oom_dump_guard(self): + """Context manager that rescues a CUDA memsnap on OOM, then re-raises. + + Wrap the GPU-compute-heavy sections (forward / backward / optimizer + step) so that an early-step OOM produces a viz-openable allocator + snapshot before the process dies. See ``_maybe_dump_oom_memsnap`` + for the gate condition. + """ + try: + yield + except torch.OutOfMemoryError as e: + self._maybe_dump_oom_memsnap(e) + raise def train_loop(self, epoch: int, profiler=contextlib.nullcontext()): """Execute training for a single epoch""" @@ -417,6 +733,12 @@ def train_loop(self, epoch: int, profiler=contextlib.nullcontext()): if hasattr(self.dataset, "set_epoch"): self.dataset.set_epoch(epoch) total = len(self.train_dataloader) + if self._resume_checkpoint_path is not None and epoch == self.start_epoch: + display_total = self._display_train_iter_total + else: + self._resume_iter_offset = 0 + self._display_train_iter_total = total + display_total = total # Drop last stop_step = total - total % self.grad_accum_steps @@ -425,12 +747,36 @@ def train_loop(self, epoch: int, profiler=contextlib.nullcontext()): gc.disable() t0 = time.time() + pending_dataset_samples = 0 + first_batch_timing_logged = False with profiler: - self.timers("interval-time", log_level=0).start(barrier=True) - self.timers("data-load", log_level=0).start(barrier=True) + self.timers("interval-time", log_level=0).start( + barrier=self.show_time_details + ) + self.timers("data-load", log_level=0).start(barrier=self.show_time_details) for i, batch in enumerate(self.train_dataloader, self.initial_step): - self.timers("data-load").stop() + # With distributed training, a slow dataloader on one rank makes + # faster ranks wait at the first forward/backward collective. + # When detailed timing is enabled, charge that cross-rank skew + # to data-load instead of compute so rank-local logs are easier + # to interpret. + self.timers("data-load").stop(barrier=self.show_time_details) + if self.show_time_details and not first_batch_timing_logged: + first_batch_timing_logged = True + self.log( + "[DataTiming] first_batch_wait_s={:.3f} " + "local_iter={} display_iter={}/{} global_step={} " + "resume_checkpoint={}".format( + time.time() - t0, + i, + self._resume_iter_offset + i, + display_total, + self.global_step, + getattr(self, "_resume_checkpoint_path", None), + ) + ) + batch_sample_count = self._batch_sample_count(batch) # Save first batch for offline profiling (rank 0 only) debug_batch_path = self.cfg.debug.save_debug_batch_path @@ -440,7 +786,10 @@ def train_loop(self, epoch: int, profiler=contextlib.nullcontext()): self.log(f"Saved debug batch to {debug_batch_path}") # Move batch to device - batch = move_batch_to_device(batch, self.device) + with zone(Z.DATA_TO_DEVICE): + self.timers("batch-to-device", log_level=0).start(barrier=False) + batch = move_batch_to_device(batch, self.device) + self.timers("batch-to-device").stop(sync=True) if self.cfg.debug.enable_mfu_profile and self.global_step == 0: self.log( @@ -449,72 +798,129 @@ def train_loop(self, epoch: int, profiler=contextlib.nullcontext()): ) # Forward pass with autocast - self.timers("forward-compute", log_level=0).start(barrier=False) - with self.autocast_context(): + self.timers("forward-compute", log_level=0).start( + barrier=self.show_time_details + ) + with zone(Z.FORWARD), self._oom_dump_guard(), self.autocast_context(): outputs = self.adapter.forward( self.model, batch, global_step=self.global_step ) - self.timers("forward-compute").stop() + self.timers("forward-compute").stop(sync=True) + if self.show_time_details: + try: + _free_bytes, _total_bytes = torch.cuda.mem_get_info(self.device) + _alloc = torch.cuda.memory_allocated(self.device) / (1024**3) + _reserved = torch.cuda.memory_reserved(self.device) / (1024**3) + _peak_alloc = torch.cuda.max_memory_allocated(self.device) / ( + 1024**3 + ) + _peak_reserved = torch.cuda.max_memory_reserved(self.device) / ( + 1024**3 + ) + self.log( + "[MemoryDetails] " + f"step={self.global_step} rank={self.rank} " + f"after_forward_allocated_gib={_alloc:.3f} " + f"after_forward_reserved_gib={_reserved:.3f} " + f"peak_allocated_gib={_peak_alloc:.3f} " + f"peak_reserved_gib={_peak_reserved:.3f} " + f"free_gib={_free_bytes / (1024 ** 3):.3f} " + f"total_gib={_total_bytes / (1024 ** 3):.3f}", + level=logging.INFO, + ) + except Exception as _mem_exc: + self.log( + f"[MemoryDetails] failed: {type(_mem_exc).__name__}: {_mem_exc}", + level=logging.WARNING, + ) loss = self.adapter.extract_loss(outputs) - # Check for NaN loss -- replace with 0 instead of `continue` - # to avoid skipping collective ops (backward, all_reduce, etc.) - # which would cause NCCL deadlock across ranks. + # Check for NaN loss -- use ``nan_to_num`` instead of + # ``zeros_like`` or ``loss * 0``: + # - ``zeros_like(loss)`` returns a fresh leaf tensor with no + # grad_fn → backward crashes with "element 0 of tensors + # does not require grad and does not have a grad_fn". + # - ``loss * 0`` preserves grad_fn but ``NaN * 0 = NaN`` per + # IEEE 754, so grads still propagate as NaN. + # ``torch.nan_to_num`` is differentiable: its forward replaces + # NaN→0, its backward zeros the upstream grad at NaN positions + # → clean grads, no NCCL deadlock, no NaN-poisoning of params. + # Skipping the step is the caller's responsibility (e.g., via + # GradScaler). nan_loss = torch.isnan(loss) if nan_loss: self.log( f"Warning: nan in loss at epoch: {epoch}, step: {i}", level=logging.WARNING, ) - loss = torch.zeros_like(loss) + loss = torch.nan_to_num(loss, nan=0.0) # Backward pass - self.timers("backward-compute", log_level=0).start(barrier=True) + self.timers("backward-compute", log_level=0).start( + barrier=self.show_time_details + ) context = ( contextlib.nullcontext() if self.sync_gradients() else self.strategy.no_sync(self.model) ) - with context: + with zone(Z.BACKWARD), context, self._oom_dump_guard(): scaled_loss = loss / self.grad_accum_steps self.backward(scaled_loss) self.timers("backward-compute").stop() + pending_dataset_samples += batch_sample_count + + # Capture sync flag before the if/else below mutates micro_step / + # global_step. Used at the end of the loop to gate profiler.step() + # on the gbs boundary so 1 ProfilerStep = grad_accum_steps × (fwd+bwd) + # + reduce-scatter + optim, comparable across grad_accum settings. + did_optim_step = self.sync_gradients() # Gradient sync and optimizer step - if self.sync_gradients(): - self.timers("optimizer", log_level=0).start(barrier=True) + if did_optim_step: + self.timers("optimizer", log_level=0).start( + barrier=self.show_time_details + ) + + # Sequence parallel: sum the SP-sharded DiT block params' + # partial grads across the SP group, before grad-norm/clip. + self._reduce_sp_partial_grads() # Per-component grad norms (before clipping) self.timers("optimizer-grad-norms", log_level=0).start( barrier=False ) self._component_grad_norms = {} - self.adapter.collect_grad_norms( - self.model, - self._component_grad_norms, - device=self.device, - reduce_tensor_fn=self.reduce_tensor, - params_sharded=self.strategy.params_sharded, - ) + with zone(Z.GRAD_NORMS): + self.adapter.collect_grad_norms( + self.model, + self._component_grad_norms, + device=self.device, + reduce_tensor_fn=self.reduce_tensor, + params_sharded=self.strategy.params_sharded, + ) self.timers("optimizer-grad-norms").stop() # Clip gradients self.timers("optimizer-clip", log_level=0).start(barrier=False) - if self.cfg.hyperparams.optimizer.enable_grad_clip: - total_norm = self.clip_grad_norm(self.max_grad_norm) - else: - self._dedicated_param_grad_clip_stats = None - total_norm = 0 + with zone(Z.GRAD_CLIP): + if self.cfg.hyperparams.optimizer.enable_grad_clip: + total_norm = self.clip_grad_norm(self.max_grad_norm) + else: + self._dedicated_param_grad_clip_stats = None + total_norm = 0 self.timers("optimizer-clip").stop() self.timers("optimizer-step", log_level=0).start(barrier=False) - self.optimizer_step() + with zone(Z.OPTIMIZER_STEP), self._oom_dump_guard(): + self.optimizer_step() self.timers("optimizer-step").stop() self.timers("optimizer-zero-grad", log_level=0).start(barrier=False) - self.optimizer_zero_grad() + with zone(Z.ZERO_GRAD): + self.optimizer_zero_grad() self.timers("optimizer-zero-grad").stop() self.timers("optimizer").stop() @@ -522,19 +928,22 @@ def train_loop(self, epoch: int, profiler=contextlib.nullcontext()): self.lr_scheduler_step() # Logging - self.timers("logging", log_level=0).start(barrier=True) - _t_metrics_start = time.time() - self._log_training_metrics( - epoch, i, total, loss, total_norm, outputs + self.timers("logging", log_level=0).start( + barrier=self.show_time_details ) - _t_metrics_ms = (time.time() - _t_metrics_start) * 1000 + with zone(Z.LOGGING): + _t_metrics_start = time.time() + self._log_training_metrics( + epoch, i, total, loss, total_norm, outputs + ) + _t_metrics_ms = (time.time() - _t_metrics_start) * 1000 t1 = time.time() self.training_log( epoch, self.num_epoch, - i, - total, + self._resume_iter_offset + i, + display_total, loss, self.get_lr(), t1 - t0, @@ -554,10 +963,17 @@ def train_loop(self, epoch: int, profiler=contextlib.nullcontext()): level=logging.INFO, ) + # The optimizer/scheduler state now includes all pending + # micro-batches, so advance the dataset cursor before any + # checkpoint for this committed step is written. + self._advance_dataset(pending_dataset_samples) + pending_dataset_samples = 0 + # Checkpoint saving (FSDP full state_dict can take 10-20s every save_interval steps) if self._should_save_checkpoint(): _t_save_start = time.time() - self.save_checkpoint(epoch, self.global_step) + with zone(Z.SAVE_CHECKPOINT): + self.save_checkpoint(epoch, self.global_step) if is_main_process(): self.log( f"[Step time breakdown] save_checkpoint took {(time.time() - _t_save_start):.1f} s at global_step={self.global_step}", @@ -571,6 +987,25 @@ def train_loop(self, epoch: int, profiler=contextlib.nullcontext()): self.global_step += 1 self.micro_step = 0 + # Dump memory snapshot at the configured step (rank 0). + if ( + self.cfg.debug.enable_memory_snapshot + and self.rank == 0 + and self.global_step + == self.cfg.debug.memory_snapshot_until_step + ): + os.makedirs(self.cfg.debug.memory_snapshot_dir, exist_ok=True) + snap_path = os.path.join( + self.cfg.debug.memory_snapshot_dir, + "memory_snapshot.pickle", + ) + torch.cuda.memory._dump_snapshot(snap_path) + torch.cuda.memory._record_memory_history(enabled=None) + self.log( + f"[MemSnap] dumped to {snap_path} at step " + f"{self.global_step}; recording stopped" + ) + # Manual GC outside timing window to avoid random stalls (can add 1-5s every gc_interval_steps) if self.global_step % gc_interval == 0: _t_gc_start = time.time() @@ -590,31 +1025,174 @@ def train_loop(self, epoch: int, profiler=contextlib.nullcontext()): if i == stop_step: break - if not isinstance(profiler, contextlib.nullcontext): + # Fire profiler.step() per optim cycle (gbs boundary). With + # grad_accum_steps > 1 this groups all micro fwd+bwd + the + # cross-node reduce-scatter + optim into one ProfilerStep. + # Without this guard, ProfilerStep tracked micro-batches and + # traces from different grad_accum settings were not comparable. + if did_optim_step and not isinstance(profiler, contextlib.nullcontext): profiler.step() if ( self.num_training_steps > 0 and self.global_step >= self.num_training_steps ): + # Set the flag so base_trainer.train() exits the outer + # epoch loop instead of starting another epoch. + self.loss_guard_should_stop = True break # Setup timers for next iteration if i < total - 1: - self.timers("interval-time", log_level=0).start(barrier=True) - self.timers("data-load", log_level=0).start(barrier=True) + self.timers("interval-time", log_level=0).start( + barrier=self.show_time_details + ) + self.timers("data-load", log_level=0).start( + barrier=self.show_time_details + ) # Re-enable automatic GC after training loop gc.enable() + # loss_guard JSON dump on main rank (only when the LOSS_GUARD_PATH env + # is set). Writes when training is ending — either we hit + # num_training_steps this epoch or this is the last configured epoch. + if ( + is_main_process() + and self.loss_guard_path + and (self.loss_guard_should_stop or epoch == self.num_epoch - 1) + ): + self._write_loss_guard() + # Reset dataloader for next epoch self.train_dataloader = self.dataset.get_train_dataloader() + def _write_loss_guard(self): + """Dump collected per-step train_loss records to LOSS_GUARD_PATH. + + Idempotent: re-callable across epochs without harm. Layout matches + acc_pi/wall-x ``fsdp_trainer._write_loss_guard`` so the two JSONs are + diffable directly. + """ + if not self.loss_guard_path: + return + dirname = os.path.dirname(self.loss_guard_path) + if dirname: + os.makedirs(dirname, exist_ok=True) + payload = { + "id": self.loss_guard_meta.get("id"), + "name": self.loss_guard_meta.get("name"), + "seed": getattr(self, "seed", None), + "num_epochs": self.num_epoch, + "global_step": int(self.global_step), + "losses": [r["train_loss"] for r in self.loss_guard_records], + "records": self.loss_guard_records, + } + with open(self.loss_guard_path, "w", encoding="utf-8") as f: + json.dump(payload, f, ensure_ascii=False, indent=2) + + def _batch_sample_count(self, batch) -> int: + """Best-effort per-rank sample count for dataset cursor advancement.""" + if isinstance(batch, dict): + for value in batch.values(): + if torch.is_tensor(value) and value.ndim > 0: + return int(value.shape[0]) + if isinstance(value, (list, tuple)): + return len(value) + return int(getattr(self.train_dataloader, "batch_size", 1) or 1) + if torch.is_tensor(batch) and batch.ndim > 0: + return int(batch.shape[0]) + if isinstance(batch, (list, tuple)): + return len(batch) + return int(getattr(self.train_dataloader, "batch_size", 1) or 1) + + @staticmethod + def _extract_dataset_consumed_samples(dataset_state) -> Optional[int]: + """Return per-rank consumed samples from backend checkpoint state.""" + if not isinstance(dataset_state, dict): + return None + + candidates = [dataset_state.get("consumed")] + train_sampler_state = dataset_state.get("train_sampler") + if isinstance(train_sampler_state, dict): + candidates.append(train_sampler_state.get("consumed")) + sampler_state = dataset_state.get("sampler") + if isinstance(sampler_state, dict): + candidates.append(sampler_state.get("consumed")) + + for value in candidates: + if value is None: + continue + try: + return int(value) + except (TypeError, ValueError): + continue + return None + + def _infer_dataloader_batch_size(self) -> Optional[int]: + """Best-effort batch-size inference for resume progress display.""" + cfg_data = getattr(self.cfg, "data", None) + for value in ( + getattr(self.train_dataloader, "batch_size", None), + getattr(self.cfg.hyperparams, "batch_size_per_gpu", None), + getattr(cfg_data, "batch_size", None), + ): + if value is None: + continue + try: + batch_size = int(value) + except (TypeError, ValueError): + continue + if batch_size > 0: + return batch_size + return None + + def _prepare_train_iter_display(self, remaining_iters: int) -> None: + """Compute display-only epoch iter offsets after dataset resume.""" + self._resume_iter_offset = 0 + consumed_samples = getattr(self, "_resume_consumed_samples", None) + batch_size = self._infer_dataloader_batch_size() + if consumed_samples is not None and batch_size: + self._resume_iter_offset = consumed_samples // batch_size + remainder = consumed_samples % batch_size + if remainder: + self.log( + "[Resume] consumed_samples_per_rank={} is not divisible by " + "dataloader_batch_size={}; display iter offset is rounded down".format( + consumed_samples, + batch_size, + ), + level=logging.WARNING, + ) + self._display_train_iter_total = self._resume_iter_offset + int(remaining_iters) + + def _advance_dataset(self, n: int) -> None: + """Advance backend-owned dataset cursors after a committed train step.""" + if n <= 0: + return + dataset = getattr(self, "dataset", None) + if dataset is not None and hasattr(dataset, "advance"): + dataset.advance(n) + elif dataset is not None and hasattr(dataset, "advance_train_sampler"): + dataset.advance_train_sampler(n) + def _log_training_metrics(self, epoch, step, total, loss, total_norm, outputs): """Collect per-step stats, delegate buffering + wandb emission.""" lr = self.get_lr() train_loss = self.reduce_tensor(loss.detach()).item() + # ── loss_guard: append per-step train_loss on main rank ── + # Stops appending once we've recorded num_training_steps points so the + # JSON payload is bounded even if extra epochs slip through. + if is_main_process() and self.loss_guard_path: + if ( + self.num_training_steps <= 0 + or len(self.loss_guard_records) < self.num_training_steps + ): + self.loss_guard_records.append( + {"step": int(self.global_step), "train_loss": float(train_loss)} + ) + step_stats = { "lr": lr, "train_loss": train_loss, @@ -648,8 +1226,28 @@ def _log_training_metrics(self, epoch, step, total, loss, total_norm, outputs): "muon_grad_clipped": float(dedicated_clip_stats["clipped"]), } ) + if "adamw_total_norm" in dedicated_clip_stats: + step_stats.update( + { + "dmuon_adamw_grad_norm": dedicated_clip_stats[ + "adamw_total_norm" + ], + "dmuon_adamw_grad_clip_coef": dedicated_clip_stats[ + "adamw_clip_coef" + ], + "dmuon_adamw_grad_clipped": float( + dedicated_clip_stats["adamw_clipped"] + ), + "dmuon_adamw_grad_found_inf": float( + dedicated_clip_stats["adamw_found_inf"] + ), + } + ) - if self.global_step % 100 == 0: + param_norm_interval = int( + getattr(self.cfg.logging, "param_norm_interval", 100) or 0 + ) + if param_norm_interval > 0 and self.global_step % param_norm_interval == 0: self.adapter.collect_param_norms( self.model, step_stats, @@ -674,27 +1272,37 @@ def training_log( time_per_step, show_time_details=False, ): - # timers.log() contains all_gather - must run on ALL ranks before the + timer_names = [ + "interval-time", + "data-load", + "batch-to-device", + "forward-compute", + "backward-compute", + "optimizer", + "optimizer-grad-norms", + "optimizer-clip", + "optimizer-step", + "optimizer-zero-grad", + "logging", + ] + timer_detail_stats = {} + # timers.log() contains all_gather — must run on ALL ranks before the # is_main_process() guard to avoid NCCL deadlock. if show_time_details: - self.timers.log( - [ - "interval-time", - "data-load", - "forward-compute", - "backward-compute", - "optimizer", - "optimizer-grad-norms", - "optimizer-clip", - "optimizer-step", - "optimizer-zero-grad", - "logging", - ], - normalizer=1, + timer_detail_stats = self._collect_timer_detail_stats( + timer_names, time_per_step ) + self.timers.log(timer_names, normalizer=1) main = is_main_process() stats = getattr(self, "_current_step_stats", None) or {} + time_stats = {"time_per_iter": float(time_per_step)} + if timer_detail_stats: + time_stats.update(timer_detail_stats) + stats.update(time_stats) + raw_stats = getattr(self, "_current_step_raw_stats", None) + if raw_stats is not None: + raw_stats.update(time_stats) # MFU is computed after step time is known, then merged into both the # console stats and the MetricsLogger buffer so wandb records it. @@ -704,7 +1312,13 @@ def training_log( self.model.module if hasattr(self.model, "module") else self.model ) mfu_info = self.adapter.compute_mfu(unwrapped, time_per_step) - if mfu_info is not None: + # Per-rank [MFU-rank] line is verbose (N_rank × N_step entries). + # Gate it on log_interval so file logs match wandb emission cadence. + log_mfu_line = ( + mfu_info is not None + and self.global_step % self.cfg.logging.log_interval == 0 + ) + if log_mfu_line: self.log( "[MFU-rank] rank={} mfu={:.2f}% step={:.3f}s " "flops_step={:.3f}T flops_fwd={:.3f}T " @@ -724,8 +1338,10 @@ def training_log( mfu_stats = { "mfu": mfu_info["mfu"], "flops_per_step_T": mfu_info["flops_per_step_T"], - "flops_forward_T": mfu_info.get("flops_forward_T"), } + flops_forward = mfu_info.get("flops_forward_T") + if flops_forward is not None: + mfu_stats["flops_forward_T"] = flops_forward profile_fwd_flops = mfu_info.get("profile_fwd_flops_T") if profile_fwd_flops is not None: mfu_stats["profile_fwd_flops_T"] = profile_fwd_flops @@ -735,8 +1351,28 @@ def training_log( raw_stats.update(mfu_stats) if main: + # Expose per-step wall time on the same wandb track as loss/lr/mfu. + # samples_per_sec 是全局吞吐（所有 rank 求和），跟 micro/accum 切分无关， + # 只看 global batch / step time。 + inv_time = 1.0 / max(float(time_per_step), 1e-9) + samples_per_step = ( + self.cfg.hyperparams.batch_size_per_gpu + * self.world_size + * self.grad_accum_steps + ) + step_time_stats = { + "step_time": float(time_per_step), + "samples_per_sec": samples_per_step * inv_time, + } + if stats is not None: + stats.update(step_time_stats) raw_stats = getattr(self, "_current_step_raw_stats", None) if raw_stats is not None: + # 跳过 warm-up：前 N 步 step_time / samples_per_sec 不进 wandb， + # 否则 step 0 ~85s 把 y 轴拉爆后续看不清。控制台 line（stats） + # 仍然显示原值。其他 key（loss / lr / mfu / grad_norm）照常写。 + if self.global_step >= self.cfg.logging.perf_log_skip_steps: + raw_stats.update(step_time_stats) self.metrics_logger.record_step(raw_stats, is_main=True) avg_stats = self.metrics_logger.flush_if_due(self.global_step) if avg_stats is not None: @@ -766,6 +1402,31 @@ def training_log( ) self.log(line) + def _collect_timer_detail_stats(self, timer_names, time_per_step): + """Collect per-step timer details before ``timers.log`` resets them.""" + try: + elapsed = self.timers._get_global_min_max_time( + timer_names, + reset=False, + barrier=False, + normalizer=1.0, + ) + except Exception as exc: + self.log( + f"[TimeDetails] failed to collect timer stats: {type(exc).__name__}: {exc}", + level=logging.WARNING, + ) + return {} + + step_time = max(float(time_per_step), 1e-9) + stats = {} + for name, (min_time, max_time) in elapsed.items(): + key = name.replace("-", "_") + stats[f"time/{key}_min_s"] = float(min_time) + stats[f"time/{key}_max_s"] = float(max_time) + stats[f"time/{key}_max_pct"] = float(max_time) / step_time + return stats + def _should_save_checkpoint(self) -> bool: """Check if checkpoint should be saved""" log = self.cfg.logging @@ -779,16 +1440,61 @@ def _should_validate(self) -> bool: """Check if validation should be performed""" log = self.cfg.logging return ( - self.global_step >= log.ignore_until_interval + log.enable_validation + and self.global_step >= log.ignore_until_interval and self.global_step % log.val_interval == 0 and self.global_step > 0 ) + def _checkpoint_config_dict(self) -> dict: + # dump_config (not asdict): persist the effective config in the same + # typed yaml schema as hand-written configs. + config = dump_config(self.cfg) + effective_data_config = getattr( + getattr(self, "dataset", None), "effective_data_config", None + ) + if effective_data_config is None: + return config + + config.setdefault("data", {})["dataset_config_effective"] = ( + self._dataset_config_params_only(effective_data_config) + ) + return config + + @staticmethod + def _dataset_config_params_only(effective_data_config: dict) -> dict: + data_config = copy.deepcopy(effective_data_config) + data_config.pop("tasks", None) + + dataset_config = data_config.get("dataset") + if not isinstance(dataset_config, dict): + return data_config + + dataset_config.pop("sources", None) + dataset_config.pop("sources_include", None) + + sampler = dataset_config.get("sampler") + if not isinstance(sampler, dict): + return data_config + + task_balance = sampler.get("task_balance") + if not isinstance(task_balance, dict): + return data_config + + task_balance.pop("counts_per_task", None) + task_balance.pop("report_path", None) + params = task_balance.get("params") + if isinstance(params, dict): + params.pop("counts_per_task", None) + if not params: + task_balance.pop("params", None) + return data_config + @torch.no_grad() def val_loop(self): """Delegate the full validation run to the adapter. - Skips silently when the dataset has no val split - adapter's + Skips silently when the dataset has no val split — adapter's run_validation iterates ``val_dataloader`` with ``tqdm(..., total=len(...))`` and cannot accept None. v2 returns None here when the YAML only declares a train split. @@ -863,7 +1569,7 @@ def save_checkpoint(self, epoch: int, step: int = 0): model=self.model, optimizer=self.optimizer, lr_scheduler=self.lr_scheduler, - config=dataclasses.asdict(self.cfg), + config=self._checkpoint_config_dict(), rank=self.rank, is_main=is_main_process(), epoch=epoch, @@ -908,7 +1614,10 @@ def resume_from_checkpoint(self): model=self.model, optimizer=self.optimizer, lr_scheduler=self.lr_scheduler, - resume_config={"ckpt": self.cfg.checkpoint.resume_from}, + resume_config={ + "ckpt": self.cfg.checkpoint.resume_from, + "try_harder": self.cfg.checkpoint.try_harder, + }, rank=self.rank, grad_scaler=self.grad_scaler, model_class=self._checkpoint_model_class(), diff --git a/wall_x/trainer/fsdp_trainer/metrics/__init__.py b/wall_x/trainer/fsdp_trainer/metrics/__init__.py deleted file mode 100644 index 31235735..00000000 --- a/wall_x/trainer/fsdp_trainer/metrics/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from wall_x.trainer.fsdp_trainer.metrics.norms import DistributedNormAccumulator - -__all__ = ["DistributedNormAccumulator"] diff --git a/wall_x/trainer/fsdp_trainer/train_fsdp.py b/wall_x/trainer/fsdp_trainer/train_fsdp.py index e30fdfb6..21496922 100644 --- a/wall_x/trainer/fsdp_trainer/train_fsdp.py +++ b/wall_x/trainer/fsdp_trainer/train_fsdp.py @@ -9,17 +9,16 @@ import logging import os import sys -from datetime import datetime - import torch -import wandb - +from datetime import datetime from wall_x.config.loader import load_config as load_typed_config from wall_x.trainer.fsdp_trainer import ( FSDPTrainer, cleanup_distributed, is_main_process, ) +from wall_x.distributed import init_sequence_parallel, destroy_parallel_groups +import wandb logger = logging.getLogger(__name__) @@ -54,8 +53,13 @@ def parse_args(): "--fsdp_sharding_strategy", type=str, default=None, - choices=["full_shard", "shard_grad_op", "no_shard", "hybrid_shard"], - help="FSDP sharding strategy (overrides config)", + choices=[ + "full_shard", + "shard_grad_op", + "hybrid_shard", + "_hybrid_shard_zero2", + ], + help="Legacy sharding strategy override for the active distributed backend", ) parser.add_argument( "--debug", @@ -117,8 +121,11 @@ def print_fsdp_config(cfg): if is_main_process(): dist = cfg.distributed logger.info("%s", "=" * 60) - logger.info("FSDP Configuration:") - logger.info(" use_fsdp: %s", dist.use_fsdp) + logger.info("Distributed Configuration:") + logger.info(" backend: %s", dist.backend) + active = dist.dmuon if dist.backend == "dmuon" else dist.fsdp + logger.info(" sharding_strategy: %s", active.sharding_strategy) + logger.info(" hsdp_replicate_size: %s", active.hsdp_replicate_size) logger.info("%s", "=" * 60) @@ -151,12 +158,18 @@ def main(): else: logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) + # Determinism (incl. the cublas workspace pin) is config-driven and applied + # in seed_all() from hyperparams.deterministic_algorithms. torch.cuda.init() local_rank = int(os.environ.get("LOCAL_RANK", 0)) device = torch.device(f"cuda:{local_rank}") torch.cuda.set_device(device) torch.distributed.init_process_group("nccl", device_id=device) cfg = load_config(args.config, cli_args=args) + # Build SP + DP process groups before the FSDP2 device mesh (constructed in + # FSDPTrainer). Called unconditionally so get_dp_rank()/get_dp_world_size() + # are valid even when sp_size==1 (the data sampler reads them). + init_sequence_parallel(cfg.distributed.sp_size) wandb_run = setup_logger(cfg) print_fsdp_config(cfg) @@ -172,6 +185,7 @@ def main(): raise finally: + destroy_parallel_groups() cleanup_distributed() if wandb_run is not None: wandb_run.finish() diff --git a/wall_x/trainer/metrics/__init__.py b/wall_x/trainer/metrics/__init__.py new file mode 100644 index 00000000..5426f27a --- /dev/null +++ b/wall_x/trainer/metrics/__init__.py @@ -0,0 +1,13 @@ +"""Training metrics: computation (``mfu/``, ``norms``) and transport (``metrics_logger``). + +Dependency rule: this package is a leaf. It must not import from +``wall_x.trainer.adapters`` or ``wall_x.trainer.fsdp_trainer`` -- both of +those import it, and a back-edge would be a circular import (the old +``fsdp_trainer/metrics`` location had exactly that problem: adapters +imported it through the eagerly-importing ``fsdp_trainer`` package). +""" + +from wall_x.trainer.metrics.metrics_logger import MetricsLogger +from wall_x.trainer.metrics.norms import DistributedNormAccumulator + +__all__ = ["MetricsLogger", "DistributedNormAccumulator"] diff --git a/wall_x/trainer/fsdp_trainer/metrics_logger.py b/wall_x/trainer/metrics/metrics_logger.py similarity index 93% rename from wall_x/trainer/fsdp_trainer/metrics_logger.py rename to wall_x/trainer/metrics/metrics_logger.py index ee00efe4..0d415c73 100644 --- a/wall_x/trainer/fsdp_trainer/metrics_logger.py +++ b/wall_x/trainer/metrics/metrics_logger.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections import deque +from numbers import Real from typing import Any, Deque, Dict, List, Optional, Tuple @@ -71,7 +72,11 @@ def _average(buf: List[Dict[str, Any]]) -> Dict[str, Any]: all_keys.update(stats.keys()) avg: Dict[str, Any] = {} for key in all_keys: - values = [s[key] for s in buf if key in s] + values = [ + s[key] + for s in buf + if key in s and s[key] is not None and isinstance(s[key], Real) + ] if values: avg[key] = sum(values) / len(values) return avg @@ -93,7 +98,7 @@ def format_training_line( """Render a line matching pre-refactor training_log output exactly. Layout: " epoch E/T | iter I/T | loss L | [adapter fields] | lr L | - time_current_backward_step Ts | [MFU P% |]" + time_per_iter Ts | [MFU P% |]" """ parts = [ " epoch {:3d}/{:3d} |".format(epoch, total_epoch), @@ -105,7 +110,7 @@ def format_training_line( if value is not None: parts.append(" {} {:{fmt}} |".format(label, value, fmt=fmt)) parts.append(" lr {:.6f} |".format(lr)) - parts.append(" time_current_backward_step {:.6f}s |".format(time_per_step)) + parts.append(" time_per_iter {:.6f}s |".format(time_per_step)) if mfu_info is not None: parts.append(" MFU {:.2f}% |".format(mfu_info["mfu"] * 100)) return "".join(parts) diff --git a/wall_x/trainer/metrics/mfu/__init__.py b/wall_x/trainer/metrics/mfu/__init__.py new file mode 100644 index 00000000..3a2bf0e2 --- /dev/null +++ b/wall_x/trainer/metrics/mfu/__init__.py @@ -0,0 +1,9 @@ +"""Per-model MFU calculators + shared GPU hardware specs. + +One module per model family (``vga_mfu``; pi0/qwen still live in their +adapters). ``hardware`` holds the shared peak-FLOPS / bandwidth tables. +""" + +from wall_x.trainer.metrics.mfu.hardware import GpuSpecs, resolve_gpu_specs + +__all__ = ["GpuSpecs", "resolve_gpu_specs"] diff --git a/wall_x/trainer/metrics/mfu/hardware.py b/wall_x/trainer/metrics/mfu/hardware.py new file mode 100644 index 00000000..4768f804 --- /dev/null +++ b/wall_x/trainer/metrics/mfu/hardware.py @@ -0,0 +1,76 @@ +"""GPU hardware specs shared by all per-model MFU calculators. + +Peak dense BF16 TFLOPS (no sparsity) and memory bandwidth per device +family. Matched by substring against ``torch.cuda.get_device_name()``. +""" + +from __future__ import annotations + +from typing import NamedTuple, Optional + +import torch + +PEAK_TFLOPS = { + "A100": 312, + "A800": 312, + "H100": 989, + "H800": 989, + "H20": 148, + "A30": 165, + "A40": 150, + "L40S": 362, + "L40": 362, + "4090": 330, + "3090": 142, + "B200": 2250, + "GB200": 2250, +} + +MEM_BW_GBS = { + "A100": 2039, + "A800": 2039, + "H100": 3350, + "H800": 3350, + "H20": 4000, + "A30": 933, + "A40": 696, + "L40S": 864, + "L40": 864, + "4090": 1008, + "3090": 936, + "B200": 8000, + "GB200": 8000, +} + +# Fallback when the device name matches nothing: A800 numbers (the fleet +# default), same behavior as the previous in-adapter table. +_DEFAULT_PEAK_TFLOPS = 312 +_DEFAULT_MEM_BW_GBS = 2039 + + +class GpuSpecs(NamedTuple): + gpu_name: str + peak_flops: float # FLOPs/s + mem_bw_gbs: float + + @property + def ops_byte_ratio(self) -> float: + return self.peak_flops / (self.mem_bw_gbs * 1e9) + + +def resolve_gpu_specs(gpu_name: Optional[str] = None) -> GpuSpecs: + """Look up peak FLOPS / memory bandwidth for a GPU. + + ``gpu_name`` defaults to the current CUDA device's name (empty string + when CUDA is unavailable, which resolves to the A800 fallback). + """ + if gpu_name is None: + gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "" + peak_tflops = _DEFAULT_PEAK_TFLOPS + mem_bw = _DEFAULT_MEM_BW_GBS + for key, tflops in PEAK_TFLOPS.items(): + if key in gpu_name: + peak_tflops = tflops + mem_bw = MEM_BW_GBS.get(key, _DEFAULT_MEM_BW_GBS) + break + return GpuSpecs(gpu_name, peak_tflops * 1e12, mem_bw) diff --git a/wall_x/trainer/fsdp_trainer/metrics/norms.py b/wall_x/trainer/metrics/norms.py similarity index 98% rename from wall_x/trainer/fsdp_trainer/metrics/norms.py rename to wall_x/trainer/metrics/norms.py index f5185341..4f1311a1 100644 --- a/wall_x/trainer/fsdp_trainer/metrics/norms.py +++ b/wall_x/trainer/metrics/norms.py @@ -8,7 +8,7 @@ class DistributedNormAccumulator: - """Accumulate L2 norms without leaking parallel-layout branches to adapters.""" + """Accumulate L2 norms.""" def __init__( self, diff --git a/wall_x/trainer/optimizer/__init__.py b/wall_x/trainer/optimizer/__init__.py index 596a6aa9..52714a9c 100644 --- a/wall_x/trainer/optimizer/__init__.py +++ b/wall_x/trainer/optimizer/__init__.py @@ -4,6 +4,8 @@ # installed so wall_x still imports; the registry will simply not have # "dmuon" and get_optimizer("dmuon", ...) will raise a clear KeyError. try: - from . import dmuon # noqa: F401 - side-effect registers "dmuon" + from . import dmuon # noqa: F401 — side-effect registers "dmuon" except ImportError: pass + +from . import naive_fsdp2_muon # noqa: F401 - side-effect registers temporary baseline diff --git a/wall_x/trainer/optimizer/dmuon/policy.py b/wall_x/trainer/optimizer/dmuon/policy.py new file mode 100644 index 00000000..986f02eb --- /dev/null +++ b/wall_x/trainer/optimizer/dmuon/policy.py @@ -0,0 +1,110 @@ +"""Wall-X DMuon parameter policy helpers.""" + +from __future__ import annotations + +from typing import Callable, Iterable, Optional + +import torch +from torch import nn + + +RouteHintFn = Callable[[str, nn.Parameter], Optional[str]] + +WALLX_ACTION_FP32_PREFIXES = ( + "action_processor.", + "action_preprocessor.", + "model.action_processor.", + "model.action_preprocessor.", +) +WALLX_ACTION_FP32_NAME_PARTS = WALLX_ACTION_FP32_PREFIXES + + +def make_dmuon_param_policy_fn( + route_hint_fn: RouteHintFn, + *, + fp32_name_prefixes: tuple[str, ...] = (), + fp32_name_parts: tuple[str, ...] = (), + default_param_dtype: torch.dtype = torch.bfloat16, +) -> Callable[[str, nn.Parameter], Optional[dict[str, object]]]: + """Convert Wall-X route selection into DMuon's structured param policy.""" + + fp32_name_prefixes = tuple(prefix for prefix in fp32_name_prefixes if prefix) + # Backward-compatible alias for older call sites. The new path uses explicit + # prefixes so qwen action-expert matrices stay on bf16 Tensor Core kernels. + if not fp32_name_prefixes and fp32_name_parts: + fp32_name_prefixes = tuple(part for part in fp32_name_parts if part) + + def param_policy_fn(name: str, param: nn.Parameter) -> Optional[dict[str, object]]: + route = route_hint_fn(name, param) + if route is None: + return None + + policy: dict[str, object] = { + "route": route, + "param_dtype": default_param_dtype, + } + if any(name.startswith(prefix) for prefix in fp32_name_prefixes): + policy.update( + { + "param_dtype": torch.float32, + "grad_dtype": torch.float32, + "cast_forward_inputs": True, + } + ) + return policy + + return param_policy_fn + + +def make_wallx_action_fp32_dmuon_param_policy_fn( + route_hint_fn: RouteHintFn, +) -> Callable[[str, nn.Parameter], Optional[dict[str, object]]]: + """Keep WallX action processor/head compute in fp32 under pure DMuon. + + The route still comes from the model-specific route selector. This helper + only overlays dtype policy: ordinary WallX parameters materialize in bf16, + while action processor/head parameters materialize gradients and forward + weights in fp32. + """ + + return make_dmuon_param_policy_fn( + route_hint_fn, + fp32_name_prefixes=WALLX_ACTION_FP32_PREFIXES, + ) + + +def summarize_dmuon_param_policy( + named_parameters: Iterable[tuple[str, nn.Parameter]], + param_policy_fn: Callable[[str, nn.Parameter], Optional[dict[str, object]]], + *, + max_examples_per_key: int = 6, +) -> tuple[dict[str, int], dict[str, list[str]]]: + """Return compact policy counts/examples for training logs.""" + + counts: dict[str, int] = {} + examples: dict[str, list[str]] = {} + for name, param in named_parameters: + if not param.requires_grad: + continue + policy = param_policy_fn(name, param) + if policy is None: + key = "unmanaged" + else: + route = policy.get("route") + param_dtype = policy.get("param_dtype") + grad_dtype = policy.get("grad_dtype") + key = f"route={route},param_dtype={param_dtype}," f"grad_dtype={grad_dtype}" + counts[key] = counts.get(key, 0) + 1 + bucket = examples.setdefault(key, []) + if len(bucket) < max_examples_per_key: + bucket.append(name) + return counts, examples + + +__all__ = [ + "WALLX_ACTION_FP32_NAME_PARTS", + "WALLX_ACTION_FP32_PREFIXES", + "make_dmuon_param_policy_fn", + "make_wallx_action_fp32_dmuon_param_policy_fn", + "summarize_dmuon_param_policy", +] diff --git a/wall_x/trainer/optimizer/dmuon/utils.py b/wall_x/trainer/optimizer/dmuon/utils.py index 3d03b826..44ef160e 100644 --- a/wall_x/trainer/optimizer/dmuon/utils.py +++ b/wall_x/trainer/optimizer/dmuon/utils.py @@ -11,10 +11,15 @@ def _emit(log_fn, message, *args, level=logging.INFO): if args: message = message % args - if log_fn is not None: - log_fn(message, level=level) - else: + if log_fn is None: _logger.log(level, message) + return + if level >= logging.WARNING: + # log_fn flattens record levels; keep a real WARNING record for + # level-based filters/alerting alongside the display path. + _logger.log(level, message) + prefix = "WARNING: " if level >= logging.WARNING else "" + log_fn(f"{prefix}{message}") def _is_rank0(): @@ -51,6 +56,40 @@ def _build_ns_backend(dmuon, opt_cfg): ) +def _maybe_add_muon_kwarg(kwargs, signature, name, value): + if name in signature.parameters: + kwargs[name] = value + return True + return False + + +def dmuon_dedicate_mesh_args(mesh): + """Return ``mesh`` and kwargs in the shape expected by DMuon.""" + + mesh_dim_names = tuple(getattr(mesh, "mesh_dim_names", ()) or ()) + if "replicate" in mesh_dim_names and "shard" in mesh_dim_names: + return mesh["shard"], {"replicate_mesh": mesh["replicate"]} + if getattr(mesh, "ndim", 1) == 1: + return mesh, {} + raise ValueError( + "Unsupported FSDP2 mesh layout for DMuon: " + f"mesh_dim_names={mesh_dim_names!r}, ndim={getattr(mesh, 'ndim', None)!r}" + ) + + +def dmuon_owner_assignment_kwargs(opt_cfg): + """Return explicit owner-assignment kwargs for ``dmuon.dedicate_params``.""" + + if getattr(opt_cfg, "optimizer_type", None) != "dmuon": + return {} + return { + "owner_strategy": getattr(opt_cfg, "owner_strategy", "lpt"), + "owner_cost_model": getattr(opt_cfg, "owner_cost_model", "optimizer"), + "pack_small_params": bool(getattr(opt_cfg, "owner_pack_small_params", True)), + "hsdp_column_balance": bool(getattr(opt_cfg, "hsdp_column_balance", True)), + } + + def get_dmuon_optimizer(model, *, opt_cfg, param_groups=None, log_fn=None): """Build dmuon.Muon from a DMuonConfig. @@ -77,12 +116,18 @@ def get_dmuon_optimizer(model, *, opt_cfg, param_groups=None, log_fn=None): ) ns_backend = _build_ns_backend(dmuon, opt_cfg) + record_step_profile = bool(getattr(opt_cfg, "record_step_profile", False)) + first_step_progress_log_limit = int( + getattr(opt_cfg, "first_step_progress_log_limit", 8) + ) _emit( log_fn, "DMuon: Muon lr=%s momentum=%s ns_steps=%s; " "AdamW lr=%s betas=%s wd=%s; " - "ns_backend=%s ns_coefficients=%s nesterov=%s", + "ns_backend=%s ns_coefficients=%s nesterov=%s " + "record_step_profile=%s first_step_progress_log_limit=%s " + "replicate_async=%s group_prepare_ahead=%s forward_prefetch_depth=%s", opt_cfg.muon_lr, opt_cfg.momentum, opt_cfg.ns_steps, @@ -92,6 +137,11 @@ def get_dmuon_optimizer(model, *, opt_cfg, param_groups=None, log_fn=None): opt_cfg.ns_backend, opt_cfg.ns_coefficients, opt_cfg.nesterov, + record_step_profile, + first_step_progress_log_limit, + opt_cfg.replicate_async, + opt_cfg.group_prepare_ahead, + opt_cfg.forward_prefetch_depth, ) if param_groups is not None: _emit( @@ -110,6 +160,25 @@ def get_dmuon_optimizer(model, *, opt_cfg, param_groups=None, log_fn=None): kwargs = {} if param_groups is not None: kwargs["param_groups"] = param_groups + optional_kwargs = { + "record_step_profile": record_step_profile, + "first_step_progress_log_limit": first_step_progress_log_limit, + "replicate_async": opt_cfg.replicate_async, + "group_prepare_ahead": opt_cfg.group_prepare_ahead, + "forward_prefetch_depth": opt_cfg.forward_prefetch_depth, + } + unsupported_kwargs = [] + for name, value in optional_kwargs.items(): + if not _maybe_add_muon_kwarg(kwargs, muon_signature, name, value): + unsupported_kwargs.append(name) + if unsupported_kwargs and _is_rank0(): + _emit( + log_fn, + "Installed DMuon does not support optional config %s; " + "continuing without them.", + unsupported_kwargs, + level=logging.WARNING, + ) optimizer = dmuon.Muon( model, diff --git a/wall_x/trainer/optimizer/naive_fsdp2_muon.py b/wall_x/trainer/optimizer/naive_fsdp2_muon.py new file mode 100644 index 00000000..f74ebb3d --- /dev/null +++ b/wall_x/trainer/optimizer/naive_fsdp2_muon.py @@ -0,0 +1,527 @@ +"""Naive FSDP2 Muon baseline. + +This is intentionally inefficient: FSDP2 keeps parameters sharded, but every +rank all-gathers each Muon gradient to a full matrix, runs Newton-Schulz +locally, then drops the full update after redistributing it back to the +parameter's DTensor layout. +""" + +from __future__ import annotations + +import logging +import os +from contextlib import contextmanager +from typing import Iterable, Optional + +import torch +import torch.distributed as dist +from torch.optim import Optimizer + +from wall_x.config.hyperparams_config import NaiveFSDP2MuonConfig +from wall_x.trainer.optimizer.dmuon.utils import _build_ns_backend +from wall_x.trainer.optimizer.utils import register_optimizer + +try: + from torch.distributed.tensor import DTensor, Replicate +except ImportError: # pragma: no cover - FSDP2 runtime provides this. + DTensor = None + Replicate = None + +_logger = logging.getLogger(__name__) + +_BLOCKED_MUON_NAME_PARTS = ( + "embed_tokens", + "pos_embed", + "lm_head", + "norm", +) + + +def _rank() -> int: + return dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + + +def _world_size() -> int: + return dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 + + +def _env_int(name: str, default: int) -> int: + try: + return int(os.environ.get(name, default)) + except (TypeError, ValueError): + return default + + +def _env_bool(name: str, default: bool) -> bool: + value = os.environ.get(name) + if value in (None, ""): + return default + return value.strip().lower() not in {"0", "false", "no", "off"} + + +def _is_dtensor(tensor: torch.Tensor) -> bool: + return DTensor is not None and isinstance(tensor, DTensor) + + +def _replicate_dtensor(tensor: torch.Tensor) -> torch.Tensor: + if not _is_dtensor(tensor): + return tensor + mesh = tensor.device_mesh + return tensor.redistribute(placements=[Replicate()] * mesh.ndim) + + +def _full_tensor(tensor: torch.Tensor) -> torch.Tensor: + replicated = _replicate_dtensor(tensor) + if _is_dtensor(replicated): + return replicated.to_local() + return replicated + + +def _redistribute_like(full_tensor: torch.Tensor, like: torch.Tensor) -> torch.Tensor: + if not _is_dtensor(like): + return full_tensor + mesh = like.device_mesh + replicated = DTensor.from_local( + full_tensor.contiguous(), + mesh, + [Replicate()] * mesh.ndim, + run_check=False, + shape=like.shape, + stride=tuple(full_tensor.stride()), + ) + return replicated.redistribute(placements=like.placements) + + +def _is_muon_target(name: str, param: torch.Tensor) -> bool: + if not getattr(param, "requires_grad", False): + return False + if param.ndim != 2 or not name.endswith(".weight"): + return False + return not any(part in name for part in _BLOCKED_MUON_NAME_PARTS) + + +def _named_params(model: torch.nn.Module) -> list[tuple[str, torch.Tensor]]: + return [ + (name, param) + for name, param in model.named_parameters() + if getattr(param, "requires_grad", False) + ] + + +def _normalize_param_groups( + model: torch.nn.Module, + opt_cfg: NaiveFSDP2MuonConfig, + param_groups: Optional[list[dict]], +) -> list[dict]: + id_to_name = {id(param): name for name, param in _named_params(model)} + + if param_groups is None: + muon_params = [] + adamw_params = [] + for name, param in _named_params(model): + if _is_muon_target(name, param): + muon_params.append(param) + else: + adamw_params.append(param) + groups = [] + if muon_params: + groups.append( + { + "params": muon_params, + "param_names": [ + id_to_name.get(id(param), "") for param in muon_params + ], + "group_name": "base/muon", + "use_muon": True, + "lr": opt_cfg.muon_lr, + "weight_decay": opt_cfg.muon_weight_decay, + "momentum": opt_cfg.momentum, + "nesterov": opt_cfg.nesterov, + "ns_steps": opt_cfg.ns_steps, + } + ) + if adamw_params: + groups.append( + { + "params": adamw_params, + "param_names": [ + id_to_name.get(id(param), "") for param in adamw_params + ], + "group_name": "base/adamw", + "use_muon": False, + "lr": opt_cfg.adamw_lr, + "weight_decay": opt_cfg.adamw_weight_decay, + "betas": tuple(opt_cfg.adamw_betas), + "eps": opt_cfg.adamw_eps, + } + ) + return groups + + groups = [] + for idx, user_group in enumerate(param_groups): + raw_params = [ + param + for param in user_group.get("params", []) + if getattr(param, "requires_grad", False) + ] + semantic_name = str( + user_group.get("group_name", user_group.get("name", f"group_{idx}")) + ) + muon_params = [] + adamw_params = [] + for param in raw_params: + name = id_to_name.get(id(param), "") + if _is_muon_target(name, param): + muon_params.append(param) + else: + adamw_params.append(param) + + if muon_params: + groups.append( + { + "params": muon_params, + "param_names": [ + id_to_name.get(id(param), "") for param in muon_params + ], + "group_name": f"{semantic_name}/muon", + "use_muon": True, + "lr": user_group.get( + "muon_lr", user_group.get("lr", opt_cfg.muon_lr) + ), + "weight_decay": user_group.get( + "muon_weight_decay", + user_group.get("weight_decay", opt_cfg.muon_weight_decay), + ), + "momentum": user_group.get("momentum", opt_cfg.momentum), + "nesterov": user_group.get("nesterov", opt_cfg.nesterov), + "ns_steps": user_group.get("ns_steps", opt_cfg.ns_steps), + } + ) + if adamw_params: + groups.append( + { + "params": adamw_params, + "param_names": [ + id_to_name.get(id(param), "") for param in adamw_params + ], + "group_name": f"{semantic_name}/adamw", + "use_muon": False, + "lr": user_group.get( + "adamw_lr", user_group.get("lr", opt_cfg.adamw_lr) + ), + "weight_decay": user_group.get( + "adamw_weight_decay", + user_group.get("weight_decay", opt_cfg.adamw_weight_decay), + ), + "betas": tuple(user_group.get("adamw_betas", opt_cfg.adamw_betas)), + "eps": user_group.get("adamw_eps", opt_cfg.adamw_eps), + } + ) + return groups + + +def _param_layout_summary(param: torch.Tensor) -> str: + if not _is_dtensor(param): + return f"type={type(param).__name__}, shape={tuple(param.shape)}" + placements = ",".join(str(placement) for placement in param.placements) + return ( + f"type=DTensor, shape={tuple(param.shape)}, " + f"local_shape={tuple(param.to_local().shape)}, placements=[{placements}]" + ) + + +class NaiveFSDP2Muon(Optimizer): + def __init__( + self, + param_groups: list[dict], + *, + ns_backend, + record_step_profile: bool = False, + log_fn=None, + ) -> None: + defaults = {"lr": 1e-3} + super().__init__(param_groups, defaults) + self._ns = ns_backend + self._record_step_profile = bool(record_step_profile) + self._debug_steps = _env_int("WALLX_NAIVE_FSDP2_MUON_DEBUG_STEPS", 1) + self._debug_ranks = _env_bool("WALLX_NAIVE_FSDP2_MUON_DEBUG_RANKS", True) + self._step_index = 0 + self._log_fn = log_fn + self._last_step_profile: dict[str, object] = {} + self._last_step_profile_events: list[ + tuple[str, torch.cuda.Event, torch.cuda.Event] + ] = [] + self._step_profile_enabled = False + + def _log(self, message: str) -> None: + if self._log_fn is not None: + self._log_fn(message) + else: + _logger.info(message) + + def _profile_begin_step(self) -> None: + self._step_profile_enabled = ( + self._record_step_profile and torch.cuda.is_available() + ) + self._last_step_profile_events = [] + self._last_step_profile = { + "enabled": self._step_profile_enabled, + "ns_matrix_count": 0, + "ns_input_numel": 0, + } + + def _profile_event_start(self, name: str): + if not self._step_profile_enabled: + return None + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + return name, start, end + + def _profile_event_end(self, token) -> None: + if token is None: + return + _name, _start, end = token + end.record() + self._last_step_profile_events.append(token) + + @contextmanager + def _profile_phase(self, name: str): + token = self._profile_event_start(name) + try: + yield + finally: + self._profile_event_end(token) + + def _profile_add(self, key: str, value: int | float) -> None: + if not self._step_profile_enabled: + return + current = self._last_step_profile.get(key, 0) + self._last_step_profile[key] = current + value + + def consume_last_step_profile(self) -> dict[str, object]: + profile = dict(self._last_step_profile) + if not profile: + return {} + if not profile.get("enabled"): + self._last_step_profile_events = [] + return profile + event_totals: dict[str, float] = {} + for name, start, end in self._last_step_profile_events: + if isinstance(start, torch.cuda.Event) and not start.query(): + start.synchronize() + if isinstance(end, torch.cuda.Event) and not end.query(): + end.synchronize() + event_totals[name] = event_totals.get(name, 0.0) + float( + start.elapsed_time(end) + ) + for name, value in event_totals.items(): + profile[f"{name}_ms"] = round(value, 6) + profile["ns_compute_ms"] = round(event_totals.get("ns_compute", 0.0), 6) + self._last_step_profile_events = [] + return profile + + @staticmethod + def _adjusted_muon_lr(lr: float, shape: Iterable[int]) -> float: + shape = tuple(int(dim) for dim in shape) + rows = shape[0] + cols = 1 + for dim in shape[1:]: + cols *= dim + return lr * 0.2 * (max(rows, cols) ** 0.5) + + @torch.no_grad() + def step(self, closure=None): + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + self._profile_begin_step() + full_gather_count = 0 + full_gather_numel = 0 + ns_matrix_count = 0 + muon_param_count = 0 + adamw_param_count = 0 + + for group in self.param_groups: + if group.get("use_muon", False): + with self._profile_phase("muon"): + stats = self._step_muon_group(group) + full_gather_count += stats["full_gather_count"] + full_gather_numel += stats["full_gather_numel"] + ns_matrix_count += stats["ns_matrix_count"] + muon_param_count += stats["param_count"] + else: + with self._profile_phase("adamw"): + adamw_param_count += self._step_adamw_group(group) + + if self._step_index < self._debug_steps and (self._debug_ranks or _rank() == 0): + print( + "[NaiveFSDP2Muon] " + f"rank={_rank()}/{_world_size()} step={self._step_index} " + f"muon_params={muon_param_count} adamw_params={adamw_param_count} " + f"full_gathers={full_gather_count} " + f"full_gather_numel={full_gather_numel} " + f"ns_matrices={ns_matrix_count}", + flush=True, + ) + self._step_index += 1 + return loss + + def _step_muon_group(self, group: dict) -> dict[str, int]: + lr = float(group["lr"]) + weight_decay = float(group.get("weight_decay", 0.0)) + momentum = float(group["momentum"]) + nesterov = bool(group.get("nesterov", True)) + ns_steps = int(group.get("ns_steps", 5)) + stats = { + "full_gather_count": 0, + "full_gather_numel": 0, + "ns_matrix_count": 0, + "param_count": 0, + } + + for param in group["params"]: + grad = param.grad + if grad is None: + continue + stats["param_count"] += 1 + + state = self.state[param] + buf = state.get("momentum_buffer") + if ( + buf is None + or tuple(buf.shape) != tuple(grad.shape) + or _is_dtensor(buf) != _is_dtensor(grad) + ): + buf = torch.zeros_like(grad) + state["momentum_buffer"] = buf + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(grad) + if nesterov: + full_grad = _full_tensor(grad).detach() + stats["full_gather_count"] += 1 + stats["full_gather_numel"] += int(full_grad.numel()) + full_buf = _full_tensor(buf).detach() + stats["full_gather_count"] += 1 + stats["full_gather_numel"] += int(full_buf.numel()) + ns_input = full_grad.add(full_buf, alpha=momentum) + else: + ns_input = _full_tensor(buf).detach() + stats["full_gather_count"] += 1 + stats["full_gather_numel"] += int(ns_input.numel()) + if ns_input.ndim > 2: + ns_input = ns_input.view(ns_input.shape[0], -1) + with self._profile_phase("ns_compute"): + update = self._ns.local(ns_input, ns_steps) + stats["ns_matrix_count"] += 1 + self._profile_add("ns_matrix_count", 1) + self._profile_add("ns_input_numel", int(ns_input.numel())) + + full_shape = tuple(int(dim) for dim in param.shape) + update = update.view(full_shape).to( + device=ns_input.device, + dtype=( + _full_tensor(param).dtype + if not _is_dtensor(param) + else param.to_local().dtype + ), + ) + update = _redistribute_like(update, param) + + if weight_decay > 0: + param.mul_(1.0 - lr * weight_decay) + param.add_(update, alpha=-self._adjusted_muon_lr(lr, full_shape)) + if nesterov: + del full_grad, full_buf + del ns_input, update + + return stats + + def _step_adamw_group(self, group: dict) -> int: + lr = float(group["lr"]) + weight_decay = float(group.get("weight_decay", 0.0)) + beta1, beta2 = tuple(group.get("betas", (0.9, 0.999))) + eps = float(group.get("eps", 1e-8)) + param_count = 0 + + for param in group["params"]: + grad = param.grad + if grad is None: + continue + param_count += 1 + state = self.state[param] + if "step" not in state: + state["step"] = 0 + state["exp_avg"] = torch.zeros_like(grad) + state["exp_avg_sq"] = torch.zeros_like(grad) + state["step"] += 1 + step = state["step"] + exp_avg = state["exp_avg"] + exp_avg_sq = state["exp_avg_sq"] + exp_avg.lerp_(grad, 1.0 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) + + denom = exp_avg_sq.sqrt().add_(eps) + bias_correction1 = 1.0 - beta1**step + bias_correction2 = 1.0 - beta2**step + step_size = lr / bias_correction1 * (bias_correction2**0.5) + if weight_decay > 0: + param.mul_(1.0 - lr * weight_decay) + param.addcdiv_(exp_avg, denom, value=-step_size) + + return param_count + + +def get_naive_fsdp2_muon_optimizer( + model: torch.nn.Module, + *, + opt_cfg, + param_groups=None, + log_fn=None, +): + if not isinstance(opt_cfg, NaiveFSDP2MuonConfig): + raise TypeError( + "get_naive_fsdp2_muon_optimizer expects " + f"NaiveFSDP2MuonConfig, got {type(opt_cfg).__name__}" + ) + + dmuon = __import__("dmuon") + ns_backend = _build_ns_backend(dmuon, opt_cfg) + if isinstance(ns_backend, str): + ns_backend = dmuon.NewtonSchulz(backend=ns_backend) + groups = _normalize_param_groups(model, opt_cfg, param_groups) + muon_count = sum(len(group["params"]) for group in groups if group["use_muon"]) + adamw_count = sum(len(group["params"]) for group in groups if not group["use_muon"]) + if _rank() == 0: + muon_examples = [] + adamw_examples = [] + for group in groups: + target = muon_examples if group["use_muon"] else adamw_examples + for name, param in zip(group.get("param_names", []), group["params"]): + if len(target) >= 6: + break + target.append(f"{name} ({_param_layout_summary(param)})") + message = ( + "NaiveFSDP2Muon: ordinary FSDP2 wrapping; every rank all-gathers " + "each Muon gradient/momentum to full matrix and runs Newton-Schulz " + "locally; optimizer momentum state remains sharded. " + f"muon_params={muon_count}, adamw_params={adamw_count}, " + f"ns_backend={opt_cfg.ns_backend}, ns_steps={opt_cfg.ns_steps}, " + f"muon_examples={muon_examples}, adamw_examples={adamw_examples}" + ) + if log_fn is not None: + log_fn(message) + else: + _logger.info(message) + + return NaiveFSDP2Muon( + groups, + ns_backend=ns_backend, + record_step_profile=bool(getattr(opt_cfg, "record_step_profile", False)), + log_fn=log_fn, + ) + + +register_optimizer("naive_fsdp2_muon", get_naive_fsdp2_muon_optimizer) diff --git a/wall_x/trainer/optimizer/utils.py b/wall_x/trainer/optimizer/utils.py index eb1c17c9..566897cd 100644 --- a/wall_x/trainer/optimizer/utils.py +++ b/wall_x/trainer/optimizer/utils.py @@ -195,6 +195,49 @@ def build_action_expert_param_groups(model, opt_cfg, action_lr_keywords): ) +def _split_dtensor_regular_groups(params_or_groups): + """Split params into DTensor and regular-Tensor groups for fused AdamW. + + PyTorch's _fused_adamw_ kernel raises if a single param_group mixes DTensor + (from FSDP2-managed shards) and plain nn.Parameter (e.g. DDP-style root + replicas under PI0's ``replicate_non_layer_modules`` path). Splitting them + into separate param_groups lets fused AdamW handle each homogeneously. + + Accepts either a flat tensor list or a list of param_group dicts. Returns + the same shape, preserving per-group metadata (lr / weight_decay / + group_name); when a mixed group is split, the regular half gets + ``" (regular)"`` appended to ``group_name`` for log clarity. + """ + from torch.distributed.tensor import DTensor + + if not params_or_groups: + return params_or_groups + + if isinstance(params_or_groups[0], dict): + out = [] + for g in params_or_groups: + dt = [p for p in g["params"] if isinstance(p, DTensor)] + reg = [p for p in g["params"] if not isinstance(p, DTensor)] + if dt and reg: + g_dt = {**g, "params": dt} + g_reg = {**g, "params": reg} + if "group_name" in g_reg: + g_reg["group_name"] = f"{g_reg['group_name']} (regular)" + out.extend([g_dt, g_reg]) + else: + out.append(g) + return out + + dt = [p for p in params_or_groups if isinstance(p, DTensor)] + reg = [p for p in params_or_groups if not isinstance(p, DTensor)] + if dt and reg: + return [ + {"params": dt, "group_name": "base"}, + {"params": reg, "group_name": "base (regular)"}, + ] + return params_or_groups + + def get_adamw_optimizer(model, *, opt_cfg, param_groups=None): """Build AdamW from AdamWConfig.""" from wall_x.config.hyperparams_config import AdamWConfig @@ -213,6 +256,9 @@ def get_adamw_optimizer(model, *, opt_cfg, param_groups=None): # downstream for per-group lr logging. params = param_groups + # _fused_adamw_ kernel cannot mix DTensor + regular Tensor in one group. + params = _split_dtensor_regular_groups(params) + kw = { "lr": opt_cfg.learning_rate, "weight_decay": opt_cfg.weight_decay, diff --git a/wall_x/trainer/scheduler/scheduler.py b/wall_x/trainer/scheduler/scheduler.py index 2b328bb7..2a413694 100644 --- a/wall_x/trainer/scheduler/scheduler.py +++ b/wall_x/trainer/scheduler/scheduler.py @@ -1,7 +1,6 @@ -import math - +from torch.optim.lr_scheduler import LambdaLR from torch.optim import Optimizer -from torch.optim.lr_scheduler import ConstantLR, LambdaLR +import math def create_cosine_scheduler( @@ -19,7 +18,7 @@ def create_cosine_scheduler( def lr_lambda(current_step: int): if current_step < num_warmup_steps: - # Start from peak_lr / (warmup_steps + 1). + # OpenPI warmup: start from peak_lr / (warmup_steps + 1) init_lr = peak_lr / (num_warmup_steps + 1) current_lr = init_lr + (peak_lr - init_lr) * current_step / num_warmup_steps return current_lr / peak_lr # LambdaLR multiplies by base_lr @@ -53,11 +52,31 @@ def lr_lambda(current_step): def create_constant_scheduler( optimizer: Optimizer, - factor: float = 1 / 3, - total_iters: int = 5, last_epoch: int = -1, + num_warmup_steps: int = 0, ): - return ConstantLR(optimizer, factor, total_iters, last_epoch) + """Constant LR at the target value, with optional linear warmup. + + With ``num_warmup_steps <= 0`` the LR is constant at the target from step 0 + (the previous behaviour). With ``num_warmup_steps = n > 0`` the LR ramps + linearly from ``1/(n+1)`` of the target up to the full target over the first + ``n`` steps, then stays constant:: + + factor(step) = 1/(n+1) + (1 - 1/(n+1)) * step / n for step < n + factor(step) = 1.0 for step >= n + + e.g. n=4: step 0..4 -> 0.2, 0.4, 0.6, 0.8, 1.0 (x target). The initial LR is + non-zero by design. Covered by test_constant_scheduler_supports_warmup_per_lr_group. + """ + + def lr_lambda(current_step: int): + if num_warmup_steps <= 0 or current_step >= num_warmup_steps: + return 1.0 + + init_lr_factor = 1.0 / (num_warmup_steps + 1) + return init_lr_factor + (1.0 - init_lr_factor) * current_step / num_warmup_steps + + return LambdaLR(optimizer, lr_lambda, last_epoch) def get_scheduler(optimizer: Optimizer, lr_scheduler_type: str, **kwargs): diff --git a/wall_x/trainer/trainer_utils.py b/wall_x/trainer/trainer_utils.py index eab59f04..f5b491f9 100644 --- a/wall_x/trainer/trainer_utils.py +++ b/wall_x/trainer/trainer_utils.py @@ -14,7 +14,7 @@ from wall_x.model.core.action.normalizer import Normalizer from wall_x.model.qact.tokenizer_mixin import get_action_tokenizer_mixin -from wall_x.utils.metrics import dtw_distance, frechet_distance, get_action_accuracy +from wall_x.utils.metrics import get_action_accuracy, dtw_distance, frechet_distance _logger = logging.getLogger(__name__) @@ -225,18 +225,20 @@ def update_model_config(train_config, model_config): "action_horizon_flow", 32 ) - if train_config.get("_attn_implementation", None) is not None: - model_config._attn_implementation = train_config["_attn_implementation"] - - if train_config.get("attn_deterministic", None) is not None: - model_config.attn_deterministic = train_config["attn_deterministic"] - model_config.vision_config.attn_deterministic = train_config[ - "attn_deterministic" - ] + attn_impl = train_config.get( + "_attn_implementation", train_config.get("attn_implementation") + ) + if attn_impl is not None: + model_config._attn_implementation = attn_impl + + # attn_deterministic is no longer a standalone knob — it is derived from the + # single ``hyperparams.deterministic_algorithms`` switch so "turn determinism + # on" enables all paths (incl. attention kernels) at once (MR !556 review). + _det = train_config.get("hyperparams", {}).get("deterministic_algorithms", False) + model_config.attn_deterministic = _det + model_config.vision_config.attn_deterministic = _det + if _det: _logger.info("Attention is using deterministic kernel for this run") - else: - model_config.attn_deterministic = True - model_config.vision_config.attn_deterministic = True if train_config.get("noise_scheduler", None) is not None: model_config.noise_scheduler = train_config["noise_scheduler"] @@ -281,15 +283,50 @@ def print_rank_last(message): _logger.info(message) -def seed_all(seed): +def seed_all(seed, deterministic_algorithms: bool = False): + """Seed RNGs; opt into strong determinism via ``deterministic_algorithms``. + + Strong determinism (deterministic algorithms + cudnn.deterministic + the + cublas workspace pin) inserts fill/elementwise kernels and disables the cudnn + autotuner, measurably lowering training MFU. Default is performance mode; + enable it only for bitwise-reproducible experiments via the yaml switch + ``hyperparams.deterministic_algorithms`` — the single determinism knob, with + no environment-variable override. See docs qwen35_mfu_report. + """ + # RNG seed 不算「确定性算法」——任何 run 都需要 seed，否则 random op 不可复现。 + # 这部分跟 flag 无关，永远调。 np.random.seed(seed) random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False os.environ["PYTHONHASHSEED"] = str(seed) + if deterministic_algorithms: + # 严格确定性套餐：cudnn 锁 deterministic kernel + benchmark 关 + torch 全 op 走 + # 确定性实现。same-seed + same-yaml 两次 run 在 pi0 路径下 step 0 forward + # byte-equal、grad 单步 drift < 1e-3。唯一开关是 yaml 的 + # hyperparams.deterministic_algorithms（无环境变量旁路）。 + # CUBLAS_WORKSPACE_CONFIG 只 setdefault 兜底——launch 脚本显式 export 的值 + # 优先；不设则 PyTorch 第一个 cuBLAS 调用就 raise RuntimeError。 + # warn_only=True 内置：bf16 训练里总有几个 op 没确定实现（如某些 scatter_add + # 路径），warn 比 crash 实用——拿大部分 deterministic 收益即可。 + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + torch.use_deterministic_algorithms(True, warn_only=True) + else: + # perf 套餐：cudnn 自由选最快 algo（含 nondet variant），但 benchmark 不默认开： + # benchmark=True 对每个新 conv 配置做一轮算法竞速，动态 shape 负载（如 Qwen3.5 + # vision Conv3d patch_embed）几乎每步撞新 shape、反复付竞速成本，step time 3.2x + # 回退（TR-017 实测 MFU 37.7%→11.7%）。固定 shape 负载（LIBERO 224 / pi05 + # SigLIP Conv2d）warm-up 一次后收益为正——已实测过收益的固定 shape 场景可用 + # WALLX_CUDNN_BENCHMARK=1 显式开启，不走全局默认。 + torch.backends.cudnn.deterministic = False + torch.backends.cudnn.benchmark = os.environ.get( + "WALLX_CUDNN_BENCHMARK", "0" + ).lower() in ("1", "true") + torch.use_deterministic_algorithms(False) + def save_text_results_to_file( current_step, all_input_texts, all_gt_texts, all_pred_texts, save_dir @@ -423,7 +460,7 @@ def compute_action_metrics(all_preds, all_actions, config, step_log={}): return step_log -def get_warmup_cosine_schedule( +def get_openpi_cosine_schedule_with_warmup( optimizer, num_warmup_steps: int, num_training_steps: int, @@ -432,7 +469,7 @@ def get_warmup_cosine_schedule( last_epoch: int = -1, ): """ - Create a schedule with linear warmup followed by cosine decay: + Create a schedule with a learning rate that follows OpenPI style: - Warmup: linearly increases from peak_lr/(warmup_steps+1) to peak_lr - Decay: cosine decay from peak_lr to end_lr @@ -454,7 +491,7 @@ def get_warmup_cosine_schedule( def lr_lambda(current_step: int): if current_step < num_warmup_steps: - # Start from peak_lr / (warmup_steps + 1). + # OpenPI warmup: start from peak_lr / (warmup_steps + 1) init_lr = peak_lr / (num_warmup_steps + 1) current_lr = init_lr + (peak_lr - init_lr) * current_step / num_warmup_steps return current_lr / peak_lr # LambdaLR multiplies by base_lr @@ -476,6 +513,8 @@ def plot_openloop( episode_index, save_path, is_static_list=None, + dim_labels=None, + skip_all_nan=False, ): """ Plot openloop action comparison visualization. @@ -487,6 +526,12 @@ def plot_openloop( episode_index: Index of the episode being visualized save_path: Directory path to save the plot is_static_list: Optional list of booleans indicating static frames + dim_labels: Optional per-dim subplot names (e.g. ``left_x`` / + ``left_gripper`` / ``left_thumb_cmc_abduction``). Each subplot title + is this name; falls back to ``Action Dimension i`` when not given. + skip_all_nan: When True, drop dims whose L1 is all-NaN (padded DOF such + as a dexterous hand's gripper or vendor-absent joints) so they don't + add empty subplots. """ import matplotlib @@ -498,10 +543,16 @@ def plot_openloop( ), "Predicted action and ground truth action must have the same shape." dim = action_pred_list[0].shape[1] - plt.figure(figsize=(12, 4 * dim)) + dims = list(range(dim)) + if skip_all_nan: + dims = [i for i in dims if np.any(~np.isnan(l1_loss[:, i]))] + if not dims: + return + n = len(dims) + plt.figure(figsize=(12, 4 * n)) - for i in range(dim): - plt.subplot(dim, 1, i + 1) + for sp, i in enumerate(dims): + plt.subplot(n, 1, sp + 1) has_labeled_static = False for j in range(len(action_gt_list)): @@ -559,13 +610,33 @@ def plot_openloop( markersize=4, ) - plt.title(f"Action Dimension {i + 1}, L1 Loss: {l1_loss[:, i].mean():.6f}") - plt.xlabel("Number of Chunk") + # Mark chunk boundaries: each chunk is an INDEPENDENT open-loop + # re-prediction (new observation + freshly sampled flow noise), so the + # curve legitimately jumps at integer x. The line within [j, j+1] is one + # chunk's full action_horizon-step prediction, not overlaid per-frame. + for b in range(1, len(action_gt_list)): + plt.axvline( + b, + color="green", + linestyle=":", + alpha=0.35, + linewidth=1, + label="Chunk boundary" if b == 1 else None, + ) + + name = ( + dim_labels[i] + if dim_labels is not None and i < len(dim_labels) + else f"Action Dimension {i + 1}" + ) + plt.title(f"{name}, L1 Loss: {np.nanmean(l1_loss[:, i]):.6f}") + plt.xlabel("Chunk index (each unit = one action_horizon prediction)") plt.ylabel("Action Value") plt.legend() plt.suptitle( - f"Openloop Action Comparison for Episode {episode_index}, L1 Loss: {l1_loss.mean():.6f}" + f"Openloop Action Comparison for Episode {episode_index}, " + f"L1 Loss: {np.nanmean(l1_loss):.6f}" ) plt.tight_layout(rect=[0, 0, 1, 0.98]) os.makedirs(save_path, exist_ok=True) diff --git a/wall_x/trainer/utils/data.py b/wall_x/trainer/utils/data.py index 8ca974cd..8ca20325 100644 --- a/wall_x/trainer/utils/data.py +++ b/wall_x/trainer/utils/data.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Mapping from typing import Any import torch @@ -32,7 +33,14 @@ def move_batch_to_device( trainer uses pinned loaders; pass False if the dataloader hasn't pinned memory. """ - if isinstance(batch, dict): + # Use Mapping (not just dict) so dict-like batches are moved too. The lerobot + # collator returns a ``transformers.BatchFeature`` (a UserDict, not a dict); + # with a plain ``isinstance(batch, dict)`` check it would fall through + # unmoved, leaving input tensors on CPU. FSDP2 masks that (it moves module + # inputs itself), but the pure-DMuon path (no FSDP2 wrap) then hits a + # CPU/CUDA device mismatch in the embedding. Rebuilding as a plain dict is + # safe: the model is called via ``model(**batch)``. + if isinstance(batch, Mapping): return { k: move_batch_to_device(v, device, non_blocking=non_blocking) for k, v in batch.items() diff --git a/wall_x/utils/checkpoint.py b/wall_x/utils/checkpoint.py new file mode 100644 index 00000000..3b7c85b3 --- /dev/null +++ b/wall_x/utils/checkpoint.py @@ -0,0 +1,56 @@ +"""Conditional activation checkpointing helper (``maybe_checkpoint``). + +Thin wrapper around ``torch.utils.checkpoint.checkpoint`` that lets callers +toggle ckpt per-call via an ``enabled`` kwarg instead of writing if/else at +every call site. Selection policy (which layers ckpt, every-Nth, etc.) +stays with the caller — this helper only does the conditional dispatch. + +Defaults follow the wall-x convention: +- ``use_reentrant=False`` — ``use_reentrant=True`` is incompatible with the + FSDP2 nested ``fully_shard`` wrapping used here and drives activation + memory up (observed OOMs). Keep it False. +- ``preserve_rng_state=False`` — our recomputed forward has no RNG ops + (dropout=0, no stochastic depth), so the RNG state save/restore is just + overhead. +""" + +from __future__ import annotations + +from typing import Any, Callable + +import torch.utils.checkpoint as _ckpt + + +def maybe_checkpoint( + func: Callable[..., Any], + *args: Any, + enabled: bool = True, + use_reentrant: bool = False, # True is incompatible with FSDP2 nested fully_shard; see module docstring + preserve_rng_state: bool = False, + **kwargs: Any, +) -> Any: + """Call ``func`` via ``torch.utils.checkpoint`` when ``enabled``, else direct. + + Args: + func: The function to run (and recompute on backward when enabled). + *args, **kwargs: Forwarded to ``func``. + enabled: When False, calls ``func(*args, **kwargs)`` directly with no + checkpointing — the activation stays in memory. + use_reentrant: Forwarded to ``torch.utils.checkpoint``. Default False; + True is incompatible with FSDP2 nested fully_shard and OOMs + (see module docstring). Do not switch. + preserve_rng_state: Forwarded to ``torch.utils.checkpoint``. Default + False — our forward paths have no RNG ops. + + Returns: + Whatever ``func`` returns. + """ + if not enabled: + return func(*args, **kwargs) + return _ckpt.checkpoint( + func, + *args, + use_reentrant=use_reentrant, + preserve_rng_state=preserve_rng_state, + **kwargs, + ) diff --git a/wall_x/utils/constant.py b/wall_x/utils/constant.py index 4c1c71f8..046f4b29 100644 --- a/wall_x/utils/constant.py +++ b/wall_x/utils/constant.py @@ -10,6 +10,7 @@ from collections.abc import Iterable + ACTION_KEY_FULL_MAPPING: dict[str, str] = {} _ACTION_KEY_FULL_MAPPING = ACTION_KEY_FULL_MAPPING diff --git a/wall_x/utils/cudagraph_wrapper.py b/wall_x/utils/cudagraph_wrapper.py deleted file mode 100644 index 6064e175..00000000 --- a/wall_x/utils/cudagraph_wrapper.py +++ /dev/null @@ -1,278 +0,0 @@ -import os - -import torch - -ENABLE_CUDA_GRAPH = os.environ.get("ENABLE_CUDA_GRAPH", "True").lower() == "true" -# When using a single master buffer + multi-bucket graph, max batch must be pre-allocated; -# otherwise resize/reallocation after capture invalidates addresses recorded by the graph. -CUDA_GRAPH_MAX_BS = int(os.environ.get("CUDA_GRAPH_MAX_BS", "128")) -_SHAPE_GUARD_FASTPATH = ( - os.environ.get("CUDAGRAPH_DISABLE_SHAPE_GUARD_FASTPATH", "0") != "1" -) - - -class CUDAGraph_Wrapper: - """ - Bucketed CUDA Graph wrapper: - - When batch size changes, select/reuse a CUDAGraph by bucket_size - - For batch_size < bucket_size: copy inputs into static buffer [:B], return outputs [:B] after replay - """ - - def __init__( - self, - model, - warm_up_times: int = 3, - enable: bool | None = None, - batch_size_key: str | None = None, - ): - self.model = model - self.warm_up_times = warm_up_times - self._enable = ENABLE_CUDA_GRAPH if enable is None else bool(enable) - self._batch_size_key = batch_size_key or "suffix_inputs_embeds" - - # graph cache: bucket_bs -> CUDAGraph - self.graphs: dict[int, torch.cuda.CUDAGraph] = {} - # Per bucket, store views into the master buffer (graph capture depends on these views' address/shape/stride) - self.static_inputs_map: dict[int, dict[str, torch.Tensor]] = {} - self.static_output_tensor: dict[int, torch.Tensor] = {} - self.graph_pool = None - - # Master buffers (shared across all buckets) - self.graph_vars: dict[str, torch.Tensor] = {} - self._max_bs: int | None = None - - # Record non-batch dimension shape signature (globally consistent) to avoid silent errors - self._shape_signature: dict[str, tuple] = {} - # Batch dimension index per input (default 0; e.g. suffix_position_ids batch is at dim=1) - self._batch_dim: dict[str, int] = {} - self._shape_guard_verified: dict[int, bool] = {} - - def _get_batch_size(self, **kwargs) -> int: - key = self._batch_size_key - if key in kwargs: - v = kwargs[key] - if isinstance(v, torch.Tensor) and v.dim() >= 1: - return int(v.shape[0]) - raise ValueError(f"CUDAGraph_Wrapper.forward requires {key} argument") - - def _select_bucket_bs(self, bs: int) -> int: - """ - Bucket selection: - - 1, 2, 4, 8 use pow2 buckets - - >8 align to 16 (16, 32, 48, ...) - """ - if bs <= 1: - return 1 - if bs <= 2: - return 2 - if bs <= 4: - return 4 - if bs <= 8: - return 8 - return int(((bs + 15) // 16) * 16) - - def _batch_dim_for_key(self, key: str, tensor: torch.Tensor) -> int: - # Minimal special-case handling for known inputs only - # suffix_position_ids: shape [3, B, T], batch at dim=1 - if key == "suffix_position_ids" and tensor.dim() >= 2: - return 1 - return 0 - - def _view_for_bucket(self, tensor: torch.Tensor, bucket_bs: int, batch_dim: int): - if tensor.dim() == 0: - return tensor - if batch_dim == 0: - return tensor[:bucket_bs] - if batch_dim == 1: - return tensor[:, :bucket_bs] - raise ValueError(f"Unsupported batch_dim={batch_dim} for cudagraph wrapper") - - def _copy_and_pad( - self, - master: torch.Tensor, - tensor: torch.Tensor, - bs: int, - bucket_bs: int, - batch_dim: int, - ): - if tensor.dim() == 0: - master.copy_(tensor) - return - if batch_dim == 0: - master[:bs].copy_(tensor) - if bs < bucket_bs: - master[bs:bucket_bs].zero_() - return - if batch_dim == 1: - master[:, :bs].copy_(tensor) - if bs < bucket_bs: - master[:, bs:bucket_bs].zero_() - return - raise ValueError(f"Unsupported batch_dim={batch_dim} for cudagraph wrapper") - - def _allocate_master_tensor( - self, key: str, tensor: torch.Tensor, max_bs: int, bs: int - ) -> torch.Tensor: - """ - Allocate a master buffer for one input (single large buffer): - - Extend only the batch dimension to max_bs; other dims unchanged - - Copy the current bs valid range first; zero the rest (avoid stale values on replay) - """ - if tensor.dim() == 0: - return tensor.clone() - batch_dim = self._batch_dim_for_key(key, tensor) - shape = list(tensor.shape) - if batch_dim >= len(shape): - raise ValueError( - f"Invalid batch_dim={batch_dim} for key={key}, tensor.shape={tuple(tensor.shape)}" - ) - shape[batch_dim] = max_bs - buf = torch.zeros(tuple(shape), device=tensor.device, dtype=tensor.dtype) - # Initial copy: pad to current bucket (at least cover the valid bs range) - self._copy_and_pad(buf, tensor, bs=bs, bucket_bs=bs, batch_dim=batch_dim) - return buf - - def _ensure_master_buffers(self, bs: int, **kwargs): - """ - Ensure master buffers are allocated once. - - max_bs from CUDA_GRAPH_MAX_BS; if unset, use the first bucket_bs as max_bs - and disallow exceeding it later (resize would invalidate captured graphs). - """ - if self._max_bs is not None: - return - - bucket_bs = self._select_bucket_bs(bs) - max_bs = CUDA_GRAPH_MAX_BS if CUDA_GRAPH_MAX_BS > 0 else bucket_bs - self._max_bs = int(max_bs) - - for k, v in kwargs.items(): - if not isinstance(v, torch.Tensor): - raise TypeError( - f"CUDAGraph_Wrapper only supports Tensor kwargs, got {k}={type(v)}" - ) - batch_dim = self._batch_dim_for_key(k, v) if v.dim() > 0 else 0 - self._batch_dim[k] = batch_dim - # Record non-batch dimensions (globally consistent) - if v.dim() == 0: - self._shape_signature[k] = tuple() - else: - sig = list(v.shape) - sig.pop(batch_dim) - self._shape_signature[k] = tuple(sig) - self.graph_vars[k] = self._allocate_master_tensor( - k, v, max_bs=self._max_bs, bs=bs - ) - - def _initialize_bucket(self, bucket_bs: int, bs: int, **kwargs): - assert self._max_bs is not None - self._shape_guard_verified[bucket_bs] = False - if bucket_bs > self._max_bs: - raise ValueError( - f"CUDAGraph bucket_bs({bucket_bs}) exceeds master max_bs({self._max_bs}). " - f"Set CUDA_GRAPH_MAX_BS >= {bucket_bs}." - ) - - print(f"[CUDAGraph] Initializing bucket_bs={bucket_bs} (current bs={bs}) ...") - - # Build view map for this bucket (fixed address/shape/stride) - static_map: dict[str, torch.Tensor] = {} - for k, v in kwargs.items(): - # Shape guard: non-batch dims must match - expected = self._shape_signature.get(k) - if v.dim() == 0: - got = tuple() - else: - batch_dim = self._batch_dim.get(k, self._batch_dim_for_key(k, v)) - got_list = list(v.shape) - got_list.pop(batch_dim) - got = tuple(got_list) - if expected is not None and got != expected: - raise ValueError( - f"CUDAGraph bucket init: non-batch dim changed: key={k}, expected={expected}, got={got}" - ) - master = self.graph_vars[k] - batch_dim = self._batch_dim.get(k, 0) - static_map[k] = self._view_for_bucket( - master, bucket_bs=bucket_bs, batch_dim=batch_dim - ) - - # Warmup (stabilize kernels/caches) - out = None - for _ in range(self.warm_up_times): - out = self.forward_naive(**static_map) - assert isinstance(out, torch.Tensor) - - # Output master buffer (allocate once) - if "_outputs" not in self.graph_vars: - out_shape = (self._max_bs,) + tuple(out.shape[1:]) - self.graph_vars["_outputs"] = torch.empty( - out_shape, device=out.device, dtype=out.dtype - ) - - self.static_inputs_map[bucket_bs] = static_map - self.static_output_tensor[bucket_bs] = self.graph_vars["_outputs"][:bucket_bs] - - # capture - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, pool=self.graph_pool): - self.static_output_tensor[bucket_bs].copy_(self.forward_naive(**static_map)) - if self.graph_pool is None: - self.graph_pool = graph.pool() - self.graphs[bucket_bs] = graph - torch.cuda.synchronize() - print(f"[CUDAGraph] Bucket {bucket_bs} captured.") - - def forward_naive(self, **kwargs): - return self.model(**kwargs) - - def forward(self, **kwargs): - if not self._enable: - return self.forward_naive(**kwargs) - bs = self._get_batch_size(**kwargs) - bucket_bs = self._select_bucket_bs(bs) - - # Initialize master buffers (first call) - self._ensure_master_buffers(bs=bs, **kwargs) - assert self._max_bs is not None - if bucket_bs > self._max_bs: - # Do not exceed master max_bs at runtime (resize would invalidate captured graphs) - return self.forward_naive(**kwargs) - - # lazy capture for this bucket - if bucket_bs not in self.graphs: - self._initialize_bucket(bucket_bs=bucket_bs, bs=bs, **kwargs) - - # Shape guard: non-batch dim change is unsafe for cudagraph; fallback - # fast-path: skip guard if already verified for this bucket (stable-state optimization) - if not ( - _SHAPE_GUARD_FASTPATH and self._shape_guard_verified.get(bucket_bs, False) - ): - for k, v in kwargs.items(): - if not isinstance(v, torch.Tensor): - raise TypeError( - f"CUDAGraph_Wrapper only supports Tensor kwargs, got {k}={type(v)}" - ) - expected = self._shape_signature.get(k) - if v.dim() == 0: - got = tuple() - else: - batch_dim = self._batch_dim.get(k, self._batch_dim_for_key(k, v)) - got_list = list(v.shape) - got_list.pop(batch_dim) - got = tuple(got_list) - if expected is not None and got != expected: - return self.forward_naive(**kwargs) - # Mark verified for fast-path on subsequent calls - self._shape_guard_verified[bucket_bs] = True - - # Copy inputs into master buffer (per batch_dim, zero padding) - for k, tensor in kwargs.items(): - master = self.graph_vars[k] - batch_dim = self._batch_dim.get(k, 0) - self._copy_and_pad( - master, tensor, bs=bs, bucket_bs=bucket_bs, batch_dim=batch_dim - ) - - graph = self.graphs[bucket_bs] - graph.replay() - return self.graph_vars["_outputs"][:bs] diff --git a/wall_x/utils/logger.py b/wall_x/utils/logger.py index c7d2afac..eac25332 100644 --- a/wall_x/utils/logger.py +++ b/wall_x/utils/logger.py @@ -54,7 +54,7 @@ def __init__( self._logger = logger - # Thin pass-throughs - callers use standard logging verbs. + # Thin pass-throughs — callers use standard logging verbs. def info(self, msg, *args, **kwargs): self._logger.info(msg, *args, **kwargs) diff --git a/wall_x/utils/metrics.py b/wall_x/utils/metrics.py index e2ffb0fb..6904f8b5 100644 --- a/wall_x/utils/metrics.py +++ b/wall_x/utils/metrics.py @@ -1,6 +1,5 @@ -from typing import List - import torch +from typing import List def get_action_accuracy( diff --git a/wall_x/utils/profile_zones.py b/wall_x/utils/profile_zones.py new file mode 100644 index 00000000..0bac2724 --- /dev/null +++ b/wall_x/utils/profile_zones.py @@ -0,0 +1,39 @@ +"""Named regions for torch.profiler — central registry of `record_function` zones. + +`torch.profiler.record_function(name)` writes a `user_annotation` slice into the +profiler trace. We use it to mark phase boundaries that would otherwise be +invisible when `with_stack=False` (the lightweight profiling mode). + +Call sites use the `zone(...)` alias and the `Z` namespace: + + from wall_x.utils.profile_zones import zone, Z + + with zone(Z.FORWARD): + outputs = self.adapter.forward(...) + +Keeping the names here makes them grep-able and prevents string-typo drift +across the ~5 files that emit them. +""" + +import torch + +# Alias — same call site ergonomics as `torch.no_grad()`. +zone = torch.profiler.record_function + + +class Z: + # All zone names use a dotted prefix so a trace viewer can disambiguate + # them with exact / prefix LIKE patterns instead of substring matches. + + # ── Training phase (fsdp_trainer.train_loop) ── + TRAIN_STEP = "train.step" + DATA_TO_DEVICE = "train.data_to_device" + FORWARD = "train.forward" + BACKWARD = "train.backward" + GRAD_NORMS = "train.grad_norms" + GRAD_CLIP = "train.grad_clip" + OPTIMIZER_STEP = "train.optimizer_step" + ZERO_GRAD = "train.zero_grad" + LOGGING = "train.logging" + SAVE_CHECKPOINT = "train.save_checkpoint" + diff --git a/wall_x/utils/timers.py b/wall_x/utils/timers.py index e7a05dcd..2ebbcc58 100644 --- a/wall_x/utils/timers.py +++ b/wall_x/utils/timers.py @@ -1,13 +1,13 @@ -import logging -import os import time +from torch.cuda import nvtx from abc import ABC, abstractmethod -from contextlib import nullcontext -from functools import wraps from typing import List import torch -from torch.cuda import nvtx +from functools import wraps +from contextlib import nullcontext +import logging +import os logger = logging.getLogger(__name__) @@ -115,7 +115,7 @@ def elapsed(self, reset=True, barrier=False): class DummyTimer(TimerBase): - """Dummy Timer - no-op placeholder used when log level exceeds threshold.""" + """Dummy Timer — no-op placeholder used when log level exceeds threshold.""" def __init__(self): super().__init__("dummy timer") @@ -228,6 +228,8 @@ def __init__(self, log_level, log_option): self._max_log_level = 2 def __call__(self, name, log_level=None): + if self._log_level < 0: + return self._dummy_timer if name in self._timers: if log_level is not None: assert log_level == self._log_levels[name], ( @@ -368,6 +370,8 @@ def get_all_timers_string( reset: Reset each timer after reading its elapsed time. barrier: Synchronize across ranks before gathering times. """ + if self._log_level < 0: + return None if names is None: names = list(self._timers.keys()) @@ -402,6 +406,8 @@ def log( reset: Reset each timer after reading. barrier: Synchronize across ranks first. """ + if self._log_level < 0: + return output_string = self.get_all_timers_string(names, normalizer, reset, barrier) world_size = _get_world_size() current_rank = _get_rank() @@ -430,6 +436,8 @@ def write( reset: Reset each timer after reading. barrier: Synchronize across ranks first. """ + if self._log_level < 0: + return assert normalizer > 0.0 name_to_min_max_time = self._get_global_min_max_time( names, reset, barrier, normalizer diff --git a/workspace/README.md b/workspace/README.md index 7d103823..46c77843 100644 --- a/workspace/README.md +++ b/workspace/README.md @@ -273,29 +273,22 @@ bash scripts/run_serving.sh \ --port 32195 ``` -By default this wrapper returns raw model action chunks, which is the expected -mode for open-loop evaluation. Pass `--serialize-actions` if your client expects -robot-serialized actions. +This (vanilla) server returns raw model action chunks — the expected mode for +open-loop evaluation. Robot-side serialization/interpolation is left to the +client, since the real-robot preprocessors are not part of the open-source cut. Or invoke the serving module directly (adjust parameters as needed): ```bash -export ENABLE_CUDA_GRAPH=True -export ENABLE_EXPERIMENTAL_INFERENCE_ENGINE=True - CKPT_PATH=/path/to/checkpoint -python -m wall_x._vendor.harrix.serving.launch_serving \ - --env X2ROBOT \ +python -m wall_x._vendor.harrix.serving.launch_vanilla_serving \ + --checkpoint-path "$CKPT_PATH" \ + --train-config-path /path/to/my_robot_config.yml \ --port 32195 \ - --no-serialize-actions \ - model-config:server-model-config \ - --model-config.checkpoint-path "$CKPT_PATH" \ - --model-config.train-config-path /path/to/my_robot_config.yml \ - --model-config.action-horizon 32 \ - --model-config.robot-action-interpolate-multiplier 1 \ - --model-config.robot-action-end-ratio 1.0 \ - --model-config.robot-type desktop + --action-horizon 32 \ + --cam-names face_view right_wrist_view \ + --image-passing-mode base64 ``` Clients connect at `ws://127.0.0.1:32195` to send observations and receive predicted actions. diff --git a/workspace/example/lerobot/qwen2_5_lerobot_template.yml b/workspace/example/lerobot/qwen2_5_lerobot_template.yml index e44d62a2..9f712c90 100644 --- a/workspace/example/lerobot/qwen2_5_lerobot_template.yml +++ b/workspace/example/lerobot/qwen2_5_lerobot_template.yml @@ -39,7 +39,13 @@ hyperparams: seed: 10233 optimizer: optimizer_type: dmuon - learning_rate: 0.0001 + # NOTE: DMuon reads muon_lr / adamw_lr, NOT learning_rate. Set them + # explicitly to the recommended LIBERO values below; do NOT rely on the + # built-in defaults (muon_lr=0.02 / adamw_lr=1e-3), which are far too large + # for this task and make training diverge (0% eval success). + muon_lr: 0.0001 # 1e-4 (recommended) + adamw_lr: 0.00005 # 5e-5 (recommended) + learning_rate: 0.0001 # unused by DMuon; kept for schema compatibility max_grad_norm: 1.0 enable_grad_clip: true betas: [0.9, 0.95] @@ -55,6 +61,12 @@ distributed: use_fsdp: true use_mixed_precision: true bf16: true + # ZeRO-2: replicate params, shard only gradients + optimizer state — skips the + # per-layer param all-gather of the full_shard (ZeRO-3) default, giving much + # higher throughput/MFU. Safe for the 3B model on any >=24GB GPU. Use + # "full_shard" only if you OOM. The fused CUDA kernels are compiled in and + # used automatically; raise batch_size_per_gpu + num_workers for higher MFU. + fsdp_sharding_strategy: shard_grad_op data: dataset_type: lerobot diff --git a/workspace/example/libero.yml b/workspace/example/libero.yml index 5a57e326..578a8084 100644 --- a/workspace/example/libero.yml +++ b/workspace/example/libero.yml @@ -62,6 +62,11 @@ model: hyperparams: num_epoch: 100 + # Throughput note: the fused CUDA kernels (rmsnorm, silu, m_rope, MoE + # pack/unpack, ...) are compiled at install time and used automatically. + # To push training MFU on 80GB GPUs, raise batch_size_per_gpu (e.g. 24-36), + # bump data.num_workers (e.g. 8), and keep distributed.fsdp_sharding_strategy + # at shard_grad_op (below). The conservative defaults here run on modest GPUs. batch_size_per_gpu: 4 gradient_accumulation_steps: 4 seed: 10222 @@ -83,6 +88,12 @@ distributed: use_fsdp: true use_mixed_precision: true bf16: true + # ZeRO-2: replicate params, shard only gradients + optimizer state. This + # skips the per-layer parameter all-gather that the full_shard (ZeRO-3) + # default pays every step, so throughput/MFU is substantially higher. Safe + # for the 3B model on any >=24GB GPU (replicated params are ~6GB in bf16). + # Switch to "full_shard" only if you OOM (many GPUs + very large batch). + fsdp_sharding_strategy: shard_grad_op data: dataset_type: lerobot diff --git a/workspace/example/libero_dmuon.yml b/workspace/example/libero_dmuon.yml new file mode 100644 index 00000000..125a5cf5 --- /dev/null +++ b/workspace/example/libero_dmuon.yml @@ -0,0 +1,149 @@ +# LIBERO single-arm finetune example using DMuon (RECOMMENDED optimizer). +# +# Same as libero.yml but with the DMuon optimizer, which converges +# significantly better than AdamW on LIBERO in our experiments. For best +# results we also recommend a global batch size of ~128 (batch_size_per_gpu x +# gradient_accumulation_steps x num_gpus) -- raise gradient_accumulation_steps +# to reach it on smaller GPU counts. +# +# Replace every /path/to/* placeholder before training: +# +# model.config_path -> model architecture JSON (mot_flash_mask_causal_xloss.json) +# model.processor_path -> Qwen2.5-VL-3B-Instruct directory +# model.pretrained_path -> same as processor_path, or HuggingFace cache path +# data.lerobot_config.repo_id -> local LeRobot dataset root (libero_all) +# data.norm_stats_path -> q01/q99 normalization JSON +# checkpoint.save_path -> writable directory for training checkpoints +# checkpoint.resume_from -> Wall-OSS-0.5 pretrained .safetensors or checkpoint directory +# +# Compute norm stats first: +# +# python scripts/compute_norm_stats.py \ +# --train_config workspace/example/libero_dmuon.yml \ +# --data_root /path/to/libero_all \ +# --output_path /path/to/libero_all_norm_stats.json +# +# Launch training (from repo root): +# +# torchrun --nproc_per_node= wall_x/trainer/fsdp_trainer/train_fsdp.py \ +# --config workspace/example/libero_dmuon.yml +# +# Strategy: keep dof / agent_pos totals at 26 to match the pretraining action space +# via ``action_padding``. The lerobot collator right-pads libero's 7-dim action / +# 8-dim state with zeros; loss does not flow through the padded tail. + +model_type: qwen2_5 + +task: + # Libero delta action: pos3 + rot3 + gripper1 = 7, plus action_padding(19) = 26. + dof_config: + master_right_ee_cartesian_pos: 3 # delta position + master_right_ee_rotation: 3 # delta rotation (ZYX euler) + master_right_gripper: 1 + action_padding: 19 + ar_dof_config: + master_right_ee_cartesian_pos: 3 + master_right_ee_rotation: 3 + master_right_gripper: 1 + action_padding: 19 + # State: pos3 + rot3 + gripper2 = 8, plus action_padding(18) = 26. + agent_pos_config: + follow_right_ee_cartesian_pos: 3 + follow_right_ee_rotation: 3 + follow_right_gripper: 2 + action_padding: 18 + action_horizon: 10 + action_horizon_flow: 10 + use_state_string_representation: false + +model: + backbone: qwen2_5 + config_path: /path/to/wall-oss-0.5/config.json + processor_path: /path/to/Qwen2.5-VL-3B-Instruct + pretrained_path: /path/to/Qwen2.5-VL-3B-Instruct + attn_deterministic: true + use_ema: false + flow_loss_weight: 1.0 + ar_loss_weight: 0.01 + +hyperparams: + num_epoch: 100 + # Throughput note: the fused CUDA kernels (rmsnorm, silu, m_rope, MoE + # pack/unpack, ...) are compiled at install time and used automatically. + # To push training MFU on 80GB GPUs, raise batch_size_per_gpu (e.g. 24-36), + # bump data.num_workers (e.g. 8), and keep distributed.fsdp_sharding_strategy + # at shard_grad_op (below). The conservative defaults here run on modest GPUs. + batch_size_per_gpu: 4 + gradient_accumulation_steps: 4 + seed: 10222 + optimizer: + optimizer_type: dmuon + # DMuon reads muon_lr / adamw_lr, NOT learning_rate. These are the + # recommended LIBERO values; do NOT rely on the built-in defaults + # (muon_lr=0.02 / adamw_lr=1e-3), which are far too large and diverge. + muon_lr: 0.0001 # 1e-4 (recommended) + adamw_lr: 0.00005 # 5e-5 (recommended) + momentum: 0.95 + ns_steps: 5 + muon_weight_decay: 0.0 + adamw_betas: [0.9, 0.999] + adamw_weight_decay: 0.01 + adamw_eps: 1.0e-8 + max_grad_norm: 1.0 + enable_grad_clip: true + scheduler: + scheduler_type: cosine + num_warmup_steps: 1000 + num_training_steps: 200000 + min_lr: 1.0e-6 + +distributed: + use_fsdp: true + use_mixed_precision: true + bf16: true + # ZeRO-2: replicate params, shard only gradients + optimizer state. This + # skips the per-layer parameter all-gather that the full_shard (ZeRO-3) + # default pays every step, so throughput/MFU is substantially higher. Safe + # for the 3B model on any >=24GB GPU (replicated params are ~6GB in bf16). + # Switch to "full_shard" only if you OOM (many GPUs + very large batch). + fsdp_sharding_strategy: shard_grad_op + +data: + dataset_type: lerobot + lerobot_config: + repo_id: /path/to/libero_all + root: null + key_mappings: + # libero_all v3.0 only has faceImg + rightImg (no leftImg). + camera: + observation.images.faceImg: face_view + observation.images.rightImg: right_wrist_view + state: observation.state + action: action + norm_stats_path: /path/to/libero_all_norm_stats.json + train_test_split: 0.95 + num_workers: 4 + max_length: 1024 + resolution: + face_view: 256 + right_wrist_view: 256 + +logging: + log_name: libero_ft + log_project: lerobot_libero_ft + log_entity: your_wandb_entity + use_wandb: true + log_interval: 10 + save_interval: 2000 + val_interval: 1000000 + epoch_save_interval: 1 + +checkpoint: + save_path: /path/to/libero + # Single-file .safetensors loads as pretrain weights before FSDP wrapping. + # Use a checkpoint directory for full resume (optimizer / scheduler / RNG). + resume_from: /path/to/wall-oss-0.5/model.safetensors + +debug: + profile: false + nvtx: false