diff --git a/.claude/docs/backends/megatron.md b/.claude/docs/backends/megatron.md index 7c1d62c6b2..81593455db 100644 --- a/.claude/docs/backends/megatron.md +++ b/.claude/docs/backends/megatron.md @@ -24,6 +24,18 @@ Key strategies: Note: Sequence parallelism is auto-enabled when `tensor_model_parallel_size > 1` — there is no separate config field for it. +## Expert MXFP8 + +Set `trainer.policy.model.expert_mxfp8.enabled=true` on SM100/SM103 to run +routed expert GEMMs with Transformer Engine MXFP8 and configure vLLM for +expert-only online MXFP8. Checkpoints and weight sync remain high precision. +Set `trainer.policy.model.expert_mxfp8.persistent=true` with +`trainer.policy.megatron_config.ddp_config.fp8_param_gather=true` to keep +routed-expert primary parameters in MXFP8 between optimizer steps. Set +`generator.inference_engine.fp8_weight_sync_mode=serialized_mxfp8` to send +routed experts to vLLM as MXFP8 data and E8M0 scales. Unmerged LoRA is not +supported. + ## Test Requirements Megatron GPU tests need: `NVTE_FLASH_ATTN=0` diff --git a/docs/content/docs/examples/quantized_rollouts.mdx b/docs/content/docs/examples/quantized_rollouts.mdx index 558678ecfb..3faab5b192 100644 --- a/docs/content/docs/examples/quantized_rollouts.mdx +++ b/docs/content/docs/examples/quantized_rollouts.mdx @@ -31,6 +31,41 @@ generator.inference_engine.engine_init_kwargs.quantization=fp8 This uses vLLM's [online dynamic FP8 quantization](https://docs.vllm.ai/en/latest/features/quantization/fp8.html), so no calibration data or pre-quantized checkpoint is required. +## Expert-only MXFP8 + +On Blackwell GPUs, the Megatron backend can run routed MoE expert GEMMs in +MXFP8 while vLLM serves the same experts with online MXFP8 quantization: + +```bash +trainer.strategy=megatron \ +trainer.policy.model.expert_mxfp8.enabled=true \ +trainer.policy.model.expert_mxfp8.persistent=true \ +trainer.policy.megatron_config.ddp_config.fp8_param_gather=true \ +generator.inference_engine.fp8_weight_sync_mode=serialized_mxfp8 +``` + +Persistent expert weights are quantized after each optimizer step and reused by +Transformer Engine until the next update. Weight sync sends routed experts as +MXFP8 data with E8M0 scales, while checkpoints remain high precision. Attention, +dense layers, shared experts, the router, and the LM head remain high precision. + +This mode requires SM100 or SM103 and does not support unmerged LoRA. Set +`training=false` or `rollout=false` under `expert_mxfp8` to enable only one +side for benchmarking. Set `persistent=false` to quantize expert weights during +each forward pass. Set `fp8_weight_sync_mode=null` to let vLLM quantize experts +after each weight sync. + +Run the Modal benchmark with identical BF16 and MXFP8 workloads: + +```bash +uv run --isolated --with modal modal run --detach \ + examples/train/megatron/modal_expert_mxfp8_benchmark.py \ + --mode both --steps 10 +``` + +Use `--mode bf16` or `--mode mxfp8` for one side. Set `MODAL_GPU` to override +the default `B200:8` allocation. + ## Enabling off-policy correction (TIS) To apply TIS, we need the inference engine to return the rollout logprobs for the generated tokens, and we configure the correction on the policy loss: diff --git a/examples/train/megatron/modal_expert_mxfp8_benchmark.py b/examples/train/megatron/modal_expert_mxfp8_benchmark.py new file mode 100644 index 0000000000..a7a03611af --- /dev/null +++ b/examples/train/megatron/modal_expert_mxfp8_benchmark.py @@ -0,0 +1,254 @@ +"""Benchmark BF16 versus expert-only MXFP8 on DAPO with one Modal node. + +Run: + uv run --isolated --with modal modal run --detach \ + examples/train/megatron/modal_expert_mxfp8_benchmark.py \ + --mode both --steps 10 +""" + +from __future__ import annotations + +import os +import pathlib + +import modal + +APP_NAME = "skyrl-expert-mxfp8-benchmark" +MODEL = "Qwen/Qwen3-30B-A3B" +REMOTE_REPO = "/root/SkyRL" +HF_HOME = "/root/hf-cache" +DATA_DIR = "/root/data/dapo" +TRAIN_FILE = f"{DATA_DIR}/dapo-math-17k-cleaned.parquet" +VAL_FILE = f"{DATA_DIR}/aime-2024-cleaned.parquet" +PROFILE_ROOT = "/root/profiles" # profiler +PROFILE_COMMIT_INTERVAL_S = 120 # profiler +GPU = os.environ.get("MODAL_GPU", "B200:8") + + +def _repo_root() -> pathlib.Path: + for start in (pathlib.Path(__file__).resolve(), pathlib.Path.cwd().resolve()): + candidate = start if start.is_dir() else start.parent + for path in (candidate, *candidate.parents): + if (path / "pyproject.toml").exists() and (path / "skyrl").exists(): + return path + raise RuntimeError("Run the benchmark from a SkyRL checkout") + + +repo_root = _repo_root() +hf_volume = modal.Volume.from_name("skyrl-hf-cache", create_if_missing=True) +data_volume = modal.Volume.from_name("skyrl-expert-mxfp8-data", create_if_missing=True) +profile_volume = modal.Volume.from_name("skyrl-expert-mxfp8-profiles", create_if_missing=True) # profiler + +image = ( + modal.Image.from_registry("nvidia/cuda:12.9.1-devel-ubuntu22.04", add_python="3.12") + .apt_install("git", "curl", "build-essential", "ca-certificates", "libnuma1", "numactl") + .pip_install("huggingface-hub") + .run_commands("curl -LsSf https://astral.sh/uv/install.sh | sh") + .env( + { + "PATH": "/root/.local/bin:/usr/local/cuda/bin:${PATH}", + "HF_HOME": HF_HOME, + "HF_XET_HIGH_PERFORMANCE": "1", + "NVTE_FLASH_ATTN": "0", + "SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S": "1800", + "SKYRL_DUMP_INFRA_LOG_TO_STDOUT": "1", + "UV_LINK_MODE": "copy", + "UV_PROJECT_ENVIRONMENT": f"{REMOTE_REPO}/.venv", + "VLLM_USE_FLASHINFER_MOE_FP16": "0", + "VLLM_USE_FLASHINFER_SAMPLER": "0", + } + ) + .add_local_dir( + str(repo_root), + REMOTE_REPO, + copy=True, + ignore=[".venv", ".git", "**/__pycache__"], + ) + .workdir(REMOTE_REPO) + .run_commands("uv sync --extra megatron", gpu="any") + .run_commands(f"rm -rf {HF_HOME}") +) + +app = modal.App(APP_NAME) + + +def _run(mode: str, steps: int, max_prompt_length: int, max_generate_length: int) -> None: + import subprocess + import time + + enabled = str(mode == "mxfp8").lower() + run_name = f"qwen3-30b-a3b-dapo-{mode}" + command = [ + "uv", + "run", + "--frozen", + "--no-sync", + "--extra", + "megatron", + "-m", + "examples.train.algorithms.dapo.main_dapo", + f"data.train_data=['{TRAIN_FILE}']", + f"data.val_data=['{VAL_FILE}']", + "trainer.strategy=megatron", + f"trainer.policy.model.path={MODEL}", + f"trainer.policy.model.expert_mxfp8.enabled={enabled}", + "trainer.placement.colocate_all=true", + "trainer.placement.policy_num_nodes=1", + "trainer.placement.policy_num_gpus_per_node=8", + "generator.inference_engine.num_engines=1", + "generator.inference_engine.tensor_parallel_size=1", + "generator.inference_engine.data_parallel_size=8", + "generator.inference_engine.expert_parallel_size=8", + "generator.inference_engine.distributed_executor_backend=mp", + "generator.inference_engine.enforce_eager=false", + "generator.inference_engine.weight_sync_backend=nccl", + "generator.inference_engine.gpu_memory_utilization=0.6", + "trainer.policy.megatron_config.tensor_model_parallel_size=1", + "trainer.policy.megatron_config.pipeline_model_parallel_size=1", + "trainer.policy.megatron_config.context_parallel_size=1", + "trainer.policy.megatron_config.expert_model_parallel_size=8", + "trainer.policy.megatron_config.expert_tensor_parallel_size=1", + "trainer.algorithm.advantage_estimator=grpo", + "trainer.algorithm.policy_loss_type=dual_clip", + "trainer.algorithm.loss_reduction=token_mean_legacy", + "trainer.algorithm.overlong_buffer_len=4096", + "trainer.algorithm.overlong_buffer_penalty_factor=1.0", + "trainer.algorithm.eps_clip_low=0.2", + "trainer.algorithm.eps_clip_high=0.28", + "trainer.algorithm.use_kl_loss=false", + "trainer.train_batch_size=64", + "trainer.policy_mini_batch_size=64", + "trainer.micro_forward_batch_size_per_gpu=1", + "trainer.micro_train_batch_size_per_gpu=1", + "trainer.eval_before_train=false", + "trainer.eval_interval=0", + "trainer.ckpt_interval=0", + "trainer.hf_save_interval=0", + "trainer.resume_mode=none", + f"trainer.max_training_steps={steps}", + f"trainer.max_prompt_length={max_prompt_length}", + f"generator.sampling_params.max_generate_length={max_generate_length}", + "generator.apply_overlong_filtering=true", + "generator.n_samples_per_prompt=8", + "generator.batched=true", + "environment.env_class=aime", + # profiler + "trainer.policy.torch_profiler_config.enable=false", + "trainer.policy.torch_profiler_config.ranks=[0]", + f"trainer.policy.torch_profiler_config.save_path={PROFILE_ROOT}/trainer/{mode}", + "trainer.policy.torch_profiler_config.skip_first=3", + "trainer.policy.torch_profiler_config.wait=0", + "trainer.policy.torch_profiler_config.warmup=1", + "trainer.policy.torch_profiler_config.active=2", + "trainer.policy.torch_profiler_config.repeat=1", + "trainer.policy.torch_profiler_config.record_shapes=false", + "trainer.policy.torch_profiler_config.with_stack=false", + "trainer.logger=wandb", + "trainer.project_name=skyrl-expert-mxfp8-dapo", + f"trainer.run_name={run_name}", + ] + if mode == "mxfp8": + command.extend( + [ + "trainer.policy.model.expert_mxfp8.persistent=true", + "trainer.policy.megatron_config.ddp_config.fp8_param_gather=true", + "trainer.policy.megatron_config.ddp_config.reuse_grad_buf_for_mxfp8_param_ag=true", + "trainer.policy.megatron_config.optimizer_config_kwargs.fp8_recipe=mxfp8", + "trainer.policy.megatron_config.optimizer_config_kwargs.reuse_grad_buf_for_mxfp8_param_ag=true", + "generator.inference_engine.fp8_weight_sync_mode=serialized_mxfp8", + ] + ) + started = time.perf_counter() + process = subprocess.Popen(command, cwd=REMOTE_REPO) + while True: + try: + returncode = process.wait(timeout=PROFILE_COMMIT_INTERVAL_S) + break + except subprocess.TimeoutExpired: + profile_volume.commit() # profiler + if returncode != 0: + raise subprocess.CalledProcessError(returncode, command) + print(f"{mode} benchmark completed in {time.perf_counter() - started:.1f}s") + + +@app.function( + image=image, + cpu=4, + memory=16384, + volumes={HF_HOME: hf_volume, "/root/data": data_volume}, + timeout=2 * 60 * 60, +) +def prepare_assets() -> None: + import shutil + import subprocess + + from huggingface_hub import hf_hub_download, snapshot_download + + snapshot_download(MODEL) + hf_volume.commit() + os.makedirs(DATA_DIR, exist_ok=True) + if not os.path.exists(TRAIN_FILE) or not os.path.exists(VAL_FILE): + train_download = hf_hub_download( + repo_id="BytedTsinghua-SIA/DAPO-Math-17k", + filename="data/dapo-math-17k.parquet", + repo_type="dataset", + ) + val_download = hf_hub_download( + repo_id="BytedTsinghua-SIA/AIME-2024", + filename="data/aime-2024.parquet", + repo_type="dataset", + ) + shutil.copyfile(train_download, f"{DATA_DIR}/dapo-math-17k.parquet") + shutil.copyfile(val_download, f"{DATA_DIR}/aime-2024.parquet") + subprocess.run( + [ + "uv", + "run", + "--frozen", + "--no-sync", + "--extra", + "megatron", + "-m", + "examples.train.algorithms.dapo.data_preprocess_dapo_aime", + "--data-dir", + DATA_DIR, + ], + cwd=REMOTE_REPO, + check=True, + ) + data_volume.commit() + + +@app.function( + image=image, + gpu=GPU, + secrets=[modal.Secret.from_name("wandb-secret")], + volumes={ + HF_HOME: hf_volume, + "/root/data": data_volume, + PROFILE_ROOT: profile_volume, + }, # profiler + timeout=24 * 60 * 60, +) +def benchmark(mode: str, steps: int, max_prompt_length: int, max_generate_length: int) -> None: + try: + _run(mode, steps, max_prompt_length, max_generate_length) + finally: + profile_volume.commit() # profiler + + +@app.local_entrypoint() +def main( + mode: str = "both", + steps: int = 10, + max_prompt_length: int = 2048, + max_generate_length: int = 4096, +) -> None: + if mode not in {"bf16", "mxfp8", "both"}: + raise ValueError("mode must be bf16, mxfp8, or both") + if max_prompt_length <= 0 or max_generate_length <= 0: + raise ValueError("max_prompt_length and max_generate_length must be positive") + prepare_assets.remote() + modes = ("bf16", "mxfp8") if mode == "both" else (mode,) + for selected in modes: + benchmark.remote(selected, steps, max_prompt_length, max_generate_length) diff --git a/examples/train/megatron/modal_vllm_mxfp8_concurrency_benchmark.py b/examples/train/megatron/modal_vllm_mxfp8_concurrency_benchmark.py new file mode 100644 index 0000000000..cd97122090 --- /dev/null +++ b/examples/train/megatron/modal_vllm_mxfp8_concurrency_benchmark.py @@ -0,0 +1,500 @@ +"""Benchmark BF16 and expert-only MXFP8 vLLM serving on DAPO prompts. + +Run: + uv run --isolated --with modal modal run --detach \ + examples/train/megatron/modal_vllm_mxfp8_concurrency_benchmark.py \ + --mode both +""" + +from __future__ import annotations + +import json +import os +import pathlib + +import modal + +APP_NAME = "skyrl-vllm-mxfp8-concurrency-benchmark" +MODEL = "Qwen/Qwen3-30B-A3B" +DAPO_DATASET = "BytedTsinghua-SIA/DAPO-Math-17k" +REMOTE_REPO = "/root/SkyRL" +HF_HOME = "/root/hf-cache" +DATA_ROOT = "/root/data/vllm-mxfp8-concurrency" +RESULT_ROOT = "/root/results/vllm-mxfp8-concurrency" +PORT = 8000 +GPU = os.environ.get("MODAL_GPU", "B200:8") + + +def _repo_root() -> pathlib.Path: + for start in (pathlib.Path(__file__).resolve(), pathlib.Path.cwd().resolve()): + candidate = start if start.is_dir() else start.parent + for path in (candidate, *candidate.parents): + if (path / "pyproject.toml").exists() and (path / "skyrl").exists(): + return path + raise RuntimeError("Run the benchmark from a SkyRL checkout") + + +repo_root = _repo_root() +hf_volume = modal.Volume.from_name("skyrl-hf-cache", create_if_missing=True) +data_volume = modal.Volume.from_name("skyrl-vllm-mxfp8-concurrency-data", create_if_missing=True) +result_volume = modal.Volume.from_name("skyrl-vllm-mxfp8-concurrency-results", create_if_missing=True) + +image = ( + modal.Image.from_registry("nvidia/cuda:12.8.1-devel-ubuntu22.04", add_python="3.12") + .apt_install("git", "curl", "build-essential", "ca-certificates", "libnuma1", "numactl") + .pip_install("huggingface-hub", "datasets", "jinja2>=3.1", "matplotlib", "transformers>=5.6.1,<=5.8.0") + .run_commands("curl -LsSf https://astral.sh/uv/install.sh | sh") + .env( + { + "PATH": "/root/.local/bin:/usr/local/cuda/bin:${PATH}", + "HF_HOME": HF_HOME, + "HF_XET_HIGH_PERFORMANCE": "1", + "MPLBACKEND": "Agg", + "UV_LINK_MODE": "copy", + "UV_PROJECT_ENVIRONMENT": f"{REMOTE_REPO}/.venv", + "VLLM_USE_FLASHINFER_MOE_FP16": "0", + "VLLM_USE_FLASHINFER_SAMPLER": "0", + } + ) + .add_local_dir( + str(repo_root), + REMOTE_REPO, + copy=True, + ignore=[".venv", ".git", "**/__pycache__"], + ) + .workdir(REMOTE_REPO) + .run_commands("uv sync --extra megatron", gpu="any") + .run_commands(f"rm -rf {HF_HOME}") +) + +app = modal.App(APP_NAME) + + +def _run_command(command: list[str]) -> None: + import subprocess + + subprocess.run(command, cwd=REMOTE_REPO, check=True) + + +@app.function( + image=image, + volumes={HF_HOME: hf_volume, DATA_ROOT: data_volume}, + timeout=2 * 60 * 60, +) +def prepare_assets(max_prompts: int, max_prompt_tokens: int) -> str: + from datasets import load_dataset + from huggingface_hub import snapshot_download + from transformers import AutoTokenizer + + snapshot_download(MODEL) + hf_volume.commit() + + os.makedirs(DATA_ROOT, exist_ok=True) + output_path = f"{DATA_ROOT}/dapo-prompts-{max_prompt_tokens}.jsonl" + tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True) + dataset = load_dataset(DAPO_DATASET, split="train", streaming=True).shuffle(seed=0, buffer_size=10_000) + + prompts: list[str] = [] + seen: set[str] = set() + for row in dataset: + messages = row["prompt"] + if hasattr(messages, "tolist"): + messages = messages.tolist() + if not isinstance(messages, list): + messages = [{"role": "user", "content": str(messages)}] + + key = json.dumps(messages, sort_keys=True, ensure_ascii=False) + if key in seen: + continue + seen.add(key) + + formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + prompt_tokens = tokenizer.encode(formatted, add_special_tokens=False) + if len(prompt_tokens) < 4 or len(prompt_tokens) > max_prompt_tokens: + continue + prompts.append(formatted) + if len(prompts) >= max_prompts: + break + + if len(prompts) < max_prompts: + raise RuntimeError(f"DAPO produced only {len(prompts)} prompts within {max_prompt_tokens} tokens") + + with open(output_path, "w", encoding="utf-8") as output_file: + for prompt in prompts: + output_file.write(json.dumps({"prompt": prompt}, ensure_ascii=False) + "\n") + data_volume.commit() + print(f"Prepared {len(prompts)} unique DAPO prompts at {output_path}") + return output_path + + +def _server_command(mode: str, max_model_len: int) -> list[str]: + command = [ + "uv", + "run", + "--frozen", + "--no-sync", + "--extra", + "megatron", + "vllm", + "serve", + MODEL, + "--host", + "0.0.0.0", + "--port", + str(PORT), + "--served-model-name", + MODEL, + "--dtype", + "bfloat16", + "--tensor-parallel-size", + "1", + "--data-parallel-size", + "8", + "--data-parallel-backend", + "mp", + "--enable-expert-parallel", + "--distributed-executor-backend", + "mp", + "--gpu-memory-utilization", + "0.6", + "--max-model-len", + str(max_model_len), + "--max-num-batched-tokens", + "8192", + "--max-num-seqs", + "1024", + "--enable-prefix-caching", + "--enable-chunked-prefill", + "--trust-remote-code", + ] + if mode == "mxfp8": + command.extend(["--quantization", "online", "--quantization-config", '{"moe":"mxfp8"}']) + return command + + +def _wait_for_server(process, log_path: str, timeout_s: int = 1800) -> None: + import time + import urllib.request + + deadline = time.monotonic() + timeout_s + health_url = f"http://127.0.0.1:{PORT}/health" + while time.monotonic() < deadline: + if process.poll() is not None: + log_text = pathlib.Path(log_path).read_text(encoding="utf-8", errors="replace") + raise RuntimeError(f"vLLM exited with code {process.returncode}\n{log_text[-20_000:]}") + try: + with urllib.request.urlopen(health_url, timeout=5) as response: + if response.status == 200: + return + except Exception: + time.sleep(5) + raise TimeoutError(f"vLLM did not become healthy within {timeout_s}s; log={log_path}") + + +def _bench_command( + dataset_path: str, + output_len: int, + concurrency: int, + num_prompts: int, + result_dir: str, + result_filename: str, +) -> list[str]: + return [ + "uv", + "run", + "--frozen", + "--no-sync", + "--extra", + "megatron", + "vllm", + "bench", + "serve", + "--backend", + "vllm", + "--base-url", + f"http://127.0.0.1:{PORT}", + "--endpoint", + "/v1/completions", + "--model", + MODEL, + "--dataset-name", + "custom", + "--dataset-path", + dataset_path, + "--custom-output-len", + str(output_len), + "--skip-chat-template", + "--disable-shuffle", + "--no-oversample", + "--num-prompts", + str(num_prompts), + "--request-rate", + "inf", + "--max-concurrency", + str(concurrency), + "--ignore-eos", + "--temperature", + "1.0", + "--top-p", + "1.0", + "--logprobs", + "1", + "--percentile-metrics", + "ttft,tpot,itl,e2el", + "--seed", + "0", + "--save-result", + "--result-dir", + result_dir, + "--result-filename", + result_filename, + ] + + +@app.function( + image=image, + gpu=GPU, + volumes={ + HF_HOME: hf_volume, + DATA_ROOT: data_volume, + RESULT_ROOT: result_volume, + }, + timeout=24 * 60 * 60, +) +def benchmark_mode( + mode: str, + dataset_path: str, + concurrencies: list[int], + output_len: int, + min_requests: int, + request_multiplier: int, + max_prompt_tokens: int, +) -> dict: + import signal + import subprocess + + if mode not in {"bf16", "mxfp8"}: + raise ValueError("mode must be bf16 or mxfp8") + + mode_result_dir = f"{RESULT_ROOT}/{mode}" + os.makedirs(mode_result_dir, exist_ok=True) + server_log_path = f"{mode_result_dir}/server.log" + max_model_len = max_prompt_tokens + output_len + + with open(server_log_path, "w", encoding="utf-8") as server_log: + process = subprocess.Popen( + _server_command(mode, max_model_len), + cwd=REMOTE_REPO, + stdout=server_log, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + + try: + _wait_for_server(process, server_log_path) + + warmup_concurrency = min(32, max(concurrencies)) + warmup_prompts = max(32, warmup_concurrency) + with open(dataset_path, encoding="utf-8") as source_file: + prompt_lines = source_file.readlines() + prompt_offset = 0 + warmup_dataset_path = f"{mode_result_dir}/warmup-prompts.jsonl" + with open(warmup_dataset_path, "w", encoding="utf-8") as warmup_file: + warmup_file.writelines(prompt_lines[prompt_offset : prompt_offset + warmup_prompts]) + prompt_offset += warmup_prompts + warmup_command = _bench_command( + warmup_dataset_path, + min(32, output_len), + warmup_concurrency, + warmup_prompts, + mode_result_dir, + "warmup.json", + ) + _run_command(warmup_command) + + points = [] + for concurrency in concurrencies: + num_prompts = max(min_requests, request_multiplier * concurrency) + if prompt_offset + num_prompts > len(prompt_lines): + raise RuntimeError(f"Need {prompt_offset + num_prompts} prompts, found {len(prompt_lines)}") + result_filename = f"{mode}-c{concurrency}.json" + point_dataset_path = f"{mode_result_dir}/prompts-c{concurrency}.jsonl" + with open(point_dataset_path, "w", encoding="utf-8") as point_dataset_file: + point_dataset_file.writelines(prompt_lines[prompt_offset : prompt_offset + num_prompts]) + prompt_offset += num_prompts + _run_command( + _bench_command( + point_dataset_path, + output_len, + concurrency, + num_prompts, + mode_result_dir, + result_filename, + ) + ) + with open(f"{mode_result_dir}/{result_filename}", encoding="utf-8") as result_file: + result = json.load(result_file) + if result["completed"] != num_prompts or result.get("failed", 0): + raise RuntimeError( + f"{mode} concurrency={concurrency}: completed={result['completed']} " + f"failed={result.get('failed', 0)} expected={num_prompts}" + ) + point = { + "concurrency": concurrency, + "num_prompts": num_prompts, + "output_throughput": result["output_throughput"], + "request_throughput": result["request_throughput"], + "mean_ttft_ms": result["mean_ttft_ms"], + "median_ttft_ms": result["median_ttft_ms"], + "p99_ttft_ms": result["p99_ttft_ms"], + "mean_tpot_ms": result["mean_tpot_ms"], + "median_tpot_ms": result["median_tpot_ms"], + "p99_tpot_ms": result["p99_tpot_ms"], + "mean_e2el_ms": result["mean_e2el_ms"], + "p99_e2el_ms": result["p99_e2el_ms"], + } + points.append(point) + print(json.dumps({"mode": mode, **point}, sort_keys=True)) + + summary = { + "mode": mode, + "model": MODEL, + "dataset": DAPO_DATASET, + "tensor_parallel_size": 1, + "expert_parallel_size": 8, + "data_parallel_size": 8, + "output_len": output_len, + "max_prompt_tokens": max_prompt_tokens, + "points": points, + } + with open(f"{mode_result_dir}/summary.json", "w", encoding="utf-8") as summary_file: + json.dump(summary, summary_file, indent=2) + result_volume.commit() + return summary + finally: + if process.poll() is None: + os.killpg(process.pid, signal.SIGTERM) + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait() + result_volume.commit() + + +@app.function( + image=image, + volumes={RESULT_ROOT: result_volume}, + timeout=30 * 60, +) +def render_results(results: list[dict]) -> dict: + import csv + + import matplotlib.pyplot as plt + + by_mode = {result["mode"]: result for result in results} + concurrencies = [point["concurrency"] for point in results[0]["points"]] + os.makedirs(RESULT_ROOT, exist_ok=True) + + rows = [] + for concurrency in concurrencies: + row = {"concurrency": concurrency} + for mode, result in by_mode.items(): + point = next(point for point in result["points"] if point["concurrency"] == concurrency) + row[f"{mode}_output_throughput"] = point["output_throughput"] + row[f"{mode}_mean_ttft_ms"] = point["mean_ttft_ms"] + row[f"{mode}_mean_tpot_ms"] = point["mean_tpot_ms"] + if "bf16" in by_mode and "mxfp8" in by_mode: + bf16 = row["bf16_output_throughput"] + row["gain_percent"] = 100 * (row["mxfp8_output_throughput"] / bf16 - 1) + rows.append(row) + + csv_path = f"{RESULT_ROOT}/comparison.csv" + with open(csv_path, "w", newline="", encoding="utf-8") as csv_file: + writer = csv.DictWriter(csv_file, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + fig, axis = plt.subplots(figsize=(9, 5)) + width = 0.38 + positions = list(range(len(concurrencies))) + modes = [mode for mode in ("bf16", "mxfp8") if mode in by_mode] + for mode_index, mode in enumerate(modes): + offsets = [position + (mode_index - (len(modes) - 1) / 2) * width for position in positions] + throughputs = [point["output_throughput"] for point in by_mode[mode]["points"]] + axis.bar(offsets, throughputs, width=width, label=mode.upper()) + axis.set_title("Qwen3-30B-A3B vLLM rollout throughput") + axis.set_xlabel("Maximum concurrent requests") + axis.set_ylabel("Output tokens per second") + axis.set_xticks(positions, [str(value) for value in concurrencies]) + axis.legend() + axis.grid(axis="y", alpha=0.25) + fig.tight_layout() + plot_path = f"{RESULT_ROOT}/comparison.png" + fig.savefig(plot_path, dpi=180) + plt.close(fig) + + combined_path = f"{RESULT_ROOT}/comparison.json" + with open(combined_path, "w", encoding="utf-8") as combined_file: + json.dump({"results": results, "comparison": rows}, combined_file, indent=2) + result_volume.commit() + return {"rows": rows, "csv": csv_path, "plot": plot_path, "json": combined_path} + + +def _print_comparison(rendered: dict) -> None: + rows = rendered["rows"] + has_both = "gain_percent" in rows[0] + if has_both: + print("| Concurrent requests | BF16 output tok/s | MXFP8 output tok/s | Gain |") + print("|---:|---:|---:|---:|") + for row in rows: + print( + f"| {row['concurrency']} | {row['bf16_output_throughput']:.1f} | " + f"{row['mxfp8_output_throughput']:.1f} | {row['gain_percent']:+.1f}% |" + ) + else: + mode = "bf16" if "bf16_output_throughput" in rows[0] else "mxfp8" + print(f"| Concurrent requests | {mode.upper()} output tok/s |") + print("|---:|---:|") + for row in rows: + print(f"| {row['concurrency']} | {row[f'{mode}_output_throughput']:.1f} |") + print(json.dumps({key: value for key, value in rendered.items() if key != "rows"}, indent=2)) + + +@app.local_entrypoint() +def main( + mode: str = "both", + concurrencies: str = "1,16,32,128,256", + output_len: int = 512, + max_prompt_tokens: int = 512, + min_requests: int = 64, + request_multiplier: int = 2, +) -> None: + if mode not in {"bf16", "mxfp8", "both"}: + raise ValueError("mode must be bf16, mxfp8, or both") + concurrency_values = [int(value) for value in concurrencies.split(",") if value] + if not concurrency_values or any(value <= 0 for value in concurrency_values): + raise ValueError("concurrencies must contain positive integers") + if output_len <= 0 or max_prompt_tokens <= 0: + raise ValueError("output_len and max_prompt_tokens must be positive") + if min_requests <= 0 or request_multiplier <= 0: + raise ValueError("min_requests and request_multiplier must be positive") + + warmup_prompts = 32 + benchmark_prompts = sum(max(min_requests, request_multiplier * value) for value in concurrency_values) + dataset_path = prepare_assets.remote(warmup_prompts + benchmark_prompts, max_prompt_tokens) + modes = ("bf16", "mxfp8") if mode == "both" else (mode,) + results = [ + benchmark_mode.remote( + selected_mode, + dataset_path, + concurrency_values, + output_len, + min_requests, + request_multiplier, + max_prompt_tokens, + ) + for selected_mode in modes + ] + rendered = render_results.remote(results) + _print_comparison(rendered) diff --git a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py index 7a3f94a9cc..9cee5b8841 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py @@ -413,12 +413,11 @@ def preprocess_packed_seqs( per row. This is the historical SkyRL behavior used by the RL path and the existing SFT path without mini-batch packing. - ``sub_seq_lengths is not None``: each row may contain multiple - sub-sequences concatenated end-to-end. ``sub_seq_lengths[r]`` lists - the per-sub-sequence valid token counts for row ``r``. Tokens - ``input_ids[r, :sum(sub_seq_lengths[r])]`` are assumed to be the - concatenated sub-sequences in order; any trailing tokens in the row - are pad. ``cu_seqlens`` enumerates every sub-sequence across every - row. + sub-sequences. ``sub_seq_lengths[r]`` lists their valid token counts. + Each sub-sequence begins at the next ``align_size`` boundary, so internal + alignment padding may separate adjacent sub-sequences; any remaining + trailing tokens are pad. ``cu_seqlens`` enumerates every sub-sequence + across every row. CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0 gets first and last chunks, GPU1 gets second and second last chunks, diff --git a/skyrl/backends/skyrl_train/distributed/megatron/packing_utils.py b/skyrl/backends/skyrl_train/distributed/megatron/packing_utils.py index f6b506361b..db6b80b311 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/packing_utils.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/packing_utils.py @@ -5,23 +5,29 @@ def is_fp8_enabled(fp8: Any) -> bool: """Return whether a Megatron/TE fp8 config value enables FP8 execution.""" if isinstance(fp8, str): - return fp8.strip().lower() not in {"", "0", "false", "none", "null", "no"} + return fp8.strip().lower() not in {"", "0", "false", "none", "null", "no", "off"} return bool(fp8) def get_packed_seq_align_size(tp_size: int, cp_size: int, fp8_enabled: bool = False) -> int: - """Return global per-subsequence padding needed for TP/CP layout.""" + """Return the global alignment unit for packed TP/CP/FP8 sequences.""" + if tp_size < 1 or cp_size < 1: + raise ValueError(f"tp_size and cp_size must be positive, got tp_size={tp_size}, cp_size={cp_size}") if cp_size > 1: layout_align = tp_size * cp_size * 2 else: layout_align = tp_size if not fp8_enabled: return layout_align - return math.lcm(layout_align, 16 * cp_size) + fp8_token_align = 128 * tp_size * cp_size if tp_size > 1 else 16 * cp_size + return math.lcm(layout_align, fp8_token_align) def get_unpacked_seq_align_size(tp_size: int, fp8_enabled: bool = False) -> int: - """Return sequence padding needed when removing microbatch padding without CP.""" + """Return the alignment unit for unpacked TP/FP8 sequences without CP.""" + if tp_size < 1: + raise ValueError(f"tp_size must be positive, got {tp_size}") if not fp8_enabled: return tp_size - return math.lcm(tp_size, 16) + fp8_token_align = 128 * tp_size if tp_size > 1 else 16 + return math.lcm(tp_size, fp8_token_align) diff --git a/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py b/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py index bf1b4c2a5f..063c2c468c 100644 --- a/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py +++ b/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py @@ -24,15 +24,139 @@ skyrl.backends.skyrl_train.inference_servers.new_inference_worker_wrap.NewInferenceWorkerWrap """ +from typing import Any + import torch from skyrl.backends.skyrl_train.inference_servers.layerwise_reload import ( LayerwiseReloadWorkerMixin, ) +from skyrl.backends.skyrl_train.weight_sync.base import cuda_uuid_to_str +from skyrl.backends.skyrl_train.weight_sync.serialized_fp8 import ( + SKYRL_BATCHED_MOE_FP8_PREFIX, +) VLLM_NEW_INFERENCE_WORKER_EXTENSION_CLS = f"{__name__}.NewInferenceWorkerWrap" +_BATCHED_MOE_TARGETS = { + ".experts.gate_proj.weight": (".experts.w13_weight", "w1"), + ".experts.up_proj.weight": (".experts.w13_weight", "w3"), + ".experts.down_proj.weight": (".experts.w2_weight", "w2"), + ".experts.gate_proj.weight_scale_inv": (".experts.w13_weight_scale_inv", "w1"), + ".experts.up_proj.weight_scale_inv": (".experts.w13_weight_scale_inv", "w3"), + ".experts.down_proj.weight_scale_inv": (".experts.w2_weight_scale_inv", "w2"), + ".experts.gate_proj.weight_scale": (".experts.w13_weight_scale", "w1"), + ".experts.up_proj.weight_scale": (".experts.w13_weight_scale", "w3"), + ".experts.down_proj.weight_scale": (".experts.w2_weight_scale", "w2"), +} + + +def _map_hf_weight_name(model: torch.nn.Module, name: str) -> str: + """Apply a top-level vLLM model's HF-to-runtime prefix mapping.""" + mapper = getattr(model, "hf_to_vllm_mapper", None) + if mapper is None: + return name + mapped = mapper.apply_list([name]) + if len(mapped) != 1: + raise ValueError(f"Unable to map batched MoE checkpoint name {name!r}") + return mapped[0] + + +def _load_batched_moe_fp8_tensor( + model: torch.nn.Module, + params_dict: dict[str, torch.nn.Parameter], + wire_name: str, + loaded_weight: torch.Tensor, +) -> bool: + """Load one expert-batched FP8 weight/scale through FusedMoE's loader. + + Returns ``False`` for ordinary checkpoint tensors. Marked tensors are + required to resolve successfully so a protocol mismatch cannot silently + leave stale rollout weights behind. + """ + if not wire_name.startswith(SKYRL_BATCHED_MOE_FP8_PREFIX): + return False + if loaded_weight.ndim != 3: + raise ValueError( + f"Batched MoE wire tensor must be 3D, got name={wire_name!r}, shape={tuple(loaded_weight.shape)}" + ) + + checkpoint_name = wire_name.removeprefix(SKYRL_BATCHED_MOE_FP8_PREFIX) + mapped_name = _map_hf_weight_name(model, checkpoint_name) + target_name = None + shard_id = None + for checkpoint_suffix, ( + target_suffix, + candidate_shard_id, + ) in _BATCHED_MOE_TARGETS.items(): + if mapped_name.endswith(checkpoint_suffix): + target_name = mapped_name[: -len(checkpoint_suffix)] + target_suffix + shard_id = candidate_shard_id + break + if target_name is None or shard_id is None: + raise ValueError(f"Unsupported batched MoE wire tensor name {wire_name!r}") + if target_name not in params_dict: + raise ValueError(f"Batched MoE target parameter {target_name!r} was not found for wire tensor {wire_name!r}") + + param = params_dict[target_name] + weight_loader = getattr(param, "weight_loader", None) + if weight_loader is None or not getattr(weight_loader, "supports_moe_loading", False): + # Layerwise reload wraps the loader with functools.wraps, which copies + # this marker from FusedMoE.weight_loader onto the wrapper. + raise ValueError(f"Parameter {target_name!r} does not expose a FusedMoE weight loader") + + if param.shape[0] == loaded_weight.shape[0]: + success = weight_loader( + param, + loaded_weight, + target_name, + shard_id=shard_id, + expert_id=0, + return_success=True, + ) + if not success: + raise ValueError(f"Fused loading failed for batched MoE tensor {wire_name!r}") + return True + + # Expert-parallel vLLM keeps only a subset locally. Retain the compact wire + # format, but let the loader map global expert IDs one view at a time. + loaded_any = False + for expert_id, expert_weight in enumerate(loaded_weight.unbind(0)): + loaded_any = ( + bool( + weight_loader( + param, + expert_weight, + target_name, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + ) + or loaded_any + ) + if not loaded_any: + raise ValueError(f"No local expert accepted batched MoE tensor {wire_name!r}") + return True + + +def _load_checkpoint_weights(model: torch.nn.Module, weights: list[tuple[str, torch.Tensor]]) -> Any: + """Load ordinary HF tensors plus SkyRL's compact batched-MoE tensors.""" + params_dict: dict[str, torch.nn.Parameter] | None = None + ordinary_weights: list[tuple[str, torch.Tensor]] = [] + for name, weight in weights: + if name.startswith(SKYRL_BATCHED_MOE_FP8_PREFIX): + if params_dict is None: + params_dict = dict(model.named_parameters()) + _load_batched_moe_fp8_tensor(model, params_dict, name, weight) + else: + ordinary_weights.append((name, weight)) + if ordinary_weights: + return model.load_weights(weights=ordinary_weights) + return set() + + class NewInferenceWorkerWrap(LayerwiseReloadWorkerMixin): """ vLLM worker extension for chunked weight sync (new inference path). @@ -72,7 +196,7 @@ def update_weights_ipc(self, update_info: dict) -> None: if self.weight_transfer_engine is None: raise RuntimeError( - "Weight transfer not configured. " "Please set weight_transfer_config to enable weight transfer." + "Weight transfer not configured. Please set weight_transfer_config to enable weight transfer." ) # --- unpack SkyRL packed CUDA IPC format --- @@ -86,9 +210,9 @@ def update_weights_ipc(self, update_info: dict) -> None: handles = pickle.loads(base64.b64decode(pickled)) device_index = torch.cuda.current_device() - physical_gpu_id = str(torch.cuda.get_device_properties(device_index).uuid) + physical_gpu_id = cuda_uuid_to_str(torch.cuda.get_device_properties(device_index).uuid) if physical_gpu_id not in handles: - raise ValueError(f"IPC handle not found for GPU UUID {physical_gpu_id}. " f"Available: {list(handles)}") + raise ValueError(f"IPC handle not found for GPU UUID {physical_gpu_id}. Available: {list(handles)}") func, args = handles[physical_gpu_id] # Remap device index to the LOCAL current-device. list_args = list(args) @@ -109,7 +233,7 @@ def update_weights_ipc(self, update_info: dict) -> None: model = self.model_runner.model with set_current_vllm_config(self.vllm_config), torch.device(self.device): if self._skyrl_is_checkpoint_format: - model.load_weights(weights=weights) + _load_checkpoint_weights(model, weights) # vLLM's load only updates the main model; the spec-decode (MTP/Eagle) # drafter is a separate module and must be reloaded from the same # checkpoint-format weights (see spec_decode_utils). @@ -165,7 +289,7 @@ def update_weights_nccl(self, update_info: dict) -> None: def _load_weights(weights): weights = list(weights) - loaded = model.load_weights(weights=weights) + loaded = _load_checkpoint_weights(model, weights) _reload_spec_decode_drafter(self.model_runner, weights) return loaded diff --git a/skyrl/backends/skyrl_train/inference_servers/utils.py b/skyrl/backends/skyrl_train/inference_servers/utils.py index c9de06c4d1..d21bd921ad 100644 --- a/skyrl/backends/skyrl_train/inference_servers/utils.py +++ b/skyrl/backends/skyrl_train/inference_servers/utils.py @@ -11,6 +11,16 @@ SKYRL_LORA_ADAPTER_NAME, ) from skyrl.backends.skyrl_train.weight_sync import get_transfer_strategy +from skyrl.backends.skyrl_train.weight_sync.serialized_fp8 import ( + SERIALIZED_BLOCKWISE_FP8, + SERIALIZED_MXFP8, + get_hf_model_type, + get_moe_architecture_spec, + get_qwen35_fp8_ignored_layers, + get_serialized_fp8_quantization_config, + get_serialized_mxfp8_quantization_config, + is_qwen35_config, +) from skyrl.train.config import ( InferenceEngineConfig, SkyRLTrainConfig, @@ -20,6 +30,130 @@ logger = logging.getLogger(__name__) +def _serialized_fp8_ignored_layers(model_path: Optional[str]) -> list[str]: + if not model_path: + raise ValueError("A model path is required when serialized FP8 weight sync is enabled") + try: + from transformers import AutoConfig + + hf_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) + except Exception as exc: + raise RuntimeError( + "Could not inspect the model config required to derive serialized FP8 ignored layers: " + f"model_path={model_path!r}" + ) from exc + if not is_qwen35_config(hf_config): + raise ValueError( + f"Serialized FP8 weight sync currently supports only Qwen3.5 checkpoint layouts; model_path={model_path!r}" + ) + return get_qwen35_fp8_ignored_layers(hf_config) + + +def _serialized_mxfp8_quantization_config(model_path: Optional[str]) -> dict: + if not model_path: + raise ValueError("A model path is required when serialized MXFP8 weight sync is enabled") + try: + from transformers import AutoConfig + + hf_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) + except Exception as exc: + raise RuntimeError( + f"Could not inspect the model config required for serialized MXFP8 weight sync: model_path={model_path!r}" + ) from exc + get_moe_architecture_spec(get_hf_model_type(hf_config)) + return get_serialized_mxfp8_quantization_config() + + +def _set_or_validate(mapping: Dict[str, Any], key: str, expected: Any, *, context: str) -> None: + if key in mapping and mapping[key] != expected: + raise ValueError( + f"{context}.{key} must be {expected!r} when serialized FP8 weight sync is enabled, got {mapping[key]!r}" + ) + mapping[key] = copy.deepcopy(expected) + + +def _apply_serialized_fp8_weight_sync_defaults( + ie_cfg: InferenceEngineConfig, + engine_kwargs: Dict[str, Any], + model_path: Optional[str] = None, +) -> None: + """Configure vLLM for checkpoint-format FP8 weight reloads.""" + + mode = ie_cfg.fp8_weight_sync_mode + if mode is None: + return + if mode == SERIALIZED_BLOCKWISE_FP8: + quantization = "fp8" + elif mode == SERIALIZED_MXFP8: + if ie_cfg.model_dtype != "bfloat16": + raise ValueError("serialized_mxfp8 weight sync requires model_dtype='bfloat16'") + quantization = "modelopt_mxfp8" + else: + raise ValueError( + f"Unsupported fp8_weight_sync_mode={mode!r}. " + f"Supported values: {SERIALIZED_BLOCKWISE_FP8!r}, {SERIALIZED_MXFP8!r}." + ) + + _set_or_validate(engine_kwargs, "quantization", quantization, context="engine_init_kwargs") + # Build FP8 modules without a bootstrap checkpoint; the first full-weight + # sync replaces the dummy values. + _set_or_validate(engine_kwargs, "load_format", "dummy", context="engine_init_kwargs") + + hf_overrides_value = engine_kwargs.get("hf_overrides") + hf_overrides = {} if hf_overrides_value is None else copy.deepcopy(hf_overrides_value) + if not isinstance(hf_overrides, dict): + raise ValueError("engine_init_kwargs.hf_overrides must be a dict when serialized FP8 weight sync is enabled") + + qcfg_value = hf_overrides.get("quantization_config") + qcfg = {} if qcfg_value is None else copy.deepcopy(qcfg_value) + if not isinstance(qcfg, dict): + raise ValueError( + "engine_init_kwargs.hf_overrides.quantization_config must be a dict " + "when serialized FP8 weight sync is enabled" + ) + + if mode == SERIALIZED_BLOCKWISE_FP8: + ignored_layers = _serialized_fp8_ignored_layers(model_path) + if ignored_layers: + logger.info( + "Serialized FP8 weight sync will leave %d vLLM modules unquantized " + "to match the Qwen3.5 serialized FP8 policy.", + len(ignored_layers), + ) + serialized_quantization_config = get_serialized_fp8_quantization_config( + ignored_layers=ignored_layers, + ) + else: + serialized_quantization_config = _serialized_mxfp8_quantization_config(model_path) + + for key, value in serialized_quantization_config.items(): + _set_or_validate( + qcfg, + key, + value, + context="engine_init_kwargs.hf_overrides.quantization_config", + ) + hf_overrides["quantization_config"] = qcfg + engine_kwargs["hf_overrides"] = hf_overrides + + +def apply_expert_mxfp8_rollout_config(args: Namespace, cfg: SkyRLTrainConfig, engine_kwargs: dict) -> None: + """Configure expert-only online MXFP8 in vLLM.""" + + expert_mxfp8 = cfg.trainer.policy.model.expert_mxfp8 + if not expert_mxfp8.enabled or not expert_mxfp8.rollout: + return + if cfg.generator.inference_engine.fp8_weight_sync_mode == SERIALIZED_MXFP8: + return + if cfg.generator.inference_engine.model_dtype not in ("bfloat16", "float16"): + raise ValueError("Expert MXFP8 rollout requires model_dtype=bfloat16 or float16") + conflicts = {"quantization", "quantization_config"} & engine_kwargs.keys() + if conflicts: + raise ValueError(f"Expert MXFP8 conflicts with engine_init_kwargs: {sorted(conflicts)}") + args.quantization = "online" + args.quantization_config = {"moe": "mxfp8"} + + def _uses_lora_weight_sync(cfg: SkyRLTrainConfig) -> bool: """Return True when the trainer syncs LoRA adapters (not merged weights). @@ -153,6 +287,12 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace: logger.info(f"vLLM speculative decoding enabled: speculative_config={spec_cfg}") engine_kwargs = get_config_as_dict(ie_cfg.engine_init_kwargs) + _apply_serialized_fp8_weight_sync_defaults( + ie_cfg, + engine_kwargs, + cfg.trainer.policy.model.path, + ) + apply_expert_mxfp8_rollout_config(args, cfg, engine_kwargs) for key, value in engine_kwargs.items(): setattr(args, key, value) diff --git a/skyrl/backends/skyrl_train/weight_sync/base.py b/skyrl/backends/skyrl_train/weight_sync/base.py index fd8753c24b..158d4fa12b 100644 --- a/skyrl/backends/skyrl_train/weight_sync/base.py +++ b/skyrl/backends/skyrl_train/weight_sync/base.py @@ -2,7 +2,7 @@ from dataclasses import asdict, dataclass, field from functools import cached_property -from typing import Any, Dict, List +from typing import Any, Dict, Iterator, List import torch @@ -84,7 +84,12 @@ class WeightChunk: def __post_init__(self): """Validate that all input lists have the same length.""" - lengths = [len(self.names), len(self.dtypes), len(self.shapes), len(self.tensors)] + lengths = [ + len(self.names), + len(self.dtypes), + len(self.shapes), + len(self.tensors), + ] if len(set(lengths)) != 1: raise ValueError( f"All lists must have the same length. Got names={len(self.names)}, " @@ -104,3 +109,53 @@ def total_numel(self) -> int: def total_size_bytes(self) -> int: """Calculate total memory footprint in bytes.""" return sum(t.numel() * t.element_size() for t in self.tensors) + + +def torch_dtype_name(dtype: torch.dtype) -> str: + """Return the dtype spelling expected by vLLM weight-transfer metadata.""" + return str(dtype).split(".")[-1] + + +def cuda_uuid_to_str(uuid: str | bytes) -> str: + """Normalize CUDA UUIDs identically on both sides of an IPC transfer.""" + return uuid.decode("ascii") if isinstance(uuid, bytes) else str(uuid) + + +def iter_single_dtype_chunks(chunk: WeightChunk) -> Iterator[WeightChunk]: + """Yield dtype-homogeneous subchunks in first-seen dtype order. + + CUDA IPC packs tensors into a typed buffer, while serialized FP8 chunks mix + FP8 weights, scale tensors, and unquantized BF16 tensors. vLLM's NCCL path + byte-packs mixed dtypes and does not require this split. + """ + by_dtype: Dict[torch.dtype, Dict[str, list]] = {} + dtype_order: List[torch.dtype] = [] + + for name, tensor in zip(chunk.names, chunk.tensors): + dtype = tensor.dtype + if dtype not in by_dtype: + dtype_order.append(dtype) + by_dtype[dtype] = {"names": [], "dtypes": [], "shapes": [], "tensors": []} + group = by_dtype[dtype] + group["names"].append(name) + group["dtypes"].append(str(dtype)) + group["shapes"].append(list(tensor.shape)) + group["tensors"].append(tensor) + + for dtype in dtype_order: + group = by_dtype[dtype] + yield WeightChunk( + names=group["names"], + dtypes=group["dtypes"], + shapes=group["shapes"], + tensors=group["tensors"], + ) + + +def get_weight_chunk_metadata(chunk: WeightChunk) -> Dict[str, List]: + """Return vLLM metadata for the tensors in a transfer chunk.""" + return { + "names": list(chunk.names), + "dtype_names": [torch_dtype_name(tensor.dtype) for tensor in chunk.tensors], + "shapes": [list(tensor.shape) for tensor in chunk.tensors], + } diff --git a/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py b/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py index 6c192c67d8..6d5c5f6b63 100644 --- a/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py +++ b/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py @@ -18,7 +18,11 @@ import ray import torch -from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk, WeightUpdateRequest +from skyrl.backends.skyrl_train.weight_sync.base import ( + WeightChunk, + WeightUpdateRequest, + get_weight_chunk_metadata, +) from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import ( WeightSyncInitInfo, WeightTransferSender, @@ -117,32 +121,35 @@ async def send_chunks( self, chunks: Iterable[WeightChunk], weight_metadata: Optional[Dict[str, list]] = None, + derive_metadata_from_chunks: bool = False, ) -> None: """Send chunks via broadcast or vLLM native NCCL. Args: chunks: Iterable of WeightChunk objects to send. - weight_metadata: Pre-computed metadata dict with "names", "dtype_names", - "shapes". Avoids materializing all chunks to collect metadata. + weight_metadata: Complete metadata for the batched update path. + derive_metadata_from_chunks: Send each chunk with derived metadata. """ - await self._send_chunks_vllm_native(chunks, weight_metadata) + if derive_metadata_from_chunks: + if weight_metadata is not None: + raise ValueError("weight_metadata must be omitted when deriving metadata from chunks") + await self._send_serialized_fp8_chunks_vllm_native(chunks) + else: + await self._send_chunks_vllm_native(chunks, weight_metadata) async def _send_chunks_vllm_native( self, chunks: Iterable[WeightChunk], - weight_metadata: Optional[Dict[str, list]] = None, + weight_metadata: Optional[Dict[str, list]], ) -> None: - """Batched path: one update_weights call + trainer_send_weights (vLLM native). + """Send one metadata-declared update through vLLM's native NCCL path. All ranks must evaluate the chunks iterator (extract_weights uses collective all-gather internally). Only rank 0 sends the gathered tensors to vLLM via the NCCL weight transfer engine. """ if weight_metadata is None: - raise ValueError( - "weight_metadata is required for vLLM native path. " - "Call weight_extractor.get_weight_metadata() and pass it to send_chunks." - ) + raise ValueError("weight_metadata is required unless derive_metadata_from_chunks=true") def weight_iterator() -> Iterator[Tuple[str, torch.Tensor]]: for chunk in chunks: @@ -165,7 +172,7 @@ def weight_iterator() -> Iterator[Tuple[str, torch.Tensor]]: update_info = {**weight_metadata, "packed": True} update_task = asyncio.create_task(self._inference_client.update_weights_nccl(update_info)) - # Run in thread so the HTTP update_task can progress concurrently + # Run in a thread so the HTTP update task can progress concurrently. await asyncio.to_thread( NCCLWeightTransferEngine.trainer_send_weights, iterator=weight_iterator(), @@ -175,12 +182,46 @@ def weight_iterator() -> Iterator[Tuple[str, torch.Tensor]]: await self._inference_client.finish_weight_update() else: - # Non-rank-0 still needs to participate in the all-gather + # Non-rank-0 still needs to participate in extractor collectives. for _ in weight_iterator(): pass torch.distributed.barrier() + async def _send_serialized_fp8_chunks_vllm_native( + self, + chunks: Iterable[WeightChunk], + ) -> None: + """Send lazy mixed-dtype serialized-FP8 chunks through vLLM NCCL.""" + if torch.distributed.get_rank() == 0: + await self._inference_client.start_weight_update(is_checkpoint_format=True) + + for chunk in chunks: + if torch.distributed.get_rank() == 0: + await self._send_chunk_vllm_native(chunk) + + if torch.distributed.get_rank() == 0: + await self._inference_client.finish_weight_update() + + torch.distributed.barrier() + + async def _send_chunk_vllm_native(self, chunk: WeightChunk) -> None: + """Send one logical chunk through vLLM's byte-packed NCCL path.""" + from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLWeightTransferEngine, + ) + + update_info = {**get_weight_chunk_metadata(chunk), "packed": True} + update_task = asyncio.create_task(self._inference_client.update_weights_nccl(update_info)) + + # Let the receiver enter its collective while the trainer broadcasts. + await asyncio.to_thread( + NCCLWeightTransferEngine.trainer_send_weights, + iterator=iter(zip(chunk.names, chunk.tensors)), + trainer_args={"group": self._model_update_group, "packed": True}, + ) + await update_task + def teardown(self) -> None: """Destroy the process group used for weight transfer.""" if self._model_update_group is not None and isinstance( diff --git a/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py b/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py index 61b4417135..02599a5ca3 100644 --- a/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py +++ b/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py @@ -28,13 +28,18 @@ import torch from torch.multiprocessing.reductions import reduce_tensor -from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk, WeightUpdateRequest +from skyrl.backends.skyrl_train.weight_sync.base import ( + WeightChunk, + WeightUpdateRequest, + cuda_uuid_to_str, + iter_single_dtype_chunks, + torch_dtype_name, +) from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import ( WeightSyncInitInfo, WeightTransferSender, WeightTransferStrategy, ) -from skyrl.train.utils.utils import str_to_torch_dtype # IPC handle type: (rebuild_func, args) returned by reduce_tensor IpcHandle = Tuple[Callable[..., torch.Tensor], Tuple[Any, ...]] @@ -148,14 +153,14 @@ async def send_chunks( self, chunks: Iterable[WeightChunk], weight_metadata: Optional[Dict[str, list]] = None, + derive_metadata_from_chunks: bool = False, ) -> None: """Send chunks via CUDA IPC. Args: chunks: Iterable of WeightChunk objects to send. - weight_metadata: Unused for IPC (metadata is derived from chunks - directly to avoid ordering mismatches). Kept for interface - compatibility with the base class. + weight_metadata: Unused; IPC derives metadata from each tensor. + derive_metadata_from_chunks: Accepted for sender interface compatibility. """ await self._send_chunks_vllm_native(chunks, weight_metadata) @@ -185,77 +190,91 @@ async def _send_chunks_vllm_native( rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() device = torch.cuda.current_device() - gpu_uuid = str(torch.cuda.get_device_properties(device).uuid) - dtype = str_to_torch_dtype(self._init_info.model_dtype_str) - dtype_name = self._init_info.model_dtype_str.split(".")[-1] - + gpu_uuid = cuda_uuid_to_str(torch.cuda.get_device_properties(device).uuid) if rank == 0: await self._inference_client.start_weight_update(is_checkpoint_format=True) torch.distributed.barrier() - for chunk in chunks: - # --- pack all tensors in this chunk into one contiguous buffer --- - # Chunk tensors share a single dtype by construction (see - # weight_extractor_utils.py), so offsets in element units are safe. - names: List[str] = [] - dtype_names: List[str] = [] - shapes: List[List[int]] = [] - sizes: List[int] = [] - - total_numel = sum(t.numel() for t in chunk.tensors) - packed_tensor = torch.empty( - total_numel, - device=device, - dtype=dtype, - requires_grad=False, - ) - - offset = 0 - for name, tensor, shape in zip(chunk.names, chunk.tensors, chunk.shapes): - size = tensor.numel() - packed_tensor[offset : offset + size].copy_(tensor.detach().reshape(-1)) - offset += size - names.append(name) - dtype_names.append(dtype_name) - shapes.append(list(shape) if not isinstance(shape, list) else shape) - sizes.append(size) - - # --- one IPC handle per rank for the packed buffer --- - ipc_handle: IpcHandle = reduce_tensor(packed_tensor) - local_handle_dict: Dict[str, IpcHandle] = {gpu_uuid: ipc_handle} - gathered: List[Optional[Dict[str, IpcHandle]]] = [None] * world_size - torch.distributed.all_gather_object(gathered, local_handle_dict) - - torch.distributed.barrier() - torch.cuda.synchronize() - - if rank == 0: - merged_handles: Dict[str, IpcHandle] = {} - for d in gathered: - if d is not None: - merged_handles.update(d) - - pickled = base64.b64encode(pickle.dumps(merged_handles)).decode("utf-8") - chunk_update_info: Dict[str, Any] = { - "names": names, - "dtype_names": dtype_names, - "shapes": shapes, - "sizes": sizes, - "ipc_handles_pickled": pickled, - } - await self._inference_client.update_weights_ipc(chunk_update_info) - - # Keep packed_tensor alive past the barrier so the receiver's - # rebuilt view has valid backing storage while it copies into - # the model. Post-barrier drops the local ref safely. - torch.distributed.barrier() - torch.cuda.ipc_collect() - torch.cuda.synchronize() + for logical_chunk in chunks: + for chunk in iter_single_dtype_chunks(logical_chunk): + await self._send_single_dtype_chunk_vllm_native( + chunk=chunk, + device=device, + gpu_uuid=gpu_uuid, + world_size=world_size, + rank=rank, + ) if rank == 0: await self._inference_client.finish_weight_update() torch.distributed.barrier() + async def _send_single_dtype_chunk_vllm_native( + self, + *, + chunk: WeightChunk, + device: int, + gpu_uuid: str, + world_size: int, + rank: int, + ) -> None: + dtype = chunk.tensors[0].dtype + dtype_name = torch_dtype_name(dtype) + if any(tensor.dtype != dtype for tensor in chunk.tensors): + raise ValueError("CUDA IPC packed chunks must contain a single tensor dtype") + + names: List[str] = [] + dtype_names: List[str] = [] + shapes: List[List[int]] = [] + sizes: List[int] = [] + + total_numel = sum(t.numel() for t in chunk.tensors) + packed_tensor = torch.empty( + total_numel, + device=device, + dtype=dtype, + requires_grad=False, + ) + + offset = 0 + for name, tensor in zip(chunk.names, chunk.tensors): + size = tensor.numel() + packed_tensor[offset : offset + size].copy_(tensor.detach().reshape(-1)) + offset += size + names.append(name) + dtype_names.append(dtype_name) + shapes.append(list(tensor.shape)) + sizes.append(size) + + ipc_handle: IpcHandle = reduce_tensor(packed_tensor) + local_handle_dict: Dict[str, IpcHandle] = {gpu_uuid: ipc_handle} + gathered: List[Optional[Dict[str, IpcHandle]]] = [None] * world_size + torch.distributed.all_gather_object(gathered, local_handle_dict) + + torch.distributed.barrier() + torch.cuda.synchronize() + + if rank == 0: + merged_handles: Dict[str, IpcHandle] = {} + for d in gathered: + if d is not None: + merged_handles.update(d) + + pickled = base64.b64encode(pickle.dumps(merged_handles)).decode("utf-8") + chunk_update_info: Dict[str, Any] = { + "names": names, + "dtype_names": dtype_names, + "shapes": shapes, + "sizes": sizes, + "ipc_handles_pickled": pickled, + } + await self._inference_client.update_weights_ipc(chunk_update_info) + + # Keep the backing tensor alive until the receiver copies from its IPC view. + torch.distributed.barrier() + torch.cuda.ipc_collect() + torch.cuda.synchronize() + def teardown(self) -> None: """No-op for CUDA IPC sender (no custom process group to clean up).""" pass diff --git a/skyrl/backends/skyrl_train/weight_sync/megatron_bridge_mappings.py b/skyrl/backends/skyrl_train/weight_sync/megatron_bridge_mappings.py new file mode 100644 index 0000000000..8514f20358 --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/megatron_bridge_mappings.py @@ -0,0 +1,46 @@ +"""Megatron-Bridge mapping overrides used by serialized weight sync.""" + +from types import MethodType + + +def get_packed_qwen3_moe_conversion_tasks(auto_bridge, model): + """Build Qwen3 MoE export tasks with experts packed by layer.""" + + from megatron.bridge.models.conversion.mapping_registry import ( + MegatronMappingRegistry, + ) + from megatron.bridge.models.conversion.param_mapping import ( + FusedExpertMapping, + FusedGatedExpertMapping, + ) + + model_bridge = auto_bridge._model_bridge + mappings = [ + mapping + for mapping in model_bridge.mapping_registry() + if ".mlp.experts." not in mapping.megatron_param + ] + mappings.extend( + [ + FusedGatedExpertMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", + hf_param="model.layers.*.mlp.experts.gate_up_proj", + ), + FusedExpertMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*", + hf_param="model.layers.*.mlp.experts.down_proj", + ), + FusedGatedExpertMapping( + megatron_param="decoder.layers.*.mlp.experts.local_experts.*.linear_fc1.weight", + hf_param="model.layers.*.mlp.experts.gate_up_proj", + ), + FusedExpertMapping( + megatron_param="decoder.layers.*.mlp.experts.local_experts.*.linear_fc2.weight", + hf_param="model.layers.*.mlp.experts.down_proj", + ), + ] + ) + registry = MegatronMappingRegistry(*mappings) + model_bridge.mapping_registry = MethodType(lambda _self: registry, model_bridge) + model_list = model if isinstance(model, list) else [model] + return model_bridge.build_conversion_tasks(auto_bridge.hf_pretrained, model_list) diff --git a/skyrl/backends/skyrl_train/weight_sync/serialized_fp8.py b/skyrl/backends/skyrl_train/weight_sync/serialized_fp8.py new file mode 100644 index 0000000000..1029aceb47 --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/serialized_fp8.py @@ -0,0 +1,549 @@ +"""Convert Megatron-exported tensors to serialized vLLM FP8 formats.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from math import isfinite +from operator import index +from typing import Any, Iterator, Sequence + +import torch + +SERIALIZED_BLOCKWISE_FP8 = "serialized_blockwise" +SERIALIZED_MXFP8 = "serialized_mxfp8" +BLOCKWISE_128X128 = "blockwise_128x128" +MXFP8_1X32 = "mxfp8_1x32" +# Internal wire-format marker for Qwen3.5 MoE tensors that remain batched over +# experts. The receiver strips this marker and routes the tensor directly to +# vLLM's fused-MoE parameter loader instead of the ordinary HF-name loader. +SKYRL_BATCHED_MOE_FP8_PREFIX = "__skyrl_batched_moe_fp8__:" + + +def use_power_2_scales_default() -> bool: + """Return whether rollout weights use power-of-two block scales. + + The setting must match Transformer Engine. Hopper defaults to FP32 scales; + Blackwell launchers select power-of-two scales by setting + ``NVTE_FP8_BLOCK_SCALING_FP32_SCALES=0``. + """ + + scale_mode = os.getenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1") + if scale_mode not in {"0", "1"}: + raise ValueError( + "NVTE_FP8_BLOCK_SCALING_FP32_SCALES must be '0' (power-of-2) " f"or '1' (FP32 scales), got {scale_mode!r}" + ) + return scale_mode == "0" + + +def use_amax_epsilon_default() -> float: + """Read the TE blockwise amax floor used by the training weight quantizer.""" + + raw_value = os.getenv("NVTE_FP8_BLOCK_AMAX_EPSILON", "0") + try: + epsilon = float(raw_value) + except ValueError as exc: + raise ValueError(f"NVTE_FP8_BLOCK_AMAX_EPSILON must be a float, got {raw_value!r}") from exc + if not isfinite(epsilon) or epsilon < 0: + raise ValueError(f"NVTE_FP8_BLOCK_AMAX_EPSILON must be finite and non-negative, got {epsilon}") + return epsilon + + +_QWEN35_FP8_WEIGHT_SUFFIXES = ( + ".self_attn.q_proj.weight", + ".self_attn.k_proj.weight", + ".self_attn.v_proj.weight", + ".self_attn.o_proj.weight", + ".mlp.gate_proj.weight", + ".mlp.up_proj.weight", + ".mlp.down_proj.weight", + ".linear_attn.in_proj_qkv.weight", + ".linear_attn.in_proj_z.weight", + ".linear_attn.out_proj.weight", + # Shared-expert linears use FP8; router and shared-expert gates remain BF16. + ".mlp.shared_expert.gate_proj.weight", + ".mlp.shared_expert.up_proj.weight", + ".mlp.shared_expert.down_proj.weight", +) +_QWEN35_UNQUANTIZED_LINEAR_SUFFIXES = ( + ".in_proj_b", + ".in_proj_a", +) +_QWEN35_LINEAR_ATTN_PREFIX_TEMPLATES = ( + "{model_prefix}.layers.{layer_idx}.linear_attn", + "{model_prefix}.language_model.layers.{layer_idx}.linear_attn", +) +_QWEN35_VISION_ATTN_PROJ_PREFIX_TEMPLATES = ("{model_prefix}.visual.blocks.{block_idx}.attn.proj",) + + +@dataclass(frozen=True) +class MoeArchitectureSpec: + """HF expert export and vLLM projection naming for one model family.""" + + gate_up_suffixes: tuple[str, ...] + down_suffix: str + gate_up_fused: bool + vllm_projection_names: tuple[str, ...] + batched: bool + + +_QWEN35_MOE_SPEC = MoeArchitectureSpec( + gate_up_suffixes=(".gate_up_proj",), + down_suffix=".down_proj", + gate_up_fused=True, + vllm_projection_names=("gate_proj", "up_proj", "down_proj"), + batched=True, +) +_QWEN3_MOE_SPEC = MoeArchitectureSpec( + gate_up_suffixes=(".gate_up_proj",), + down_suffix=".down_proj", + gate_up_fused=True, + vllm_projection_names=("gate_proj", "up_proj", "down_proj"), + batched=True, +) +MOE_ARCHITECTURE_SPECS = { + "qwen3_moe": _QWEN3_MOE_SPEC, + "qwen3_5_moe": _QWEN35_MOE_SPEC, + "qwen3_5_moe_text": _QWEN35_MOE_SPEC, +} +_EXPERT_ONLY_MXFP8_IGNORED_MODULES = ( + "*.self_attn.*", + "*.linear_attn.*", + "*.mlp.gate", + "*.mlp.gate_up_proj", + "*.mlp.down_proj", + "*.mlp.shared_expert*", + "*lm_head*", + "*.visual.*", + "mtp.*", +) + + +def get_hf_model_type(hf_config: Any) -> str: + """Return the text model type used for expert serialization dispatch.""" + + text_config = getattr(hf_config, "text_config", None) or getattr(hf_config, "language_config", None) or hf_config + return str(getattr(text_config, "model_type", "") or getattr(hf_config, "model_type", "")) + + +def get_moe_architecture_spec(model_type: str) -> MoeArchitectureSpec: + """Return the registered MoE export layout.""" + + try: + return MOE_ARCHITECTURE_SPECS[model_type] + except KeyError as exc: + supported = ", ".join(sorted(MOE_ARCHITECTURE_SPECS)) + raise ValueError( + f"Serialized MXFP8 does not support model_type={model_type!r}; supported: {supported}" + ) from exc + + +def _normalize_block_size(block_size: Sequence[int]) -> tuple[int, int]: + try: + raw_values = tuple(block_size) + if any(isinstance(value, bool) for value in raw_values): + raise TypeError + values = tuple(index(value) for value in raw_values) + except (TypeError, ValueError) as exc: + raise ValueError(f"weight_block_size must contain exactly two positive integers, got {block_size!r}") from exc + if len(values) != 2 or any(value <= 0 for value in values): + raise ValueError(f"weight_block_size must contain exactly two positive integers, got {block_size!r}") + return values + + +@dataclass(frozen=True) +class SerializedFp8Config: + """Configuration for serialized FP8 rollout weight sync.""" + + scaling_mode: str = BLOCKWISE_128X128 + expert_only: bool = False + model_type: str | None = None + weight_block_size: tuple[int, int] = (128, 128) + power_2_scale: bool = field(default_factory=use_power_2_scales_default) + amax_epsilon: float = field(default_factory=use_amax_epsilon_default) + + def __post_init__(self) -> None: + if self.scaling_mode not in (BLOCKWISE_128X128, MXFP8_1X32): + raise ValueError(f"scaling_mode must be {BLOCKWISE_128X128!r} or {MXFP8_1X32!r}, got {self.scaling_mode!r}") + if type(self.expert_only) is not bool: + raise ValueError(f"expert_only must be a bool, got {self.expert_only!r}") + object.__setattr__(self, "weight_block_size", _normalize_block_size(self.weight_block_size)) + if type(self.power_2_scale) is not bool: + raise ValueError(f"power_2_scale must be a bool, got {self.power_2_scale!r}") + if not isfinite(self.amax_epsilon) or self.amax_epsilon < 0: + raise ValueError(f"amax_epsilon must be finite and non-negative, got {self.amax_epsilon}") + if self.expert_only: + if not self.model_type: + raise ValueError("expert_only serialized FP8 requires model_type") + get_moe_architecture_spec(self.model_type) + + +def serialized_fp8_config_for_mode(mode: str, *, model_type: str | None = None) -> SerializedFp8Config: + """Build the serializer configuration for a weight-sync mode.""" + + if mode == SERIALIZED_BLOCKWISE_FP8: + return SerializedFp8Config(model_type=model_type) + if mode == SERIALIZED_MXFP8: + return SerializedFp8Config( + scaling_mode=MXFP8_1X32, + expert_only=True, + model_type=model_type, + weight_block_size=(1, 32), + ) + raise ValueError(f"Unsupported fp8_weight_sync_mode={mode!r}") + + +def get_serialized_fp8_quantization_config( + weight_block_size: Sequence[int] = (128, 128), + ignored_layers: Sequence[str] | None = None, +) -> dict: + """Return vLLM's Hugging Face quantization config for serialized FP8.""" + + block_m, block_n = _normalize_block_size(weight_block_size) + qconfig = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [block_m, block_n], + } + if ignored_layers: + qconfig["ignored_layers"] = list(ignored_layers) + return qconfig + + +def get_serialized_mxfp8_quantization_config() -> dict: + """Return vLLM's ModelOpt config for expert-only serialized MXFP8.""" + + return { + "quant_method": "modelopt", + "quant_algo": "MXFP8", + "ignore": list(_EXPERT_ONLY_MXFP8_IGNORED_MODULES), + } + + +def is_qwen35_config(hf_config: Any) -> bool: + """Return whether an HF config uses the supported Qwen3.5 text layout.""" + + return get_hf_model_type(hf_config) in {"qwen3_5", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text"} + + +def get_qwen35_fp8_ignored_layers(hf_config: Any, model_prefix: str = "model") -> list[str]: + """Return Qwen3.5 vLLM module prefixes excluded from serialized FP8. + + Serialized sync excludes GDN ``in_proj_a`` and ``in_proj_b``. vLLM requires + every shard of the fused module to share a quantization scheme, so both + prefixes are ignored for text-only and conditional-generation checkpoints. + """ + + text_config = getattr(hf_config, "text_config", None) or getattr(hf_config, "language_config", None) or hf_config + if not is_qwen35_config(hf_config): + return [] + + layer_types = list(getattr(text_config, "layer_types", []) or []) + ignored: list[str] = [] + for layer_idx, layer_type in enumerate(layer_types): + if layer_type != "linear_attention": + continue + layer_prefixes = [] + for template in _QWEN35_LINEAR_ATTN_PREFIX_TEMPLATES: + prefix = template.format(model_prefix=model_prefix, layer_idx=layer_idx) + if prefix not in layer_prefixes: + layer_prefixes.append(prefix) + + for layer_prefix in layer_prefixes: + for suffix in _QWEN35_UNQUANTIZED_LINEAR_SUFFIXES: + ignored.append(f"{layer_prefix}{suffix}") + + # vLLM 0.23 may instantiate the vision tower for text-only runs. Its TP2 + # attention output is incompatible with 128-wide blocks, and ignore matching + # requires each block's exact module prefix. + vision_config = getattr(hf_config, "vision_config", None) or getattr(hf_config, "visual_config", None) + vision_depth = 0 + if vision_config is not None: + for attr in ("depth", "num_hidden_layers", "num_layers"): + value = getattr(vision_config, attr, None) + if isinstance(value, int) and value > 0: + vision_depth = value + break + for block_idx in range(vision_depth): + for template in _QWEN35_VISION_ATTN_PROJ_PREFIX_TEMPLATES: + ignored.append(template.format(model_prefix=model_prefix, block_idx=block_idx)) + return ignored + + +def should_use_serialized_fp8(mode: str | None) -> bool: + return mode in (SERIALIZED_BLOCKWISE_FP8, SERIALIZED_MXFP8) + + +def is_quantizable_weight(name: str, tensor: torch.Tensor, *, expert_only: bool = False) -> bool: + """Return whether an exported HF tensor should be serialized as FP8. + + vLLM's FP8 config applies to Linear modules. HF checkpoints also contain 2D + embedding/output weights, so keep known non-Linear weight tables unquantized. + """ + + if expert_only or not name.endswith(".weight") or tensor.ndim != 2: + return False + + return is_quantizable_weight_shape(name, tensor.shape) + + +def is_quantizable_weight_shape(name: str, shape: Sequence[int]) -> bool: + if not name.endswith(".weight") or len(shape) != 2: + return False + return name.endswith(_QWEN35_FP8_WEIGHT_SUFFIXES) + + +def scale_name_for_weight(name: str) -> str: + if not name.endswith(".weight"): + raise ValueError(f"FP8 scale can only be derived from .weight tensors: {name}") + return name[: -len(".weight")] + ".weight_scale_inv" + + +def mxfp8_scale_name_for_weight(name: str) -> str: + if not name.endswith(".weight"): + raise ValueError(f"MXFP8 scale can only be derived from .weight tensors: {name}") + return name[: -len(".weight")] + ".weight_scale" + + +def blockwise_cast_to_fp8( + weight: torch.Tensor, + block_size: Sequence[int], + power_2_scale: bool = False, + amax_epsilon: float = 0.0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2D tensor to vLLM's blockwise E4M3 checkpoint format. + + Returns ``weight_scale_inv`` such that + ``weight ~= qweight.float() * scale``. Power-of-two mode rounds scales up + to match Transformer Engine's UE8M0 rule. + """ + + if weight.ndim != 2: + raise ValueError(f"Blockwise FP8 expects a 2D tensor, got shape={tuple(weight.shape)}") + if not isfinite(amax_epsilon) or amax_epsilon < 0: + raise ValueError(f"amax_epsilon must be finite and non-negative, got {amax_epsilon}") + + block_m, block_n = _normalize_block_size(block_size) + rows, cols = weight.shape + padded_rows = ((rows + block_m - 1) // block_m) * block_m + padded_cols = ((cols + block_n - 1) // block_n) * block_n + + fp8_info = torch.finfo(torch.float8_e4m3fn) + weight_fp32 = weight.detach().to(torch.float32) + if padded_rows != rows or padded_cols != cols: + padded = weight_fp32.new_zeros((padded_rows, padded_cols)) + padded[:rows, :cols].copy_(weight_fp32) + else: + padded = weight_fp32 + + blocks = padded.view(padded_rows // block_m, block_m, padded_cols // block_n, block_n) + blocks = blocks.permute(0, 2, 1, 3) + # Match TE's amax floor, with a nonzero fallback for all-zero blocks. + scale = blocks.abs().amax(dim=(2, 3)).clamp(min=max(amax_epsilon, 1e-10)) / fp8_info.max + if power_2_scale: + # Rounding up preserves range and matches TE's power-of-two scale rule. + scale = torch.pow(2.0, torch.ceil(torch.log2(scale))) + q_blocks = (blocks / scale[:, :, None, None]).clamp(min=fp8_info.min, max=fp8_info.max) + q_blocks = q_blocks.to(torch.float8_e4m3fn) + q_padded = q_blocks.permute(0, 2, 1, 3).contiguous().view(padded_rows, padded_cols) + q_weight = q_padded[:rows, :cols].contiguous() + return q_weight, scale.to(torch.float32).contiguous() + + +def batched_blockwise_cast_to_fp8( + weight: torch.Tensor, + block_size: Sequence[int], + power_2_scale: bool = False, + amax_epsilon: float = 0.0, + expert_batch_size: int = 8, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 3D ``[experts, rows, cols]`` tensor blockwise. + + Quantizing several experts per operation avoids launching the full 2D + conversion pipeline once per expert, while bounded batches limit peak FP32 + workspace. + """ + + if weight.ndim != 3: + raise ValueError(f"Batched blockwise FP8 expects a 3D tensor, got shape={tuple(weight.shape)}") + if not isfinite(amax_epsilon) or amax_epsilon < 0: + raise ValueError(f"amax_epsilon must be finite and non-negative, got {amax_epsilon}") + if isinstance(expert_batch_size, bool) or not isinstance(expert_batch_size, int) or expert_batch_size <= 0: + raise ValueError(f"expert_batch_size must be a positive integer, got {expert_batch_size!r}") + + block_m, block_n = _normalize_block_size(block_size) + num_experts, rows, cols = weight.shape + padded_rows = ((rows + block_m - 1) // block_m) * block_m + padded_cols = ((cols + block_n - 1) // block_n) * block_n + row_blocks = padded_rows // block_m + col_blocks = padded_cols // block_n + + fp8_info = torch.finfo(torch.float8_e4m3fn) + q_weight = torch.empty(weight.shape, dtype=torch.float8_e4m3fn, device=weight.device) + scales = torch.empty( + (num_experts, row_blocks, col_blocks), + dtype=torch.float32, + device=weight.device, + ) + + for start in range(0, num_experts, expert_batch_size): + end = min(start + expert_batch_size, num_experts) + weight_fp32 = weight[start:end].detach().to(torch.float32).contiguous() + if padded_rows != rows or padded_cols != cols: + padded = weight_fp32.new_zeros((end - start, padded_rows, padded_cols)) + padded[:, :rows, :cols].copy_(weight_fp32) + else: + padded = weight_fp32 + + blocks = padded.view(end - start, row_blocks, block_m, col_blocks, block_n) + blocks = blocks.permute(0, 1, 3, 2, 4) + scale = blocks.abs().amax(dim=(3, 4)).clamp(min=max(amax_epsilon, 1e-10)) / fp8_info.max + if power_2_scale: + scale = torch.pow(2.0, torch.ceil(torch.log2(scale))) + q_blocks = (blocks / scale[:, :, :, None, None]).clamp(min=fp8_info.min, max=fp8_info.max) + q_blocks = q_blocks.to(torch.float8_e4m3fn) + q_padded = q_blocks.permute(0, 1, 3, 2, 4).contiguous().view(end - start, padded_rows, padded_cols) + q_weight[start:end].copy_(q_padded[:, :rows, :cols]) + scales[start:end].copy_(scale) + + return q_weight, scales + + +def _power_2_scales_to_e8m0(scales: torch.Tensor) -> torch.Tensor: + """Encode positive power-of-two FP32 scales as biased E8M0 exponents.""" + + exponent_bits = (scales.contiguous().view(torch.int32) >> 23) & 0xFF + return exponent_bits.to(torch.uint8) + + +def mxfp8_cast_to_fp8(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2D tensor to vLLM's row-major MXFP8 checkpoint format.""" + + if weight.ndim != 2: + raise ValueError(f"MXFP8 expects a 2D tensor, got shape={tuple(weight.shape)}") + if weight.shape[-1] % 32 != 0: + raise ValueError(f"MXFP8 requires the last dimension to be divisible by 32, got shape={tuple(weight.shape)}") + q_weight, scales = blockwise_cast_to_fp8(weight, (1, 32), power_2_scale=True) + return q_weight, _power_2_scales_to_e8m0(scales) + + +def batched_mxfp8_cast_to_fp8( + weight: torch.Tensor, + expert_batch_size: int = 8, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 3D expert tensor to row-major MXFP8 weights and E8M0 scales.""" + + if weight.ndim != 3: + raise ValueError(f"Batched MXFP8 expects a 3D tensor, got shape={tuple(weight.shape)}") + if isinstance(expert_batch_size, bool) or not isinstance(expert_batch_size, int) or expert_batch_size <= 0: + raise ValueError(f"expert_batch_size must be a positive integer, got {expert_batch_size!r}") + if weight.shape[-1] % 32 != 0: + raise ValueError(f"MXFP8 requires the last dimension to be divisible by 32, got shape={tuple(weight.shape)}") + + q_weight, scales = batched_blockwise_cast_to_fp8( + weight, + (1, 32), + power_2_scale=True, + expert_batch_size=expert_batch_size, + ) + return q_weight, _power_2_scales_to_e8m0(scales) + + +def batched_moe_expert_spec( + name: str, + model_type: str | None = None, +) -> tuple[str, tuple[str, ...], bool] | None: + """Parse a registered Megatron Bridge batched MoE tensor name. + + Returns ``(experts_base, projection_names, split_gate_up)`` or ``None``. + """ + + spec = _QWEN35_MOE_SPEC if model_type is None else get_moe_architecture_spec(model_type) + if not spec.batched or ".mlp.experts" not in name: + return None + gate_up_suffix = spec.gate_up_suffixes[0] + if name.endswith(gate_up_suffix): + return name[: -len(gate_up_suffix)], spec.vllm_projection_names[:2], spec.gate_up_fused + if name.endswith(spec.down_suffix): + return name[: -len(spec.down_suffix)], (spec.vllm_projection_names[-1],), False + return None + + +def iter_batched_moe_expert_fp8_tensors( + name: str, + tensor: torch.Tensor, + config: SerializedFp8Config, +) -> Iterator[tuple[str, torch.Tensor]]: + """Convert a batched expert tensor without expanding expert names. + + The old wire format emitted one weight and one scale tensor for every + expert/projection pair. Keeping the expert dimension intact reduces each + routed MoE layer from ``6 * num_experts`` tensors to six and lets vLLM use + its fused 3D loader. + """ + spec = batched_moe_expert_spec(name, config.model_type) + if spec is None: + raise ValueError(f"Not a batched MoE expert tensor: {name}") + if tensor.ndim != 3: + raise ValueError(f"Batched MoE expert tensor must be 3D, got shape={tuple(tensor.shape)}") + experts_base, proj_names, is_gate_up = spec + if is_gate_up: + if tensor.shape[1] % 2 != 0: + raise ValueError(f"Batched MoE gate_up_proj output dimension must be even, got shape={tuple(tensor.shape)}") + half = tensor.shape[1] // 2 + projection_tensors = (tensor[:, :half], tensor[:, half:]) + else: + projection_tensors = (tensor,) + + for proj, projection_tensor in zip(proj_names, projection_tensors): + if config.scaling_mode == MXFP8_1X32: + q_weight, scale = batched_mxfp8_cast_to_fp8(projection_tensor) + else: + q_weight, scale = batched_blockwise_cast_to_fp8( + projection_tensor, + config.weight_block_size, + config.power_2_scale, + config.amax_epsilon, + ) + weight_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{experts_base}.{proj}.weight" + yield weight_name, q_weight + scale_name = ( + mxfp8_scale_name_for_weight(weight_name) + if config.scaling_mode == MXFP8_1X32 + else scale_name_for_weight(weight_name) + ) + yield scale_name, scale + + +def iter_serialized_fp8_tensors( + name: str, + tensor: torch.Tensor, + target_dtype: torch.dtype, + config: SerializedFp8Config, +) -> Iterator[tuple[str, torch.Tensor]]: + """Yield vLLM checkpoint tensors for one Megatron-exported weight.""" + + if batched_moe_expert_spec(name, config.model_type) is not None: + yield from iter_batched_moe_expert_fp8_tensors(name, tensor, config) + return + + if config.expert_only and ".mlp.experts." in name and tensor.ndim >= 2: + raise ValueError(f"Unsupported routed-expert export tensor for model_type={config.model_type!r}: {name}") + + if is_quantizable_weight(name, tensor, expert_only=config.expert_only): + if config.scaling_mode == MXFP8_1X32: + q_weight, scale = mxfp8_cast_to_fp8(tensor) + scale_name = mxfp8_scale_name_for_weight(name) + else: + q_weight, scale = blockwise_cast_to_fp8( + tensor, + config.weight_block_size, + config.power_2_scale, + config.amax_epsilon, + ) + scale_name = scale_name_for_weight(name) + yield name, q_weight + yield scale_name, scale + return + + yield name, tensor.to(dtype=target_dtype) diff --git a/skyrl/backends/skyrl_train/weight_sync/transfer_strategy.py b/skyrl/backends/skyrl_train/weight_sync/transfer_strategy.py index 175074b6c6..d7b188598e 100644 --- a/skyrl/backends/skyrl_train/weight_sync/transfer_strategy.py +++ b/skyrl/backends/skyrl_train/weight_sync/transfer_strategy.py @@ -38,6 +38,7 @@ async def send_chunks( self, chunks: Iterable[WeightChunk], weight_metadata: Optional[Dict[str, list]] = None, + derive_metadata_from_chunks: bool = False, ) -> None: """Send chunks using this transfer strategy. @@ -47,8 +48,7 @@ async def send_chunks( Args: chunks: Iterable of WeightChunk objects to send. weight_metadata: Optional pre-computed metadata (names, dtype_names, shapes). - When provided, allows the sender to avoid materializing all chunks - to collect metadata upfront. + derive_metadata_from_chunks: Derive metadata from each transferred chunk. """ ... diff --git a/skyrl/backends/skyrl_train/workers/megatron/_fp8_block_amax_epsilon_patch.py b/skyrl/backends/skyrl_train/workers/megatron/_fp8_block_amax_epsilon_patch.py new file mode 100644 index 0000000000..0abfd456b1 --- /dev/null +++ b/skyrl/backends/skyrl_train/workers/megatron/_fp8_block_amax_epsilon_patch.py @@ -0,0 +1,76 @@ +"""Apply an optional amax floor to Transformer Engine blockwise FP8 recipes.""" + +from __future__ import annotations + +import os +from dataclasses import replace +from math import isfinite + +from loguru import logger + +_ENV_VAR = "NVTE_FP8_BLOCK_AMAX_EPSILON" +_PATCH_FLAG = "_skyrl_amax_epsilon_patched" +_QPARAM_FIELDS = ("fp8_quant_fwd_inp", "fp8_quant_fwd_weight", "fp8_quant_bwd_grad") +_MISSING = object() + + +def _configured_amax_epsilon() -> float | None: + raw_value = os.getenv(_ENV_VAR) + if raw_value is None or not raw_value.strip(): + return None + try: + epsilon = float(raw_value) + except ValueError as exc: + raise ValueError(f"{_ENV_VAR} must be a float, got {raw_value!r}") from exc + if not isfinite(epsilon) or epsilon < 0: + raise ValueError(f"{_ENV_VAR} must be finite and non-negative, got {epsilon}") + return epsilon or None + + +def apply_fp8_block_amax_epsilon_patch() -> None: + """Idempotently set ``amax_epsilon`` on TE 2.11 blockwise QParams. + + Once configured, fail if the TE API cannot be patched and verified; rollout + and training quantizers must use the same amax floor. + """ + + epsilon = _configured_amax_epsilon() + if epsilon is None: + return + + try: + from transformer_engine.common.recipe import Float8BlockScaling, Format + except Exception as exc: + raise RuntimeError( + f"{_ENV_VAR} is set, but Transformer Engine's Float8BlockScaling recipe is unavailable" + ) from exc + + if getattr(Float8BlockScaling, _PATCH_FLAG, None) == epsilon: + return + + original_qparams = {} + original_flag = getattr(Float8BlockScaling, _PATCH_FLAG, _MISSING) + try: + original_qparams = {name: getattr(Float8BlockScaling, name) for name in _QPARAM_FIELDS} + patched_qparams = {name: replace(qparams, amax_epsilon=epsilon) for name, qparams in original_qparams.items()} + for name, qparams in patched_qparams.items(): + setattr(Float8BlockScaling, name, qparams) + setattr(Float8BlockScaling, _PATCH_FLAG, epsilon) + + probe = Float8BlockScaling(fp8_format=Format.E4M3) + observed = {name: getattr(probe, name).amax_epsilon for name in _QPARAM_FIELDS} + if any(value != epsilon for value in observed.values()): + raise RuntimeError(f"fresh Float8BlockScaling instance reported {observed}") + except Exception as exc: + for name, qparams in original_qparams.items(): + setattr(Float8BlockScaling, name, qparams) + if original_flag is _MISSING: + try: + delattr(Float8BlockScaling, _PATCH_FLAG) + except AttributeError: + pass + else: + setattr(Float8BlockScaling, _PATCH_FLAG, original_flag) + raise RuntimeError(f"Failed to apply {_ENV_VAR}={epsilon} to Transformer Engine") from exc + + logger.info("Applied {}={} to TE Float8BlockScaling blockwise quantizers", _ENV_VAR, epsilon) diff --git a/skyrl/backends/skyrl_train/workers/megatron/expert_mxfp8.py b/skyrl/backends/skyrl_train/workers/megatron/expert_mxfp8.py new file mode 100644 index 0000000000..6788bdabae --- /dev/null +++ b/skyrl/backends/skyrl_train/workers/megatron/expert_mxfp8.py @@ -0,0 +1,138 @@ +"""Expert-only MXFP8 setup for Megatron.""" + +from __future__ import annotations + +import torch +from loguru import logger + + +def expert_mxfp8_recipe_dict(*, persistent: bool = False) -> dict: + """Return the per-module Transformer Engine recipe.""" + + high_precision = { + "transformer_engine_config_type": "TEQuantizationParams", + "training_recipe": {"override_quantized_autocast": True}, + "evaluation_recipe": {"override_quantized_autocast": True}, + } + mxfp8 = { + "transformer_engine_config_type": "TEQuantizationParams", + "training_recipe": { + "fp8_quantization_recipe": "mxfp8", + "fp8_format": "e4m3", + "override_quantized_autocast": True, + "override_nonquantized_autocast": True, + }, + "evaluation_recipe": { + "fp8_quantization_recipe": "mxfp8", + "fp8_format": "e4m3", + "override_quantized_autocast": True, + "override_nonquantized_autocast": True, + }, + } + if persistent: + mxfp8["training_recipe"]["fp8_param"] = True + mxfp8["evaluation_recipe"]["fp8_param"] = True + return { + "configs": {"expert_mxfp8": mxfp8, "high_precision": high_precision}, + "matchers": { + "routed_fc1": { + "type": "glob", + "pattern": "*.experts.linear_fc1", + "config": "expert_mxfp8", + "enabled": True, + }, + "routed_fc2": { + "type": "glob", + "pattern": "*.experts.linear_fc2", + "config": "expert_mxfp8", + "enabled": True, + }, + "default": { + "type": "glob", + "pattern": "*", + "config": "high_precision", + "enabled": True, + }, + }, + } + + +def configure_expert_mxfp8_provider(provider, *, persistent: bool = False) -> None: + """Configure routed experts for MXFP8 compute.""" + + from megatron.core.quantization.quant_config import RecipeConfig + + if getattr(provider, "quant_recipe", None) is not None: + raise ValueError("Expert MXFP8 conflicts with an existing quant_recipe") + if getattr(provider, "fp8", None) not in (None, False): + raise ValueError("Expert MXFP8 conflicts with global FP8 configuration") + if not provider.moe_grouped_gemm: + raise ValueError("Expert MXFP8 requires moe_grouped_gemm=true") + if provider.moe_router_dtype != "fp32": + raise ValueError("Expert MXFP8 requires moe_router_dtype=fp32") + for option in ("fp8_dot_product_attention", "fp8_multi_head_attention", "fp8_output_proj"): + if getattr(provider, option, False): + raise ValueError(f"Expert MXFP8 requires {option}=false") + provider.fp8 = "e4m3" + provider.fp8_recipe = "mxfp8" + # Pad each local expert immediately before GroupedLinear. Router-side + # padding can lose MXFP8 alignment after EP all-to-all redistribution. + provider.moe_router_padding_for_quantization = False + provider.quant_recipe = RecipeConfig.from_config_dict(expert_mxfp8_recipe_dict(persistent=persistent)) + + +def validate_expert_mxfp8_hardware() -> None: + """Require native Blackwell MXFP8 support.""" + + capability = torch.cuda.get_device_capability() + if capability not in ((10, 0), (10, 3)): + raise RuntimeError(f"Expert MXFP8 requires SM100 or SM103, got SM{capability[0]}{capability[1]}") + + +def is_routed_expert_linear(name: str) -> bool: + """Return whether a module name identifies routed expert FC1 or FC2.""" + + parts = name.split(".") + return any( + parts[index] == "experts" and parts[index + 1] in ("linear_fc1", "linear_fc2") + for index in range(len(parts) - 1) + ) + + +def audit_expert_mxfp8_modules(model_chunks) -> int: + """Verify that only routed expert linears execute quantized.""" + + matched = [] + errors = [] + for chunk in model_chunks: + for name, module in chunk.named_modules(): + if not hasattr(module, "will_execute_quantized"): + continue + targeted = is_routed_expert_linear(name) + quantized = module.will_execute_quantized(True) + quant_params = getattr(module, "te_quant_params", None) + recipe = getattr(getattr(quant_params, "training_recipe", None), "fp8_quantization_recipe", None) + uses_mxfp8 = getattr(recipe, "value", recipe) == "mxfp8" + if targeted and (not quantized or not uses_mxfp8): + errors.append(f"{name} did not enable MXFP8") + elif not targeted and (quantized or uses_mxfp8): + errors.append(f"{name} unexpectedly enabled quantization") + elif targeted: + matched.append(name) + + state = torch.tensor( + [bool(errors), len(matched)], + device=torch.cuda.current_device(), + dtype=torch.int32, + ) + torch.distributed.all_reduce(state) + if state[0].item(): + rank_errors = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(rank_errors, errors) + messages = [message for rank in rank_errors for message in rank] + raise RuntimeError("; ".join(messages)) + if state[1].item() == 0: + raise RuntimeError("Expert MXFP8 recipe matched no routed expert modules") + if torch.distributed.get_rank() == 0: + logger.info(f"Expert MXFP8 enabled for {state[1].item()} routed grouped-linear modules") + return state[1].item() diff --git a/skyrl/backends/skyrl_train/workers/megatron/fp8_param.py b/skyrl/backends/skyrl_train/workers/megatron/fp8_param.py new file mode 100644 index 0000000000..4e5dca3a34 --- /dev/null +++ b/skyrl/backends/skyrl_train/workers/megatron/fp8_param.py @@ -0,0 +1,148 @@ +"""Initialize optimizer state for persistent Transformer Engine FP8 parameters.""" + +from collections.abc import Mapping +from typing import Any + +import torch + +from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import is_fp8_enabled + + +def is_fp8_param_enabled(transformer_config_kwargs: Mapping[str, Any]) -> bool: + """Return whether dictionary config enables persistent FP8 parameters.""" + return is_fp8_enabled(transformer_config_kwargs.get("fp8_param", False)) + + +def _copy_model_shards_to_main_params( + megatron_optimizer: Any, + state_dict: Mapping[str, Any], +) -> int: + """Reload MCore FP32 master shards from converted checkpoint tensors. + + HybridDeviceOptimizer bypasses Megatron's state-dict-aware reload path. + Copying from ``state_dict`` preserves unquantized values for persistent FP8 + parameters and refreshes ordinary BF16 parameters. + """ + model_groups = getattr(megatron_optimizer, "model_float16_groups", None) + main_groups = getattr(megatron_optimizer, "shard_fp32_from_float16_groups", None) + model_fp32_groups = getattr(megatron_optimizer, "model_fp32_groups", None) + main_fp32_groups = getattr(megatron_optimizer, "shard_fp32_groups", None) + get_range = getattr(megatron_optimizer, "_get_model_param_range_map", None) + if model_groups is None or main_groups is None or not callable(get_range): + raise TypeError( + "Persistent FP8 parameter training with CPU optimizer offload requires " + "Megatron DistributedOptimizer shard metadata." + ) + + build_state_dict_map = getattr(megatron_optimizer, "_build_model_param_to_state_dict_param_map", None) + if not callable(build_state_dict_map): + raise TypeError( + "Persistent FP8 parameter training requires Megatron's " + "_build_model_param_to_state_dict_param_map() to reload exact checkpoint masters." + ) + state_dict_params = build_state_dict_map(state_dict) + + group_pairs = [(model_groups, main_groups)] + if model_fp32_groups is not None and main_fp32_groups is not None: + group_pairs.append((model_fp32_groups, main_fp32_groups)) + + copied = 0 + for grouped_model_params, grouped_main_params in group_pairs: + for model_group, main_group in zip(grouped_model_params, grouped_main_params): + for model_param, main_param in zip(model_group, main_group): + if main_param is None: + continue + + param_range = get_range(model_param)["param"] + if param_range.size != main_param.numel(): + raise RuntimeError( + "Persistent FP8 master-shard range does not match its FP32 master tensor: " + f"{param_range.size=} {main_param.numel()=}." + ) + + source_param = state_dict_params[model_param] + source = source_param.detach().reshape(-1)[param_range.start : param_range.end] + main_param.data.copy_(source.to(device=main_param.device, dtype=main_param.dtype)) + copied += 1 + return copied + + +def _sync_hybrid_device_optimizer_masters(hybrid_optimizer: Any) -> int: + """Refresh HybridDeviceOptimizer's secondary masters after checkpoint import. + + CPU copies are created before checkpoint import and are not refreshed by + MCore's public reload helper. Stale copies would overwrite weights on the + first optimizer step. + """ + copied = 0 + for cpu_param, gpu_param in getattr(hybrid_optimizer, "cpu_copys_map_gpu_param", {}).items(): + cpu_param.data.copy_( + gpu_param.detach().to(device=cpu_param.device, dtype=cpu_param.dtype), + non_blocking=False, + ) + copied += 1 + + for param, fp32_param in getattr(hybrid_optimizer, "param_to_fp32_param", {}).items(): + fp32_param.data.copy_( + param.detach().to(device=fp32_param.device, dtype=fp32_param.dtype), + non_blocking=False, + ) + copied += 1 + return copied + + +def _uses_hybrid_device_optimizer(megatron_optimizer: Any) -> bool: + """Identify MCore's CPU-offload optimizer without importing it eagerly.""" + optimizer = getattr(megatron_optimizer, "optimizer", None) + return hasattr(optimizer, "cpu_copys_map_gpu_param") and hasattr(optimizer, "param_to_fp32_param") + + +def initialize_fp8_param_optimizer_masters( + optimizer: Any, + *, + fp8_param: bool, + fp8_param_gather: bool, + state_dict: Mapping[str, Any] | None = None, +) -> int: + """Initialize optimizer masters from unquantized checkpoint shards. + + Persistent FP8 parameters cannot seed exact FP32 masters. For CPU offload, + refresh both distributed shards and HybridDeviceOptimizer's secondary copies. + """ + if not fp8_param: + return 0 + if not fp8_param_gather: + raise ValueError( + "Persistent FP8 parameters require ddp_config.fp8_param_gather=true " + "so updated FP32 master weights are requantized into FP8 compute weights." + ) + if state_dict is None: + raise ValueError( + "Persistent FP8 optimizer masters require Megatron-Bridge's exact unquantized checkpoint state." + ) + + optimizers = getattr(optimizer, "chained_optimizers", None) + if optimizers is None: + optimizers = [optimizer] + + initialized = 0 + with torch.no_grad(): + for megatron_optimizer in optimizers: + if _uses_hybrid_device_optimizer(megatron_optimizer): + copied = _copy_model_shards_to_main_params(megatron_optimizer, state_dict) + if copied == 0: + raise RuntimeError( + "Persistent FP8 parameter training did not find any model shards " + "to seed into CPU-offloaded optimizer masters." + ) + _sync_hybrid_device_optimizer_masters(megatron_optimizer.optimizer) + else: + reload_main_params = getattr(megatron_optimizer, "_copy_model_params_to_main_params", None) + if not callable(reload_main_params): + raise TypeError( + "Persistent FP8 parameter training requires a Megatron optimizer " + "with _copy_model_params_to_main_params()." + ) + reload_main_params(state_dict=state_dict) + initialized += 1 + return initialized diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 76e82b78d3..02430341be 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -1,3 +1,4 @@ +import gc import os import shutil from collections import defaultdict @@ -49,11 +50,27 @@ WeightChunk, WeightExtractor, ) +from skyrl.backends.skyrl_train.weight_sync.megatron_bridge_mappings import ( + get_packed_qwen3_moe_conversion_tasks, +) +from skyrl.backends.skyrl_train.weight_sync.serialized_fp8 import ( + get_hf_model_type, + iter_serialized_fp8_tensors, + serialized_fp8_config_for_mode, + should_use_serialized_fp8, +) +from skyrl.backends.skyrl_train.workers.megatron._fp8_block_amax_epsilon_patch import ( + apply_fp8_block_amax_epsilon_patch, +) from skyrl.backends.skyrl_train.workers.megatron.adapter_store import ( AdapterStore, LoraSignature, iter_opts, ) +from skyrl.backends.skyrl_train.workers.megatron.fp8_param import ( + initialize_fp8_param_optimizer_masters, + is_fp8_param_enabled, +) from skyrl.backends.skyrl_train.workers.megatron.megatron_model_wrapper import ( MegatronModelWrapper, ) @@ -107,12 +124,24 @@ def __init__( enable_bucketing: bool = False, bucket_size_threshold_GB: float = 1.0, training_dtype: torch.dtype = torch.bfloat16, + fp8_weight_sync_mode: Optional[str] = None, + model_type: Optional[str] = None, ): self.bridge = bridge self.actor_module = actor_module self.enable_bucketing = enable_bucketing self.bucket_size_threshold_GB = bucket_size_threshold_GB self.training_dtype = training_dtype + if fp8_weight_sync_mode is None: + self.serialized_fp8_config = None + elif should_use_serialized_fp8(fp8_weight_sync_mode): + self.serialized_fp8_config = serialized_fp8_config_for_mode( + fp8_weight_sync_mode, + model_type=model_type, + ) + else: + raise ValueError(f"Unsupported fp8_weight_sync_mode={fp8_weight_sync_mode!r}") + self._packed_qwen3_moe_export = self.serialized_fp8_config is not None and model_type == "qwen3_moe" # Defer bucket init to first extract_weights call. # At __init__ time the model may be CPU-offloaded (colocate_all), @@ -122,6 +151,11 @@ def __init__( self.bucket_index_groups = None self._buckets_initialized = False + def _get_conversion_tasks(self): + if self._packed_qwen3_moe_export: + return get_packed_qwen3_moe_conversion_tasks(self.bridge, self.actor_module) + return self.bridge.get_conversion_tasks(self.actor_module) + def _init_param_buckets(self): """Compute bucket boundaries (index groups) from parameter sizes. @@ -139,7 +173,7 @@ def _init_param_buckets(self): present in one call; splitting them across buckets causes expert weights to never be yielded. """ - weight_conversion_tasks = self.bridge.get_conversion_tasks(self.actor_module) + weight_conversion_tasks = self._get_conversion_tasks() def calculate_size_in_bytes(param, tp_size, ep_size): if param is None: @@ -210,6 +244,11 @@ def get_weight_metadata(self, dtype: torch.dtype) -> dict: (tensors are discarded immediately). Result is cached for subsequent calls. TODO (aaron): find a better way to get all metadata without materializing tensors. """ + if self.serialized_fp8_config is not None: + raise RuntimeError( + "Serialized FP8 metadata depends on quantized tensor contents; " + "consume extract_weights() chunks instead." + ) if hasattr(self, "_weight_metadata_cache"): return self._weight_metadata_cache @@ -233,7 +272,7 @@ def get_weight_metadata(self, dtype: torch.dtype) -> dict: else: # Build fresh tasks each sync so mapping objects have clean # PP-collective caches; reuse the pre-computed bucket structure. - fresh_tasks = self.bridge.get_conversion_tasks(self.actor_module) + fresh_tasks = self._get_conversion_tasks() for index_group in self.bucket_index_groups: bucket_tasks = [fresh_tasks[i] for i in index_group] for name, tensor in self.bridge.export_hf_weights( @@ -257,6 +296,18 @@ def _ensure_buckets_initialized(self): self._init_param_buckets() self._buckets_initialized = True + def _iter_sync_tensors( + self, + name: str, + tensor: torch.Tensor, + dtype: torch.dtype, + device: int, + ): + if self.serialized_fp8_config is not None: + tensor = tensor.to(device=device, non_blocking=True) + return iter_serialized_fp8_tensors(name, tensor, dtype, self.serialized_fp8_config) + return [(name, tensor.to(device=device, dtype=dtype, non_blocking=True))] + def extract_weights(self, dtype: torch.dtype): """Extract weights from Megatron model. @@ -274,22 +325,29 @@ def extract_weights(self, dtype: torch.dtype): hf_params_generator = self.bridge.export_hf_weights( self.actor_module, show_progress=False, - conversion_tasks=None, + conversion_tasks=self._get_conversion_tasks() if self._packed_qwen3_moe_export else None, ) for name, tensor in hf_params_generator: - tensor = tensor.to(device=device, dtype=dtype, non_blocking=True) + tensor_iter = self._iter_sync_tensors(name, tensor, dtype, device) - yield WeightChunk( - names=[name], - dtypes=[str(dtype)], - shapes=[list(tensor.shape)], - tensors=[tensor], - ) + names = [] + dtypes = [] + shapes = [] + tensors = [] + for out_name, out_tensor in tensor_iter: + out_tensor = out_tensor.contiguous() + names.append(out_name) + dtypes.append(str(out_tensor.dtype)) + shapes.append(list(out_tensor.shape)) + tensors.append(out_tensor) + + if tensors: + yield WeightChunk(names=names, dtypes=dtypes, shapes=shapes, tensors=tensors) else: # Build fresh tasks each sync so mapping objects have clean # PP-collective caches; reuse the pre-computed bucket structure. - fresh_tasks = self.bridge.get_conversion_tasks(self.actor_module) + fresh_tasks = self._get_conversion_tasks() for index_group in self.bucket_index_groups: bucket_tasks = [fresh_tasks[i] for i in index_group] @@ -306,13 +364,14 @@ def extract_weights(self, dtype: torch.dtype): tensors = [] for name, tensor in hf_params_generator: - # Move to device and convert dtype - tensor = tensor.to(device=device, dtype=dtype, non_blocking=True) + tensor_iter = self._iter_sync_tensors(name, tensor, dtype, device) - names.append(name) - dtypes_list.append(str(dtype)) - shapes.append(list(tensor.shape)) - tensors.append(tensor) + for out_name, out_tensor in tensor_iter: + out_tensor = out_tensor.contiguous() + names.append(out_name) + dtypes_list.append(str(out_tensor.dtype)) + shapes.append(list(out_tensor.shape)) + tensors.append(out_tensor) # Yield one chunk containing all parameters in this bucket if tensors: @@ -382,6 +441,8 @@ def init_configs( enable_mtp=False, language_model_only=False, bridge_weights_path=None, + expert_mxfp8=False, + expert_mxfp8_persistent=False, ): """ Initialize the Megatron-Bridge bridge and provider objects + hf_config and tokenizer @@ -393,6 +454,7 @@ def init_configs( masters and fake-quantizes them in the forward pass. Tokenizer + HF config (the logical model identity) still come from ``model_path``. """ + apply_fp8_block_amax_epsilon_patch() tokenizer = get_tokenizer(model_path, trust_remote_code=True) hf_config_original = AutoConfig.from_pretrained(model_path, trust_remote_code=True) @@ -422,6 +484,7 @@ def init_configs( for key in ("recompute_granularity", "recompute_method", "recompute_num_layers"): transformer_config_kwargs[key] = None + fp8_param_enabled = is_fp8_param_enabled(transformer_config_kwargs) or expert_mxfp8_persistent bridge_source = bridge_weights_path or model_path if bridge_weights_path: logger.info( @@ -439,7 +502,11 @@ def init_configs( "(native GDN thd packing path; vision tower dropped)" ) - provider = bridge.to_megatron_provider() + # Defer persistent-FP8 checkpoint import until the bridge can expose + # unquantized converted shards for optimizer-master initialization. + provider = bridge.to_megatron_provider(load_weights=not fp8_param_enabled) + if fp8_param_enabled: + provider.perform_initialization = False if not enable_mtp and getattr(provider, "mtp_num_layers", None): logger.info(f"Disabling MTP for training (mtp_num_layers={provider.mtp_num_layers} -> None)") @@ -489,6 +556,13 @@ def init_configs( for k, v in transformer_config_kwargs.items(): setattr(provider, k, v) + if expert_mxfp8: + from skyrl.backends.skyrl_train.workers.megatron.expert_mxfp8 import ( + configure_expert_mxfp8_provider, + ) + + configure_expert_mxfp8_provider(provider, persistent=expert_mxfp8_persistent) + # MTP head count: megatron-bridge infers provider.mtp_num_layers from the model's HF config. if not enable_mtp: provider.mtp_num_layers = None @@ -532,6 +606,8 @@ def init_configs( # the bridge weights path only under fake-INT4 QAT (INT4 model.path, BF16 # bridge weights); used so saved LoRA adapters reference the INT4 base. self._logical_model_path = model_path + self._deferred_fp8_param_weight_load = fp8_param_enabled + self._fp8_param_unquantized_state_dict = None # strategy.hf_config is the on-disk source-of-truth used by # save_hf_configs and must NOT carry runtime overrides like @@ -540,6 +616,40 @@ def init_configs( self.tokenizer = tokenizer self.enable_router_replay = megatron_config.moe_enable_routing_replay + def _load_deferred_fp8_param_weights(self, *, retain_unquantized_state: bool = False) -> None: + """Load deferred FP8 weights and optionally retain exact master tensors.""" + if not self._deferred_fp8_param_weight_load: + return + + original_export_dtype = self.bridge.export_weight_dtype + try: + # The FP8 export path retains converted, unquantized local shards. + # Only policy workers with an optimizer keep this additional copy. + if retain_unquantized_state: + self.bridge.export_weight_dtype = "fp8" + self.bridge.load_hf_weights(self.actor_module) + state_dict = getattr(self.bridge, "unquantized_state_dict", None) + if retain_unquantized_state and not state_dict: + raise RuntimeError( + "Megatron-Bridge did not capture unquantized checkpoint shards " + "for persistent FP8 optimizer initialization." + ) + self._fp8_param_unquantized_state_dict = state_dict if retain_unquantized_state else None + finally: + self.bridge.export_weight_dtype = original_export_dtype + + def _release_fp8_param_unquantized_state(self) -> None: + """Release temporary checkpoint shards after deferred weight import.""" + if not self._deferred_fp8_param_weight_load: + return + self._fp8_param_unquantized_state_dict = None + if self.bridge is not None: + # Clear both owners; either reference keeps the unquantized shard in HBM. + self.bridge.unquantized_state_dict = None + self._deferred_fp8_param_weight_load = False + gc.collect() + torch.cuda.empty_cache() + def configure_lora(self, lora_config, lora_type: Optional[str] = "lora"): if lora_type == "lora": self.lora_cls = LoRA( @@ -889,6 +999,13 @@ def init_model(self, model_path, num_training_steps: int = 1e9): # Fake-INT4 QAT: install the MoE expert fake-quant hook and (when the # served checkpoint is INT4) redirect the trainer's BF16 master weights. bridge_weights_path = self._maybe_setup_fake_int4_qat() + expert_mxfp8 = self.cfg.policy.model.expert_mxfp8 + if expert_mxfp8.enabled and expert_mxfp8.training: + from skyrl.backends.skyrl_train.workers.megatron.expert_mxfp8 import ( + validate_expert_mxfp8_hardware, + ) + + validate_expert_mxfp8_hardware() # initialize the bridge and provider objects self.init_configs( @@ -901,6 +1018,8 @@ def init_model(self, model_path, num_training_steps: int = 1e9): language_model_only=self.cfg.policy.language_model_only, bridge_weights_path=bridge_weights_path, enable_mtp=self.cfg.mtp.enabled, + expert_mxfp8=expert_mxfp8.enabled and expert_mxfp8.training, + expert_mxfp8_persistent=expert_mxfp8.enabled and expert_mxfp8.training and expert_mxfp8.persistent, ) if self.enable_router_replay: @@ -926,6 +1045,12 @@ def init_model(self, model_path, num_training_steps: int = 1e9): lora_type=self.cfg.policy.megatron_config.lora_config.lora_type, bf16=self.cfg.bf16, ) + if expert_mxfp8.enabled and expert_mxfp8.training: + from skyrl.backends.skyrl_train.workers.megatron.expert_mxfp8 import ( + audit_expert_mxfp8_modules, + ) + + audit_expert_mxfp8_modules(self.actor_module) if self._local_rank == 0 and not os.path.exists( model_path @@ -933,6 +1058,8 @@ def init_model(self, model_path, num_training_steps: int = 1e9): snapshot_download(model_path) # will be no-op if already downloaded torch.distributed.barrier() + self._load_deferred_fp8_param_weights(retain_unquantized_state=not self.cfg.policy.inference_only_init) + if self._rank == 0: print_model_size(self.actor_module[0]) @@ -950,7 +1077,18 @@ def init_model(self, model_path, num_training_steps: int = 1e9): self.cfg.policy.optimizer_config, self.cfg.policy.megatron_config.optimizer_config_kwargs ) self.optimizer = get_megatron_optimizer(self.actor_module, optim_config) - + fp8_param_masters = initialize_fp8_param_optimizer_masters( + self.optimizer, + fp8_param=is_fp8_param_enabled(self.cfg.policy.megatron_config.transformer_config_kwargs) + or (expert_mxfp8.enabled and expert_mxfp8.training and expert_mxfp8.persistent), + fp8_param_gather=self.cfg.policy.megatron_config.ddp_config.fp8_param_gather, + state_dict=self._fp8_param_unquantized_state_dict, + ) + if fp8_param_masters: + logger.info( + "Initialized {} persistent-FP8 optimizer master shard group(s) from exact checkpoint shards.", + fp8_param_masters, + ) # create scheduler self.scheduler = get_megatron_optimizer_param_scheduler( optimizer=self.optimizer, @@ -969,6 +1107,8 @@ def init_model(self, model_path, num_training_steps: int = 1e9): f"({n_local} head main params on rank {self._rank}; 0 is normal under DP sharding)" ) + self._release_fp8_param_unquantized_state() + # create worker model self.model = MegatronModelWrapper( config=self.cfg, @@ -1360,6 +1500,8 @@ async def init_weight_sync_state(self, inference_engine_client, inference_engine enable_bucketing=True, bucket_size_threshold_GB=inference_engine_cfg.weight_transfer_threshold_cuda_ipc_GB, training_dtype=torch.bfloat16 if self.cfg.bf16 else torch.float32, + fp8_weight_sync_mode=inference_engine_cfg.fp8_weight_sync_mode, + model_type=get_hf_model_type(self.strategy.hf_config), ) async def _save_lora_adapters_and_sync( @@ -1457,10 +1599,15 @@ async def broadcast_to_inference_engines( # CUDA-IPC path calls cudaIpcGetMemHandle, which is incompatible with the # VMM addresses expandable segments uses. with self._expandable_segments_disabled_for_sync(): - weight_metadata = self.weight_extractor.get_weight_metadata(generator_dtype) + weight_iterator = self.weight_extractor.extract_weights(generator_dtype) + derive_metadata_from_chunks = self.weight_extractor.serialized_fp8_config is not None + weight_metadata = ( + None if derive_metadata_from_chunks else self.weight_extractor.get_weight_metadata(generator_dtype) + ) await self._weight_transfer_sender.send_chunks( - self.weight_extractor.extract_weights(generator_dtype), + weight_iterator, weight_metadata=weight_metadata, + derive_metadata_from_chunks=derive_metadata_from_chunks, ) if cache_reset_task is not None: @@ -1640,6 +1787,9 @@ def init_model(self, model_path, num_training_steps: int = 1e9): snapshot_download(model_path) # will be no-op if already downloaded torch.distributed.barrier() + self._load_deferred_fp8_param_weights() + self._release_fp8_param_unquantized_state() + # load weights if self._rank == 0: print_model_size(self.actor_module[0]) diff --git a/skyrl/train/config/__init__.py b/skyrl/train/config/__init__.py index c00f629035..408ba00c85 100644 --- a/skyrl/train/config/__init__.py +++ b/skyrl/train/config/__init__.py @@ -10,6 +10,7 @@ DPPOConfig, DynamicSamplingConfig, EnvironmentConfig, + ExpertMxfp8Config, FSDPConfig, FullyAsyncConfig, GeneratorConfig, @@ -60,6 +61,7 @@ "InferenceEngineConfig", "MTPConfig", "EnvironmentConfig", + "ExpertMxfp8Config", "ModelConfig", "SkyRLLoraConfig", "OptimizerConfig", diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 86ab816f26..fe8eb427d0 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -100,6 +100,16 @@ class SkyRLLoraConfig(BaseConfig): >= ``max_loras`` if explicitly set.""" +@dataclass +class ExpertMxfp8Config(BaseConfig): + """MXFP8 compute for routed MoE experts.""" + + enabled: bool = False + training: bool = True + rollout: bool = True + persistent: bool = False + + @dataclass class FakeInt4QatConfig(BaseConfig): """Fake-INT4 quantization-aware training for MoE experts (Megatron only). @@ -138,6 +148,7 @@ class FakeInt4QatConfig(BaseConfig): class ModelConfig(BaseConfig): path: Optional[str] = None lora: SkyRLLoraConfig = field(default_factory=SkyRLLoraConfig) + expert_mxfp8: ExpertMxfp8Config = field(default_factory=ExpertMxfp8Config) fake_int4_qat: FakeInt4QatConfig = field(default_factory=FakeInt4QatConfig) def __post_init__(self) -> None: @@ -196,6 +207,8 @@ class MegatronDDPConfig(BaseConfig): grad_reduce_in_fp32: bool = True overlap_grad_reduce: bool = False overlap_param_gather: bool = False + fp8_param_gather: bool = False + reuse_grad_buf_for_mxfp8_param_ag: bool = False average_in_collective: bool = True @@ -744,6 +757,10 @@ class InferenceEngineConfig(BaseConfig): model_dtype: str = "bfloat16" """Should match the dtype used by the inference engine.""" + fp8_weight_sync_mode: Optional[str] = None + """Optional rollout weight format. ``"serialized_blockwise"`` serializes + supported linears with 128x128 scales; ``"serialized_mxfp8"`` serializes + routed experts with 1x32 E8M0 scales.""" run_engines_locally: bool = True num_engines: int = 1 backend: str = "vllm" @@ -1041,6 +1058,46 @@ def __post_init__(self): "sync preserves the inference engine's INT4 base weights." ) + expert_mxfp8 = self.policy.model.expert_mxfp8 + if expert_mxfp8.enabled: + if self.strategy != "megatron": + raise ValueError("Expert MXFP8 is only supported with trainer.strategy=megatron") + if not expert_mxfp8.training and not expert_mxfp8.rollout: + raise ValueError("Expert MXFP8 must enable training, rollout, or both") + if expert_mxfp8.persistent: + if not expert_mxfp8.training: + raise ValueError("Persistent expert MXFP8 requires expert_mxfp8.training=true") + if not self.policy.megatron_config.ddp_config.fp8_param_gather: + raise ValueError( + "Persistent expert MXFP8 requires policy.megatron_config.ddp_config.fp8_param_gather=true" + ) + if self.policy.model.fake_int4_qat.enabled: + raise ValueError("Expert MXFP8 and fake INT4 QAT are mutually exclusive") + if expert_mxfp8.training: + if not self.policy.megatron_config.moe_grouped_gemm: + raise ValueError("Expert MXFP8 training requires moe_grouped_gemm=true") + if self.policy.megatron_config.moe_router_dtype != "fp32": + raise ValueError("Expert MXFP8 training requires moe_router_dtype=fp32") + transformer_overrides = self.policy.megatron_config.transformer_config_kwargs + if transformer_overrides.get("moe_grouped_gemm") is False: + raise ValueError("Expert MXFP8 training requires moe_grouped_gemm=true") + if transformer_overrides.get("moe_router_dtype", "fp32") != "fp32": + raise ValueError("Expert MXFP8 training requires moe_router_dtype=fp32") + conflicts = {"fp8", "quant_recipe"} & transformer_overrides.keys() + conflicts.update( + option + for option in ("fp8_dot_product_attention", "fp8_multi_head_attention", "fp8_output_proj") + if transformer_overrides.get(option) + ) + if conflicts: + raise ValueError(f"Expert MXFP8 conflicts with transformer_config_kwargs: {sorted(conflicts)}") + if ( + expert_mxfp8.rollout + and self.policy.model.lora.rank > 0 + and not self.policy.megatron_config.lora_config.merge_lora + ): + raise ValueError("Expert MXFP8 rollout does not support unmerged LoRA") + if self.logprobs_chunk_size is not None and ( not isinstance(self.logprobs_chunk_size, int) or self.logprobs_chunk_size <= 0 ): diff --git a/skyrl/train/dataset/collators.py b/skyrl/train/dataset/collators.py index 45965f695a..742d046faa 100644 --- a/skyrl/train/dataset/collators.py +++ b/skyrl/train/dataset/collators.py @@ -71,7 +71,8 @@ class PackedDataCollator: Flow: 1. Compute per-example sequence lengths. - 2. FFD-pack with ``bin_capacity = max_tokens_per_microbatch``, + 2. FFD-pack using each sequence's alignment-padded footprint and + ``bin_capacity = max(max_tokens_per_microbatch, align_size)``, ``min_bin_count = dp_size``, ``bin_count_multiple = dp_size``. 3. Round-robin assign bins to DP shards (this happens implicitly inside ``MeshDispatch.dispatch`` because the rows are laid out in shard-major @@ -132,20 +133,23 @@ def __call__(self, examples: list, batch_size: int) -> TrainingInputBatch: pp_size = self.pp_size cp_size = self.cp_size # Each sub-seq's padded length must satisfy these divisibility - # constraints, which is why ``align_size`` carries all factors: + # constraints, which is why ``align_size`` carries all required factors: # - Sequence Parallelism (auto-on when tp>1) shards along the seq # dim, so each segment must be divisible by ``tp_size``. # - Context Parallelism splits each segment into ``2*cp_size`` equal # load-balanced causal chunks, so each segment must be divisible by # ``2*cp_size``. - # - When FP8 is enabled, Transformer Engine GEMMs require each CP - # rank's local token slab to be 16-aligned; globally this means - # ``16*cp_size``. + # - FP8 requires 16-token local slabs for TP=1. With sequence + # parallelism, TE quantizes all-gather inputs in 128-token blocks, + # so the global segment includes the TP and CP shard factors. # This MUST stay in lockstep with the worker's preprocess_packed_seqs # (megatron_utils.py): if the divisors drift, the per-rank CP/SP # gather/scatter offsets silently corrupt loss/grads (no crash). align_size = get_packed_seq_align_size(tp_size, cp_size, fp8_enabled=self.fp8_enabled) + def _round_up(x: int, multiple: int) -> int: + return ((x + multiple - 1) // multiple) * multiple + dp_size = self.dp_size # ------------------------------------------------------------------ @@ -174,14 +178,18 @@ def __call__(self, examples: list, batch_size: int) -> TrainingInputBatch: # same number of micro-batches. Forcing the global bin count to a # multiple of ``dp_size`` makes the per-DP-rank bin count (and thus # ``num_microbatches``) identical across ranks. + # Pack aligned footprints so per-sequence padding cannot exceed the row + # budget. Allow at least one alignment unit per bin. + packing_lengths = [_round_up(length, align_size) for length in seq_lengths] + packing_capacity = max(bin_capacity, align_size) bin_count_multiple = dp_size packer = make_seq_packer( "first_fit_decreasing", - bin_capacity=bin_capacity, + bin_capacity=packing_capacity, min_bin_count=bin_count_multiple, bin_count_multiple=bin_count_multiple, ) - bins: List[List[int]] = packer.pack(seq_lengths) + bins: List[List[int]] = packer.pack(packing_lengths) # Assign bins to DP shards via round-robin (bin_idx % shards). # Concretely we want the resulting layout to be shard-major: @@ -201,9 +209,6 @@ def __call__(self, examples: list, batch_size: int) -> TrainingInputBatch: # 3. Compute packed-row lengths (with align_size padding per sub-seq) # and the global max packed length (for PP > 1 uniform padding). # ------------------------------------------------------------------ - def _round_up(x: int, m: int) -> int: - return ((x + m - 1) // m) * m - bin_packed_lengths: List[int] = [] bin_subseq_lengths: List[List[int]] = [] # one list per bin row for bin_indices in flat_bins: @@ -239,7 +244,7 @@ def _round_up(x: int, m: int) -> int: n_samples = len(examples) logger.info( f"sequence packing | packed {n_samples} samples into {num_bins} bins " - f"(~{num_bins // dp_size}/DP rank, bin_capacity={bin_capacity} tokens)" + f"(~{num_bins // dp_size}/DP rank, bin_capacity={packing_capacity} aligned tokens)" ) sequences = torch.full((num_bins, max_packed_len), pad_token_id, dtype=torch.long) @@ -279,8 +284,7 @@ def _round_up(x: int, m: int) -> int: # p_local = s - 1 (last token of sub-seq): mask = 0. # Already zero by initialization. - # Advance row_offset, padding sub-seq to the TP/CP layout - # multiple, plus FP8's 16-token local-rank multiple when active. + # Match the aligned footprint consumed by preprocess_packed_seqs. row_offset += _round_up(s, align_size) # The total_nonpad we just counted matches sum(loss_mask). diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index 78dab4e84f..667f2bf3f8 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -1511,12 +1511,14 @@ def _execute_training_step(self, model: str, data: TrainingInputBatch) -> Dict[s # Training loop over epochs and mini-batches for _epoch in range(self.cfg.trainer.update_epochs_per_batch): for chunk_refs in all_chunk_refs: - status = self.dispatch.forward_backward_from_staged(model, chunk_refs) + with Timer(f"{model}_forward_backward", self.all_timings): + status = self.dispatch.forward_backward_from_staged(model, chunk_refs) for k, v in status.metrics.items(): all_metrics[k].append(v) # Optimizer step after each mini batch - grad_norm = self.dispatch.optim_step(model) + with Timer(f"{model}_optim_step", self.all_timings): + grad_norm = self.dispatch.optim_step(model) if grad_norm is not None: all_metrics["grad_norm"].append(grad_norm) diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index 777704a3a3..bf345c6ed9 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -20,6 +20,11 @@ ) from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import is_fp8_enabled +from skyrl.backends.skyrl_train.weight_sync.serialized_fp8 import ( + SERIALIZED_BLOCKWISE_FP8, + SERIALIZED_MXFP8, +) from skyrl.env_vars import ( SKYRL_DUMP_INFRA_LOG_TO_STDOUT, SKYRL_LD_LIBRARY_PATH_EXPORT, @@ -198,6 +203,18 @@ def validate_megatron_cfg(cfg: SkyRLTrainConfig): assert ie_cfg.backend == "vllm", "only vllm is supported for with megatron" assert cfg.trainer.critic.model.path is None, "only GRPO training is currently supported for megatron" + policy_cfg = cfg.trainer.policy + policy_fp8_param = is_fp8_enabled(policy_cfg.megatron_config.transformer_config_kwargs.get("fp8_param")) + if ( + policy_fp8_param + and not policy_cfg.inference_only_init + and not policy_cfg.megatron_config.ddp_config.fp8_param_gather + ): + raise ValueError( + "Persistent policy fp8_param training requires " + "trainer.policy.megatron_config.ddp_config.fp8_param_gather=true" + ) + if cfg.trainer.policy.megatron_config.moe_enable_routing_replay: assert ( cfg.generator.inference_engine.enable_return_routed_experts @@ -522,6 +539,22 @@ def validate_inference_engine_cfg(cfg: SkyRLTrainConfig): """ ie_cfg = cfg.generator.inference_engine + serialized_fp8_modes = (SERIALIZED_BLOCKWISE_FP8, SERIALIZED_MXFP8) + if ie_cfg.fp8_weight_sync_mode not in (None, *serialized_fp8_modes): + raise ValueError( + f"Unsupported fp8_weight_sync_mode={ie_cfg.fp8_weight_sync_mode!r}; " + f"expected one of {serialized_fp8_modes!r} or None" + ) + if ie_cfg.fp8_weight_sync_mode in serialized_fp8_modes: + if cfg.trainer.strategy != "megatron": + raise ValueError(f"{ie_cfg.fp8_weight_sync_mode} FP8 weight sync requires trainer.strategy='megatron'") + lora_cfg = cfg.trainer.policy.model.lora + if lora_cfg.rank > 0 and not cfg.trainer.policy.megatron_config.lora_config.merge_lora: + raise ValueError( + f"{ie_cfg.fp8_weight_sync_mode} FP8 weight sync requires full-weight updates; " + "Megatron LoRA with merge_lora=false syncs adapters only" + ) + if ie_cfg.enable_pd: assert ie_cfg.num_prefill > 0, "num_prefill must be > 0 when enable_pd=True" assert ( @@ -602,8 +635,7 @@ def _validate_new_inference_cfg(cfg: SkyRLTrainConfig): if not cfg.generator.inference_engine.run_engines_locally and not (has_external_proxy or has_external_servers): raise ValueError( - "generator.inference_engine.run_engines_locally=false requires " - "external_proxy_url or external_server_urls." + "generator.inference_engine.run_engines_locally=false requires external_proxy_url or external_server_urls." ) @@ -805,6 +837,54 @@ def prepare_runtime_environment(cfg: SkyRLTrainConfig) -> dict[str, str]: ) env_vars["SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S"] = health_timeout + # Forward one block-scale contract to all Ray actors. Hopper defaults to FP32; + # Blackwell launchers explicitly select power-of-two scales. + serialized_fp8 = cfg.generator.inference_engine.fp8_weight_sync_mode == SERIALIZED_BLOCKWISE_FP8 + use_ref_model = cfg.trainer.algorithm.use_kl_loss or cfg.trainer.algorithm.use_kl_in_reward + policy_megatron_config = getattr(cfg.trainer.policy, "megatron_config", None) + ref_megatron_config = getattr(cfg.trainer.ref, "megatron_config", None) + policy_transformer_kwargs = getattr(policy_megatron_config, "transformer_config_kwargs", None) or {} + ref_transformer_kwargs = getattr(ref_megatron_config, "transformer_config_kwargs", None) or {} + policy_fp8_param = is_fp8_enabled(policy_transformer_kwargs.get("fp8_param")) + ref_fp8_param = use_ref_model and is_fp8_enabled(ref_transformer_kwargs.get("fp8_param")) + fp8_compute = is_fp8_enabled(policy_transformer_kwargs.get("fp8")) or ( + use_ref_model and is_fp8_enabled(ref_transformer_kwargs.get("fp8")) + ) + fp8_contract_enabled = serialized_fp8 or fp8_compute or policy_fp8_param or ref_fp8_param + + fp8_env_defaults: dict[str, str] = {} + configured_scale_mode = os.environ.get("NVTE_FP8_BLOCK_SCALING_FP32_SCALES") + if fp8_contract_enabled or configured_scale_mode is not None: + scale_mode = configured_scale_mode or "1" + if scale_mode not in {"0", "1"}: + raise ValueError("NVTE_FP8_BLOCK_SCALING_FP32_SCALES must be '0' (power-of-2) or '1' (FP32 scales).") + + if scale_mode == "0" and (policy_fp8_param or ref_fp8_param): + raise ValueError( + "Persistent fp8_param requires FP32 block scales. Blackwell only supports " + "power-of-2 block scales, so use fp8_param=false on Blackwell." + ) + + if fp8_contract_enabled: + fp8_env_defaults["NVTE_FP8_BLOCK_SCALING_FP32_SCALES"] = scale_mode + if serialized_fp8 and scale_mode == "1": + e8m0_mode = os.environ.get("VLLM_USE_DEEP_GEMM_E8M0", "0") + if e8m0_mode != "0": + raise ValueError( + "FP32 block scales require VLLM_USE_DEEP_GEMM_E8M0=0 so vLLM " + "does not requantize them to power-of-2 scales." + ) + fp8_env_defaults["VLLM_USE_DEEP_GEMM_E8M0"] = e8m0_mode + + for var_name in ( + "NVTE_FP8_BLOCK_SCALING_FP32_SCALES", + "NVTE_FP8_BLOCK_AMAX_EPSILON", + "VLLM_USE_DEEP_GEMM_E8M0", + ): + if value := os.environ.get(var_name, fp8_env_defaults.get(var_name)): + logger.info(f"Exporting `{var_name}` to ray runtime env: {value}") + env_vars[var_name] = value + return env_vars @@ -1018,7 +1098,7 @@ def str_to_torch_dtype(dtype: str) -> torch.dtype: def format_gib(mem_bytes: int) -> str: - return f"{mem_bytes / (1024 ** 3):.2f} GiB" + return f"{mem_bytes / (1024**3):.2f} GiB" def print_mem(tag: str, mem: dict): diff --git a/tests/backends/skyrl_train/distributed/test_packing_utils.py b/tests/backends/skyrl_train/distributed/test_packing_utils.py index eb99bd2d6d..092ecd10bc 100644 --- a/tests/backends/skyrl_train/distributed/test_packing_utils.py +++ b/tests/backends/skyrl_train/distributed/test_packing_utils.py @@ -30,10 +30,25 @@ def test_packed_alignment_uses_layout_only_without_fp8(): def test_packed_alignment_adds_fp8_local_rank_multiple(): - assert get_packed_seq_align_size(tp_size=4, cp_size=1, fp8_enabled=True) == 16 + assert get_packed_seq_align_size(tp_size=4, cp_size=1, fp8_enabled=True) == 512 assert get_packed_seq_align_size(tp_size=1, cp_size=2, fp8_enabled=True) == 32 + assert get_packed_seq_align_size(tp_size=2, cp_size=1, fp8_enabled=True) == 256 + assert get_packed_seq_align_size(tp_size=2, cp_size=2, fp8_enabled=True) == 512 def test_unpacked_alignment_adds_fp8_multiple_only_when_enabled(): assert get_unpacked_seq_align_size(tp_size=4) == 4 - assert get_unpacked_seq_align_size(tp_size=4, fp8_enabled=True) == 16 + assert get_unpacked_seq_align_size(tp_size=1, fp8_enabled=True) == 16 + assert get_unpacked_seq_align_size(tp_size=2, fp8_enabled=True) == 256 + assert get_unpacked_seq_align_size(tp_size=4, fp8_enabled=True) == 512 + + +@pytest.mark.parametrize(("tp_size", "cp_size"), [(0, 1), (1, 0), (-1, 1)]) +def test_packed_alignment_rejects_nonpositive_parallel_sizes(tp_size, cp_size): + with pytest.raises(ValueError, match="must be positive"): + get_packed_seq_align_size(tp_size, cp_size, fp8_enabled=True) + + +def test_unpacked_alignment_rejects_nonpositive_tp_size(): + with pytest.raises(ValueError, match="must be positive"): + get_unpacked_seq_align_size(0, fp8_enabled=True) diff --git a/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_cp.py b/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_cp.py index 6c53135367..7a0144e475 100644 --- a/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_cp.py +++ b/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_cp.py @@ -142,3 +142,47 @@ def test_short_seq_no_crash(self, tp_size, cp_size, real_tokens): assert result_ids.shape[1] % 16 == 0 assert packed_params.max_seqlen_q % (16 * cp_size) == 0 assert packed_params.qkv_format == "thd" + + def test_remove_left_padding_tp1_aligns_only_for_fp8(self): + """TP1 applies 16-token alignment only when FP8 is enabled.""" + from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( + remove_left_padding, + ) + + input_ids = torch.arange(7000).unsqueeze(0) + attention_mask = torch.zeros((1, 7000), dtype=torch.bool) + attention_mask[0, :6541] = True + position_ids = torch.arange(7000).unsqueeze(0) + + with patch("skyrl.backends.skyrl_train.distributed.megatron.megatron_utils.mpu") as mock_mpu: + mock_mpu.get_tensor_model_parallel_world_size.return_value = 1 + mock_mpu.get_context_parallel_world_size.return_value = 1 + + bf16_ids, bf16_mask, _ = remove_left_padding(input_ids, attention_mask, position_ids, fp8_enabled=False) + fp8_ids, fp8_mask, _ = remove_left_padding(input_ids, attention_mask, position_ids, fp8_enabled=True) + + assert bf16_ids.shape == (1, 6541) + assert int(bf16_mask.sum().item()) == 6541 + assert fp8_ids.shape == (1, 6544) + assert int(fp8_mask.sum().item()) == 6541 + + def test_remove_left_padding_tp_gt_1_fp8_uses_local_128_alignment(self): + """TP2 rounds 8,552 tokens to the 8,704-token global alignment.""" + from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( + remove_left_padding, + ) + + input_ids = torch.arange(9000).unsqueeze(0) + attention_mask = torch.zeros((1, 9000), dtype=torch.bool) + attention_mask[0, :8552] = True + position_ids = torch.arange(9000).unsqueeze(0) + + with patch("skyrl.backends.skyrl_train.distributed.megatron.megatron_utils.mpu") as mock_mpu: + mock_mpu.get_tensor_model_parallel_world_size.return_value = 2 + mock_mpu.get_context_parallel_world_size.return_value = 1 + + fp8_ids, fp8_mask, _ = remove_left_padding(input_ids, attention_mask, position_ids, fp8_enabled=True) + + assert fp8_ids.shape == (1, 8704) + assert int(fp8_mask.sum().item()) == 8552 + assert fp8_ids.shape[1] % (128 * 2) == 0 diff --git a/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_multiseq.py b/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_multiseq.py index 27d5e0a920..1ac32ee525 100644 --- a/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_multiseq.py +++ b/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_multiseq.py @@ -185,24 +185,24 @@ def test_multiseq_with_tp_alignment(self): The intra-row offsets read by preprocess must match the collator's row layout, which advances ``row_offset += round_up(s, align_size)`` between sub-seqs. So with sub-seqs of length 3 and - 5 and tp_size=4, the collator places sub-seq 1 at row column 16 - (after the FP8/TP alignment pad gap), NOT row column 3. + 5 and tp_size=4, the collator places sub-seq 1 at row column + ``align_size`` after the FP8/TP alignment gap, not row column 3. """ from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( preprocess_packed_seqs, ) - seq_len = 48 batch_size = 1 - # Two sub-seqs of length 3 and 5; tp_size=4 still pads each to 16 - # for Transformer Engine FP8 compatibility. + # Each sequence uses a global footprint that leaves TP-local inputs + # aligned to 128 tokens. align_size = _get_align_size(tp_size=4, cp_size=1, fp8_enabled=True) + seq_len = 2 * align_size # Row layout mirrors what PackedDataCollator produces: - # row[0:3] = sub-seq 0 tokens - # row[3:16] = alignment pad (zero) - # row[16:21] = sub-seq 1 tokens - # row[21:32] = alignment pad (zero) + # row[0:3] = sub-seq 0 tokens + # row[3:align_size] = alignment pad (zero) + # row[align_size:align_size + 5] = sub-seq 1 tokens + # row[align_size + 5:2*align_size] = alignment pad (zero) input_ids = torch.zeros(batch_size, seq_len, dtype=torch.long) input_ids[0, :3] = torch.tensor([1, 2, 3]) input_ids[0, align_size : align_size + 5] = torch.tensor([10, 11, 12, 13, 14]) @@ -222,11 +222,11 @@ def test_multiseq_with_tp_alignment(self): fp8_enabled=True, ) - assert params.cu_seqlens_q.tolist() == [0, 16, 32] - assert packed.shape == (1, 32) + assert params.cu_seqlens_q.tolist() == [0, align_size, 2 * align_size] + assert packed.shape == (1, 2 * align_size) assert packed[0, :3].tolist() == [1, 2, 3] - assert packed[0, 3:16].tolist() == [0] * 13 - assert packed[0, 16:21].tolist() == [10, 11, 12, 13, 14] + assert packed[0, 3:align_size].tolist() == [0] * (align_size - 3) + assert packed[0, align_size : align_size + 5].tolist() == [10, 11, 12, 13, 14] def test_multiple_bin_rows(self): """Two bin rows, each with two sub-seqs, produce 4+1 cu_seqlens entries.""" diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_worker.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_worker.py index 2b5a1bf571..b201a49f34 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_worker.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_worker.py @@ -38,6 +38,7 @@ # this might be a model specific mbridge issue - see if this persists when we transition to Megatron-Bridge # MOE_MODEL_NAME = "Qwen/Qwen1.5-MoE-A2.7B" MOE_MODEL_NAME = "Qwen/Qwen3-30B-A3B" +TINY_MOE_MODEL_NAME = "eatang/qwen3-moe-tiny-random" def get_test_actor_config(model_name=MODEL_NAME) -> SkyRLTrainConfig: @@ -196,6 +197,89 @@ async def test_megatron_policy_weight_sync( ray.shutdown() +@pytest.mark.asyncio +@pytest.mark.megatron +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="Expert MXFP8 requires SM100 or SM103", +) +async def test_expert_mxfp8_training_and_weight_sync(ray_init_fixture): + cfg = SkyRLTrainConfig.from_cli_overrides( + [ + "trainer.strategy=megatron", + f"trainer.policy.model.path={TINY_MOE_MODEL_NAME}", + "trainer.policy.model.expert_mxfp8.enabled=true", + "trainer.policy.model.expert_mxfp8.persistent=true", + "trainer.policy.megatron_config.ddp_config.fp8_param_gather=true", + "trainer.placement.policy_num_gpus_per_node=1", + "generator.inference_engine.fp8_weight_sync_mode=serialized_mxfp8", + "generator.inference_engine.num_engines=1", + "generator.inference_engine.tensor_parallel_size=1", + "trainer.micro_forward_batch_size_per_gpu=1", + "trainer.micro_train_batch_size_per_gpu=1", + ] + ) + + try: + async with InferenceEngineState.create( + cfg=cfg, + model=TINY_MOE_MODEL_NAME, + use_local=True, + backend="vllm", + sleep_level=2, + ) as engines: + client, pg = engines.client, engines.pg + await client.sleep() + policy = init_worker_with_type( + "policy", + shared_pg=pg, + colocate_all=True, + num_gpus_per_node=1, + cfg=cfg, + ) + ray.get( + policy.async_run_ray_method( + "pass_through", + "init_weight_sync_state", + client, + cfg.generator.inference_engine, + ) + ) + + batch = get_test_training_batch() + batch.metadata["global_step"] = 0 + ray.get(policy.async_run_ray_method("mesh", "forward_backward", batch)) + ray.get(policy.async_run_ray_method("pass_through", "optim_step")) + + await client.wake_up(tags=["weights"]) + for _ in range(2): + ray.get( + policy.async_run_ray_method( + "pass_through", + "broadcast_to_inference_engines", + client, + cfg.generator.inference_engine, + ) + ) + policy.offload_to_cpu() + await client.wake_up(tags=["kv_cache"]) + + tokenizer = AutoTokenizer.from_pretrained(TINY_MOE_MODEL_NAME, trust_remote_code=True) + sampling_params = get_sampling_params_for_backend( + cfg.generator.inference_engine.backend, + cfg.generator.sampling_params, + ) + outputs = await run_inference( + client, + get_test_prompts(TINY_MOE_MODEL_NAME), + sampling_params, + tokenizer=tokenizer, + ) + assert outputs["responses"] + finally: + ray.shutdown() + + @pytest.mark.asyncio @pytest.mark.parametrize( ( diff --git a/tests/backends/skyrl_train/inference_servers/test_batched_moe_fp8_reload.py b/tests/backends/skyrl_train/inference_servers/test_batched_moe_fp8_reload.py new file mode 100644 index 0000000000..5f4881baea --- /dev/null +++ b/tests/backends/skyrl_train/inference_servers/test_batched_moe_fp8_reload.py @@ -0,0 +1,77 @@ +from types import SimpleNamespace + +import torch + +from skyrl.backends.skyrl_train.inference_servers.new_inference_worker_wrap import ( + _load_batched_moe_fp8_tensor, +) +from skyrl.backends.skyrl_train.weight_sync.serialized_fp8 import ( + SKYRL_BATCHED_MOE_FP8_PREFIX, +) + + +def test_batched_moe_tensor_uses_one_full_expert_loader_call(): + calls = [] + + def weight_loader(param, loaded_weight, weight_name, *, shard_id, expert_id, return_success): + calls.append((param, loaded_weight, weight_name, shard_id, expert_id, return_success)) + return True + + weight_loader.supports_moe_loading = True + param = torch.nn.Parameter(torch.empty(3, 8, 4), requires_grad=False) + param.weight_loader = weight_loader + target_name = "model.layers.0.mlp.experts.w13_weight" + loaded_weight = torch.randn(3, 4, 4) + wire_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}model.layers.0.mlp.experts.gate_proj.weight" + + loaded = _load_batched_moe_fp8_tensor( + SimpleNamespace(), + {target_name: param}, + wire_name, + loaded_weight, + ) + + assert loaded + assert len(calls) == 1 + assert calls[0][1] is loaded_weight + assert calls[0][2:] == (target_name, "w1", 0, True) + + +def test_batched_moe_scale_maps_to_fused_scale_parameter(): + calls = [] + + def weight_loader(param, loaded_weight, weight_name, *, shard_id, expert_id, return_success): + calls.append((weight_name, shard_id, tuple(loaded_weight.shape))) + return True + + weight_loader.supports_moe_loading = True + param = torch.nn.Parameter(torch.empty(2, 6, 3), requires_grad=False) + param.weight_loader = weight_loader + target_name = "language_model.model.layers.2.mlp.experts.w13_weight_scale_inv" + loaded_weight = torch.randn(2, 3, 3) + mapper = SimpleNamespace( + apply_list=lambda names: [names[0].replace("model.language_model.", "language_model.model.", 1)] + ) + model = SimpleNamespace(hf_to_vllm_mapper=mapper) + wire_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}model.language_model.layers.2.mlp.experts.up_proj.weight_scale_inv" + + assert _load_batched_moe_fp8_tensor(model, {target_name: param}, wire_name, loaded_weight) + assert calls == [(target_name, "w3", (2, 3, 3))] + + +def test_batched_moe_mxfp8_scale_maps_to_modelopt_scale_parameter(): + calls = [] + + def weight_loader(param, loaded_weight, weight_name, *, shard_id, expert_id, return_success): + calls.append((weight_name, shard_id, loaded_weight.dtype)) + return True + + weight_loader.supports_moe_loading = True + param = torch.nn.Parameter(torch.empty(2, 6, 3, dtype=torch.uint8), requires_grad=False) + param.weight_loader = weight_loader + target_name = "model.layers.2.mlp.experts.w2_weight_scale" + loaded_weight = torch.zeros(2, 3, 3, dtype=torch.uint8) + wire_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}model.layers.2.mlp.experts.down_proj.weight_scale" + + assert _load_batched_moe_fp8_tensor(SimpleNamespace(), {target_name: param}, wire_name, loaded_weight) + assert calls == [(target_name, "w2", torch.uint8)] diff --git a/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py b/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py index ad3372e06f..e24c9854e0 100644 --- a/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py +++ b/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py @@ -1,14 +1,186 @@ """Tests for build_vllm_cli_args on GPU-less hosts.""" +from types import SimpleNamespace + import pytest from skyrl.backends.skyrl_train.inference_servers.utils import ( + _apply_serialized_fp8_weight_sync_defaults, build_vllm_cli_args, resolve_policy_model_name, ) from skyrl.train.config import SkyRLTrainConfig +def test_serialized_fp8_weight_sync_defaults_configure_vllm_checkpoint_fp8(monkeypatch): + import skyrl.backends.skyrl_train.inference_servers.utils as inference_utils + + monkeypatch.setattr(inference_utils, "_serialized_fp8_ignored_layers", lambda _model_path: []) + cfg = SkyRLTrainConfig() + ie_cfg = cfg.generator.inference_engine + ie_cfg.fp8_weight_sync_mode = "serialized_blockwise" + engine_kwargs = {"hf_overrides": {"rope_theta": 10000.0}} + + _apply_serialized_fp8_weight_sync_defaults(ie_cfg, engine_kwargs, model_path="qwen35-test") + + assert engine_kwargs["quantization"] == "fp8" + assert engine_kwargs["load_format"] == "dummy" + assert engine_kwargs["hf_overrides"]["rope_theta"] == 10000.0 + assert engine_kwargs["hf_overrides"]["quantization_config"] == { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + } + + +def test_serialized_mxfp8_weight_sync_configures_modelopt_experts(monkeypatch): + import skyrl.backends.skyrl_train.inference_servers.utils as inference_utils + + monkeypatch.setattr( + inference_utils, + "_serialized_mxfp8_quantization_config", + lambda _model_path: { + "quant_method": "modelopt", + "quant_algo": "MXFP8", + "ignore": ["*.self_attn.*"], + }, + ) + cfg = SkyRLTrainConfig() + ie_cfg = cfg.generator.inference_engine + ie_cfg.fp8_weight_sync_mode = "serialized_mxfp8" + engine_kwargs = {} + + _apply_serialized_fp8_weight_sync_defaults(ie_cfg, engine_kwargs, model_path="qwen3-moe-test") + + assert engine_kwargs == { + "quantization": "modelopt_mxfp8", + "load_format": "dummy", + "hf_overrides": { + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "MXFP8", + "ignore": ["*.self_attn.*"], + } + }, + } + + +def test_serialized_mxfp8_requires_bfloat16(): + cfg = SkyRLTrainConfig() + ie_cfg = cfg.generator.inference_engine + ie_cfg.fp8_weight_sync_mode = "serialized_mxfp8" + ie_cfg.model_dtype = "float16" + + with pytest.raises(ValueError, match="model_dtype='bfloat16'"): + _apply_serialized_fp8_weight_sync_defaults(ie_cfg, {}, model_path="qwen3-moe-test") + + +@pytest.mark.parametrize( + "engine_kwargs", + [ + {"quantization": "awq"}, + {"load_format": "safetensors"}, + {"hf_overrides": {"quantization_config": {"weight_block_size": [64, 128]}}}, + ], +) +def test_serialized_fp8_weight_sync_rejects_conflicting_vllm_settings(engine_kwargs, monkeypatch): + import skyrl.backends.skyrl_train.inference_servers.utils as inference_utils + + monkeypatch.setattr(inference_utils, "_serialized_fp8_ignored_layers", lambda _model_path: []) + cfg = SkyRLTrainConfig() + cfg.generator.inference_engine.fp8_weight_sync_mode = "serialized_blockwise" + + with pytest.raises(ValueError, match="serialized FP8 weight sync"): + _apply_serialized_fp8_weight_sync_defaults( + cfg.generator.inference_engine, + engine_kwargs, + model_path="qwen35-test", + ) + + +@pytest.mark.parametrize( + "engine_kwargs", + [ + {"hf_overrides": []}, + {"hf_overrides": {"quantization_config": []}}, + ], +) +def test_serialized_fp8_weight_sync_rejects_non_mapping_overrides(engine_kwargs): + cfg = SkyRLTrainConfig() + cfg.generator.inference_engine.fp8_weight_sync_mode = "serialized_blockwise" + + with pytest.raises(ValueError, match="must be a dict"): + _apply_serialized_fp8_weight_sync_defaults( + cfg.generator.inference_engine, + engine_kwargs, + model_path="qwen35-test", + ) + + +def test_serialized_fp8_requires_model_path(): + cfg = SkyRLTrainConfig() + cfg.generator.inference_engine.fp8_weight_sync_mode = "serialized_blockwise" + + with pytest.raises(ValueError, match="model path is required"): + _apply_serialized_fp8_weight_sync_defaults(cfg.generator.inference_engine, {}) + + +def test_serialized_fp8_fails_when_model_config_cannot_be_inspected(monkeypatch): + import transformers + + cfg = SkyRLTrainConfig() + cfg.generator.inference_engine.fp8_weight_sync_mode = "serialized_blockwise" + + def fail_config_load(*_args, **_kwargs): + raise OSError("missing config") + + monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", fail_config_load) + with pytest.raises(RuntimeError, match="Could not inspect the model config"): + _apply_serialized_fp8_weight_sync_defaults( + cfg.generator.inference_engine, + {}, + model_path="missing-model", + ) + + +def test_serialized_fp8_rejects_unsupported_model_layout(monkeypatch): + import transformers + + cfg = SkyRLTrainConfig() + cfg.generator.inference_engine.fp8_weight_sync_mode = "serialized_blockwise" + monkeypatch.setattr( + transformers.AutoConfig, + "from_pretrained", + lambda *_args, **_kwargs: SimpleNamespace(model_type="llama"), + ) + + with pytest.raises(ValueError, match="only Qwen3.5"): + _apply_serialized_fp8_weight_sync_defaults( + cfg.generator.inference_engine, + {}, + model_path="unsupported-model", + ) + + +def test_serialized_mxfp8_rejects_unsupported_model_layout(monkeypatch): + import transformers + + cfg = SkyRLTrainConfig() + cfg.generator.inference_engine.fp8_weight_sync_mode = "serialized_mxfp8" + monkeypatch.setattr( + transformers.AutoConfig, + "from_pretrained", + lambda *_args, **_kwargs: SimpleNamespace(model_type="llama"), + ) + + with pytest.raises(ValueError, match="does not support model_type"): + _apply_serialized_fp8_weight_sync_defaults( + cfg.generator.inference_engine, + {}, + model_path="unsupported-model", + ) + + @pytest.mark.vllm def test_build_vllm_cli_args_succeeds_on_gpu_less_host(monkeypatch): import vllm.platforms @@ -25,7 +197,13 @@ def test_build_vllm_cli_args_succeeds_on_gpu_less_host(monkeypatch): cfg = SkyRLTrainConfig() cfg.generator.inference_engine.served_model_name = "served-alias" cfg.generator.inference_engine.engine_init_kwargs = { - "hf_overrides": {"rope_parameters": {"rope_type": "linear", "factor": 2.0, "rope_theta": 10000.0}} + "hf_overrides": { + "rope_parameters": { + "rope_type": "linear", + "factor": 2.0, + "rope_theta": 10000.0, + } + } } args = build_vllm_cli_args(cfg) @@ -33,7 +211,11 @@ def test_build_vllm_cli_args_succeeds_on_gpu_less_host(monkeypatch): assert args.model == cfg.trainer.policy.model.path assert args.served_model_name == ["served-alias"] assert args.tensor_parallel_size == cfg.generator.inference_engine.tensor_parallel_size - assert args.hf_overrides["rope_parameters"] == {"rope_type": "linear", "factor": 2.0, "rope_theta": 10000.0} + assert args.hf_overrides["rope_parameters"] == { + "rope_type": "linear", + "factor": 2.0, + "rope_theta": 10000.0, + } assert vllm.platforms.current_platform.device_type == "cuda" # NOTE: the MTP speculative_config wiring test lives in diff --git a/tests/backends/skyrl_train/test_expert_mxfp8.py b/tests/backends/skyrl_train/test_expert_mxfp8.py new file mode 100644 index 0000000000..a0cddc5f48 --- /dev/null +++ b/tests/backends/skyrl_train/test_expert_mxfp8.py @@ -0,0 +1,77 @@ +from argparse import Namespace +from fnmatch import fnmatch + +import pytest + +from skyrl.backends.skyrl_train.inference_servers.utils import ( + apply_expert_mxfp8_rollout_config, +) +from skyrl.backends.skyrl_train.workers.megatron.expert_mxfp8 import ( + expert_mxfp8_recipe_dict, + is_routed_expert_linear, +) +from skyrl.train.config import SkyRLTrainConfig + + +def _enabled_config() -> SkyRLTrainConfig: + return SkyRLTrainConfig.from_cli_overrides( + [ + "trainer.strategy=megatron", + "trainer.policy.model.expert_mxfp8.enabled=true", + ] + ) + + +def test_expert_recipe_targets_only_routed_experts(): + recipe = expert_mxfp8_recipe_dict() + fc1_pattern = recipe["matchers"]["routed_fc1"]["pattern"] + fc2_pattern = recipe["matchers"]["routed_fc2"]["pattern"] + assert fnmatch("decoder.layers.0.mlp.experts.linear_fc1", fc1_pattern) + assert fnmatch("decoder.layers.0.mlp.experts.linear_fc2", fc2_pattern) + assert not fnmatch("decoder.layers.0.mlp.shared_experts.linear_fc1", fc1_pattern) + assert not fnmatch("decoder.layers.0.mlp.shared_experts.linear_fc2", fc2_pattern) + assert is_routed_expert_linear("decoder.layers.0.mlp.experts.linear_fc1") + assert is_routed_expert_linear("decoder.layers.0.mlp.experts.linear_fc2.base_layer") + assert not is_routed_expert_linear("decoder.layers.0.mlp.shared_experts.linear_fc1") + assert not is_routed_expert_linear("decoder.layers.0.self_attention.linear_qkv") + + +def test_persistent_expert_recipe_enables_fp8_params_only_for_routed_experts(): + default_recipe = expert_mxfp8_recipe_dict() + persistent_recipe = expert_mxfp8_recipe_dict(persistent=True) + + for phase in ("training_recipe", "evaluation_recipe"): + assert "fp8_param" not in default_recipe["configs"]["expert_mxfp8"][phase] + assert persistent_recipe["configs"]["expert_mxfp8"][phase]["fp8_param"] is True + assert "fp8_param" not in persistent_recipe["configs"]["high_precision"][phase] + + +def test_rollout_config_enables_expert_only_mxfp8(): + args = Namespace() + apply_expert_mxfp8_rollout_config(args, _enabled_config(), {}) + assert args.quantization == "online" + assert args.quantization_config == {"moe": "mxfp8"} + + +def test_rollout_config_preserves_serialized_mxfp8_settings(): + cfg = _enabled_config() + cfg.generator.inference_engine.fp8_weight_sync_mode = "serialized_mxfp8" + args = Namespace() + + apply_expert_mxfp8_rollout_config(args, cfg, {"quantization": "modelopt_mxfp8"}) + + assert not hasattr(args, "quantization") + assert not hasattr(args, "quantization_config") + + +def test_rollout_config_rejects_float32(): + cfg = _enabled_config() + cfg.generator.inference_engine.model_dtype = "float32" + with pytest.raises(ValueError, match="model_dtype"): + apply_expert_mxfp8_rollout_config(Namespace(), cfg, {}) + + +@pytest.mark.parametrize("key", ["quantization", "quantization_config"]) +def test_rollout_config_rejects_quantization_overrides(key): + with pytest.raises(ValueError, match="engine_init_kwargs"): + apply_expert_mxfp8_rollout_config(Namespace(), _enabled_config(), {key: "conflict"}) diff --git a/tests/backends/skyrl_train/weight_sync/test_serialized_fp8.py b/tests/backends/skyrl_train/weight_sync/test_serialized_fp8.py new file mode 100644 index 0000000000..e5aca5114c --- /dev/null +++ b/tests/backends/skyrl_train/weight_sync/test_serialized_fp8.py @@ -0,0 +1,392 @@ +from types import SimpleNamespace + +import pytest +import torch + +from skyrl.backends.skyrl_train.weight_sync.serialized_fp8 import ( + MXFP8_1X32, + SERIALIZED_MXFP8, + SKYRL_BATCHED_MOE_FP8_PREFIX, + SerializedFp8Config, + batched_blockwise_cast_to_fp8, + batched_moe_expert_spec, + batched_mxfp8_cast_to_fp8, + blockwise_cast_to_fp8, + get_moe_architecture_spec, + get_qwen35_fp8_ignored_layers, + get_serialized_fp8_quantization_config, + get_serialized_mxfp8_quantization_config, + is_quantizable_weight, + is_quantizable_weight_shape, + iter_serialized_fp8_tensors, + mxfp8_cast_to_fp8, + serialized_fp8_config_for_mode, +) + + +def test_blockwise_cast_to_fp8_emits_weight_and_fp32_scale(): + weight = torch.arange(257 * 129, dtype=torch.float32).reshape(257, 129) / 1000 + + q_weight, scale = blockwise_cast_to_fp8(weight, [128, 128]) + + assert q_weight.shape == weight.shape + assert q_weight.dtype == torch.float8_e4m3fn + assert scale.shape == (3, 2) + assert scale.dtype == torch.float32 + + +def test_blockwise_cast_defaults_to_exact_fp32_scales(): + torch.manual_seed(7) + weight = torch.randn(256, 256, dtype=torch.float32) + + default_weight, default_scale = blockwise_cast_to_fp8(weight, [128, 128]) + exact_weight, exact_scale = blockwise_cast_to_fp8(weight, [128, 128], power_2_scale=False) + + assert torch.equal(default_weight.view(torch.uint8), exact_weight.view(torch.uint8)) + assert torch.equal(default_scale, exact_scale) + + +def test_blockwise_cast_uses_training_amax_epsilon_for_near_zero_blocks(monkeypatch): + monkeypatch.setenv("NVTE_FP8_BLOCK_AMAX_EPSILON", "1e-4") + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1") + config = SerializedFp8Config() + weight = torch.full((128, 128), 1e-8, dtype=torch.float32) + + _, scale = blockwise_cast_to_fp8( + weight, + config.weight_block_size, + config.power_2_scale, + config.amax_epsilon, + ) + + assert scale.item() == pytest.approx(config.amax_epsilon / torch.finfo(torch.float8_e4m3fn).max) + + +def test_blockwise_cast_pow2_scales_match_te_ue8m0_rule(): + # TE's UE8M0 mode rounds dequantization scales up to powers of two. + torch.manual_seed(0) + weight = torch.randn(256, 384, dtype=torch.float32) + + _, pow2_scale = blockwise_cast_to_fp8(weight, [128, 128], power_2_scale=True) + _, exact_scale = blockwise_cast_to_fp8(weight, [128, 128], power_2_scale=False) + + log2 = torch.log2(pow2_scale) + assert torch.allclose(log2, log2.round(), atol=0.0) + expected = torch.pow(2.0, torch.ceil(torch.log2(exact_scale))) + assert torch.allclose(pow2_scale, expected) + assert torch.all(pow2_scale >= exact_scale) + + +def test_serialized_config_power_2_scale_follows_te_env(monkeypatch): + monkeypatch.delenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", raising=False) + assert SerializedFp8Config().power_2_scale is False + + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1") + assert SerializedFp8Config().power_2_scale is False + + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "0") + assert SerializedFp8Config().power_2_scale is True + + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "invalid") + with pytest.raises(ValueError, match="must be '0'.*or '1'"): + SerializedFp8Config() + + +@pytest.mark.parametrize("block_size", [[], [128], [128, 128, 128], [128, 0], [128, 1.5], [True, 128]]) +def test_blockwise_cast_rejects_invalid_block_size(block_size): + with pytest.raises(ValueError, match="exactly two positive integers"): + blockwise_cast_to_fp8(torch.ones((2, 2)), block_size) + + +def test_quantizable_weight_filter_keeps_embeddings_in_target_dtype(): + config = SerializedFp8Config() + linear = torch.ones((256, 256), dtype=torch.bfloat16) + embedding = torch.ones((32000, 256), dtype=torch.bfloat16) + + assert is_quantizable_weight("model.layers.0.mlp.down_proj.weight", linear) + assert is_quantizable_weight("model.layers.0.linear_attn.in_proj_qkv.weight", linear) + assert is_quantizable_weight("model.layers.0.linear_attn.in_proj_z.weight", linear) + assert is_quantizable_weight("model.layers.0.linear_attn.out_proj.weight", linear) + assert not is_quantizable_weight("model.layers.0.linear_attn.conv1d.weight", linear) + assert not is_quantizable_weight("model.layers.0.linear_attn.in_proj_b.weight", linear) + assert not is_quantizable_weight("model.layers.0.linear_attn.in_proj_a.weight", linear) + assert not is_quantizable_weight("model.embed_tokens.weight", embedding) + + tensors = list( + iter_serialized_fp8_tensors( + "model.embed_tokens.weight", + embedding, + torch.bfloat16, + config, + ) + ) + assert [(name, tensor.dtype) for name, tensor in tensors] == [("model.embed_tokens.weight", torch.bfloat16)] + + +def test_vllm_serialized_fp8_quantization_config(): + assert get_serialized_fp8_quantization_config() == { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + } + assert get_serialized_fp8_quantization_config(ignored_layers=["model.layers.0.linear_attn.in_proj_b"]) == { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + "ignored_layers": ["model.layers.0.linear_attn.in_proj_b"], + } + + +def test_vllm_serialized_mxfp8_quantization_config_is_expert_only(): + config = get_serialized_mxfp8_quantization_config() + + assert config["quant_method"] == "modelopt" + assert config["quant_algo"] == "MXFP8" + assert "*.self_attn.*" in config["ignore"] + assert "*.mlp.shared_expert*" in config["ignore"] + assert not any(pattern == "*.mlp.experts*" for pattern in config["ignore"]) + + +def test_serialized_mxfp8_mode_selects_registered_expert_layout(): + config = serialized_fp8_config_for_mode(SERIALIZED_MXFP8, model_type="qwen3_moe") + + assert config.scaling_mode == MXFP8_1X32 + assert config.expert_only is True + assert config.weight_block_size == (1, 32) + assert get_moe_architecture_spec(config.model_type).batched is True + + +def test_serialized_mxfp8_rejects_unknown_model_type(): + with pytest.raises(ValueError, match="does not support model_type"): + serialized_fp8_config_for_mode(SERIALIZED_MXFP8, model_type="llama") + + +def test_qwen35_fp8_ignored_layers_use_linear_attention_layers(): + hf_config = SimpleNamespace( + model_type="qwen3_5_text", + layer_types=[ + "linear_attention", + "full_attention", + "linear_attention", + ], + ) + + assert get_qwen35_fp8_ignored_layers(hf_config) == [ + "model.layers.0.linear_attn.in_proj_b", + "model.layers.0.linear_attn.in_proj_a", + "model.language_model.layers.0.linear_attn.in_proj_b", + "model.language_model.layers.0.linear_attn.in_proj_a", + "model.layers.2.linear_attn.in_proj_b", + "model.layers.2.linear_attn.in_proj_a", + "model.language_model.layers.2.linear_attn.in_proj_b", + "model.language_model.layers.2.linear_attn.in_proj_a", + ] + + +def test_qwen35_ignored_layers_include_only_checkpoint_vision_prefixes(): + hf_config = SimpleNamespace( + model_type="qwen3_5", + text_config=SimpleNamespace(model_type="qwen3_5_text", layer_types=[]), + vision_config=SimpleNamespace(depth=2), + ) + + assert get_qwen35_fp8_ignored_layers(hf_config) == [ + "model.visual.blocks.0.attn.proj", + "model.visual.blocks.1.attn.proj", + ] + + +def test_qwen35_ignored_layers_are_not_inferred_from_unrelated_hybrid_config(): + hf_config = SimpleNamespace(model_type="unrelated_hybrid", layer_types=["linear_attention"]) + + assert get_qwen35_fp8_ignored_layers(hf_config) == [] + + +def test_moe_batched_expert_spec_recognizes_and_splits_gate_up(): + base = "model.language_model.layers.5.mlp.experts" + assert batched_moe_expert_spec(f"{base}.gate_up_proj") == (base, ("gate_proj", "up_proj"), True) + assert batched_moe_expert_spec(f"{base}.down_proj") == (base, ("down_proj",), False) + assert batched_moe_expert_spec("model.layers.5.mlp.gate.weight") is None + assert batched_moe_expert_spec("model.layers.5.self_attn.q_proj.weight") is None + + +def test_moe_shared_expert_quantized_router_and_gate_bf16(): + lin = (256, 256) + assert is_quantizable_weight_shape("model.layers.3.mlp.shared_expert.gate_proj.weight", lin) + assert is_quantizable_weight_shape("model.layers.3.mlp.shared_expert.up_proj.weight", lin) + assert is_quantizable_weight_shape("model.layers.3.mlp.shared_expert.down_proj.weight", lin) + assert not is_quantizable_weight_shape("model.layers.3.mlp.gate.weight", (256, 8)) + assert not is_quantizable_weight_shape("model.layers.3.mlp.shared_expert_gate.weight", (256, 1)) + + +def test_batched_blockwise_cast_matches_independent_expert_casts(): + torch.manual_seed(11) + weight = torch.randn(5, 257, 129, dtype=torch.bfloat16) + + q_batched, scale_batched = batched_blockwise_cast_to_fp8( + weight, + [128, 128], + power_2_scale=True, + expert_batch_size=2, + ) + + for expert_id in range(weight.shape[0]): + q_expected, scale_expected = blockwise_cast_to_fp8( + weight[expert_id], + [128, 128], + power_2_scale=True, + ) + assert torch.equal(q_batched[expert_id].view(torch.uint8), q_expected.view(torch.uint8)) + assert torch.equal(scale_batched[expert_id], scale_expected) + + +def test_mxfp8_cast_emits_row_major_e8m0_scales(): + weight = torch.cat( + ( + torch.full((2, 32), 1.0, dtype=torch.bfloat16), + torch.full((2, 32), 8.0, dtype=torch.bfloat16), + ), + dim=1, + ) + + q_weight, scale = mxfp8_cast_to_fp8(weight) + expected_weight, expected_scale = blockwise_cast_to_fp8(weight, (1, 32), power_2_scale=True) + + assert q_weight.dtype == torch.float8_e4m3fn + assert q_weight.shape == weight.shape + assert scale.dtype == torch.uint8 + assert scale.shape == (2, 2) + assert torch.equal(scale, torch.tensor([[119, 122], [119, 122]], dtype=torch.uint8)) + assert torch.equal(q_weight.view(torch.uint8), expected_weight.view(torch.uint8)) + assert torch.equal(torch.exp2(scale.float() - 127), expected_scale) + + +def test_batched_mxfp8_cast_matches_batched_blockwise_cast(): + torch.manual_seed(17) + weight = torch.randn(5, 96, 64, dtype=torch.bfloat16) + + q_batched, scale_batched = batched_mxfp8_cast_to_fp8(weight, expert_batch_size=2) + q_expected, scale_expected = batched_blockwise_cast_to_fp8( + weight, + (1, 32), + power_2_scale=True, + expert_batch_size=2, + ) + + assert torch.equal(q_batched.view(torch.uint8), q_expected.view(torch.uint8)) + assert torch.equal(torch.exp2(scale_batched.float() - 127), scale_expected) + + +def test_batched_moe_experts_remain_fused_with_pow2_scales(): + num_experts, moe_inter, hidden = 3, 128, 256 + config = SerializedFp8Config(weight_block_size=(128, 128), power_2_scale=True) + base = "model.language_model.layers.7.mlp.experts" + + torch.manual_seed(0) + gate_up = torch.randn(num_experts, 2 * moe_inter, hidden, dtype=torch.bfloat16) + emitted = list(iter_serialized_fp8_tensors(f"{base}.gate_up_proj", gate_up, torch.bfloat16, config)) + by_name = dict(emitted) + assert len(emitted) == 4 + for projection_idx, proj in enumerate(("gate_proj", "up_proj")): + weight_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.{proj}.weight" + w = by_name[weight_name] + s = by_name[f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.{proj}.weight_scale_inv"] + assert w.dtype == torch.float8_e4m3fn and tuple(w.shape) == (num_experts, moe_inter, hidden) + assert s.dtype == torch.float32 and tuple(s.shape) == (num_experts, 1, 2) + log2 = torch.log2(s) + assert torch.allclose(log2, log2.round(), atol=0.0) + for expert_id in range(num_experts): + source = gate_up[expert_id, projection_idx * moe_inter : (projection_idx + 1) * moe_inter] + q_expected, scale_expected = blockwise_cast_to_fp8( + source, + config.weight_block_size, + config.power_2_scale, + ) + assert torch.equal(w[expert_id].view(torch.uint8), q_expected.view(torch.uint8)) + assert torch.equal(s[expert_id], scale_expected) + + down = torch.randn(num_experts, hidden, moe_inter, dtype=torch.bfloat16) + emitted_d = dict(iter_serialized_fp8_tensors(f"{base}.down_proj", down, torch.bfloat16, config)) + down_weight_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.down_proj.weight" + assert set(emitted_d) == { + down_weight_name, + f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.down_proj.weight_scale_inv", + } + assert emitted_d[down_weight_name].shape == down.shape + + +def test_batched_qwen35_moe_experts_serialize_as_mxfp8(): + config = serialized_fp8_config_for_mode(SERIALIZED_MXFP8, model_type="qwen3_5_moe_text") + base = "model.language_model.layers.7.mlp.experts" + gate_up = torch.randn(3, 256, 64, dtype=torch.bfloat16) + + emitted = dict(iter_serialized_fp8_tensors(f"{base}.gate_up_proj", gate_up, torch.bfloat16, config)) + + gate_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.gate_proj.weight" + assert emitted[gate_name].dtype == torch.float8_e4m3fn + assert emitted[f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.gate_proj.weight_scale"].dtype == torch.uint8 + assert emitted[f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.gate_proj.weight_scale"].shape == (3, 128, 2) + + +def test_qwen3_packed_expert_weights_serialize_as_mxfp8_only(): + config = serialized_fp8_config_for_mode(SERIALIZED_MXFP8, model_type="qwen3_moe") + base = "model.layers.2.mlp.experts" + gate_up = torch.randn(3, 64, 64, dtype=torch.bfloat16) + attention = torch.randn(128, 64, dtype=torch.bfloat16) + + expert_tensors = dict( + iter_serialized_fp8_tensors(f"{base}.gate_up_proj", gate_up, torch.bfloat16, config) + ) + attention_tensors = list( + iter_serialized_fp8_tensors( + "model.layers.2.self_attn.q_proj.weight", + attention, + torch.bfloat16, + config, + ) + ) + + assert set(expert_tensors) == { + f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.gate_proj.weight", + f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.gate_proj.weight_scale", + f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.up_proj.weight", + f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.up_proj.weight_scale", + } + assert [(name, tensor.dtype) for name, tensor in attention_tensors] == [ + ("model.layers.2.self_attn.q_proj.weight", torch.bfloat16) + ] + + +def test_serialized_mxfp8_rejects_unregistered_expert_export_name(): + config = serialized_fp8_config_for_mode(SERIALIZED_MXFP8, model_type="qwen3_moe") + + with pytest.raises(ValueError, match="Unsupported routed-expert export tensor"): + list( + iter_serialized_fp8_tensors( + "model.layers.2.mlp.experts.unexpected.weight", + torch.randn(128, 64), + torch.bfloat16, + config, + ) + ) + + +def test_batched_moe_gate_up_rejects_odd_output_dimension(): + name = "model.layers.0.mlp.experts.gate_up_proj" + tensor = torch.randn(2, 255, 128, dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="output dimension must be even"): + list(iter_serialized_fp8_tensors(name, tensor, torch.bfloat16, SerializedFp8Config())) + + +def test_batched_moe_expert_iterator_rejects_non_3d_input(): + name = "model.layers.0.mlp.experts.down_proj" + with pytest.raises(ValueError, match="must be 3D"): + list( + iter_serialized_fp8_tensors( + name, + torch.ones((128, 128)), + torch.bfloat16, + SerializedFp8Config(), + ) + ) diff --git a/tests/backends/skyrl_train/weight_sync/test_transfer_strategies.py b/tests/backends/skyrl_train/weight_sync/test_transfer_strategies.py index a5920462f9..9391cfff4b 100644 --- a/tests/backends/skyrl_train/weight_sync/test_transfer_strategies.py +++ b/tests/backends/skyrl_train/weight_sync/test_transfer_strategies.py @@ -1,4 +1,7 @@ +import asyncio + import pytest +import torch from skyrl.backends.skyrl_train.weight_sync import ( BroadcastInitInfo, @@ -10,6 +13,10 @@ LoraLoadRequest, get_transfer_strategy_cls, ) +from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk +from skyrl.backends.skyrl_train.weight_sync.broadcast_strategy import ( + BroadcastWeightTransferSender, +) from skyrl.train.config import InferenceEngineConfig @@ -55,7 +62,7 @@ def _make_ie_cfg( ) def test_cuda_ipc_create_init_info(self): - """CudaIpcTransferStrategy.create_init_info should create CudaIpcInitInfo with model_dtype_str.""" + """Preserve model-dtype metadata in CUDA IPC initialization.""" ie_cfg = self._make_ie_cfg(model_dtype="torch.float32") init_info = CudaIpcTransferStrategy.create_init_info(ie_cfg) @@ -122,6 +129,89 @@ def test_mismatched_lengths_raises(self): ) +def test_broadcast_sender_preserves_mixed_dtype_logical_chunk(monkeypatch): + """vLLM's byte-packed NCCL path accepts one mixed-dtype logical update.""" + import skyrl.backends.skyrl_train.weight_sync.broadcast_strategy as broadcast_module + + class FakeInferenceClient: + def __init__(self): + self.events = [] + + async def start_weight_update(self, is_checkpoint_format): + self.events.append(("start", is_checkpoint_format)) + + async def finish_weight_update(self): + self.events.append(("finish",)) + + client = FakeInferenceClient() + sender = BroadcastWeightTransferSender( + init_info=BroadcastInitInfo( + master_addr="127.0.0.1", + master_port=12345, + rank_offset=1, + world_size=2, + override_existing_receiver=False, + ), + model_update_group=object(), + inference_client=client, + ) + sent_chunks = [] + + async def record_chunk(chunk): + sent_chunks.append((list(chunk.names), [tensor.dtype for tensor in chunk.tensors])) + + monkeypatch.setattr(sender, "_send_chunk_vllm_native", record_chunk) + monkeypatch.setattr(broadcast_module.torch.distributed, "get_rank", lambda: 0) + monkeypatch.setattr(broadcast_module.torch.distributed, "barrier", lambda: None) + + mixed_chunk = WeightChunk( + names=["w0", "scale", "w1", "norm"], + dtypes=["ignored"] * 4, + shapes=[[4], [1], [8], [2]], + tensors=[ + torch.empty((4,), dtype=torch.float8_e4m3fn), + torch.empty((1,), dtype=torch.float32), + torch.empty((8,), dtype=torch.float8_e4m3fn), + torch.empty((2,), dtype=torch.bfloat16), + ], + ) + + asyncio.run(sender.send_chunks(iter([mixed_chunk]), derive_metadata_from_chunks=True)) + + assert client.events == [("start", True), ("finish",)] + assert sent_chunks == [ + ( + ["w0", "scale", "w1", "norm"], + [torch.float8_e4m3fn, torch.float32, torch.float8_e4m3fn, torch.bfloat16], + ) + ] + + +def test_broadcast_sender_retains_precomputed_metadata_path(monkeypatch): + sender = BroadcastWeightTransferSender( + init_info=BroadcastInitInfo( + master_addr="127.0.0.1", + master_port=12345, + rank_offset=1, + world_size=2, + override_existing_receiver=False, + ), + model_update_group=object(), + inference_client=object(), + ) + calls = [] + + async def record_batched(chunks, weight_metadata): + calls.append((list(chunks), weight_metadata)) + + monkeypatch.setattr(sender, "_send_chunks_vllm_native", record_batched) + metadata = {"names": ["w"], "dtype_names": ["bfloat16"], "shapes": [[1]]} + + asyncio.run(sender.send_chunks(iter([]), weight_metadata=metadata)) + + assert calls == [([], metadata)] + + class TestCudaIpcWeightUpdateRequest: """Tests for CudaIpcWeightUpdateRequest.""" diff --git a/tests/backends/skyrl_train/weight_sync/test_weight_chunk.py b/tests/backends/skyrl_train/weight_sync/test_weight_chunk.py index a66c4f5159..6251d3273e 100644 --- a/tests/backends/skyrl_train/weight_sync/test_weight_chunk.py +++ b/tests/backends/skyrl_train/weight_sync/test_weight_chunk.py @@ -2,6 +2,11 @@ import torch from skyrl.backends.skyrl_train.weight_sync import WeightChunk +from skyrl.backends.skyrl_train.weight_sync.base import ( + cuda_uuid_to_str, + get_weight_chunk_metadata, + iter_single_dtype_chunks, +) class TestWeightChunk: @@ -95,3 +100,50 @@ def test_total_size_bytes_mixed_dtypes(self): ) assert chunk.total_size_bytes == 40 + 20 + + +def test_single_dtype_chunk_partition_preserves_first_seen_dtype_order(): + chunk = WeightChunk( + names=["w0", "s0", "w1", "b0"], + dtypes=["ignored"] * 4, + shapes=[[4], [1], [8], [2]], + tensors=[ + torch.empty((4,), dtype=torch.float8_e4m3fn), + torch.ones((1,), dtype=torch.float32), + torch.empty((8,), dtype=torch.float8_e4m3fn), + torch.ones((2,), dtype=torch.bfloat16), + ], + ) + + chunks = list(iter_single_dtype_chunks(chunk)) + + assert [subchunk.names for subchunk in chunks] == [["w0", "w1"], ["s0"], ["b0"]] + assert [[tensor.dtype for tensor in subchunk.tensors] for subchunk in chunks] == [ + [torch.float8_e4m3fn, torch.float8_e4m3fn], + [torch.float32], + [torch.bfloat16], + ] + + +def test_weight_chunk_metadata_uses_actual_tensor_dtypes(): + chunk = WeightChunk( + names=["w", "scale", "norm"], + dtypes=["stale"] * 3, + shapes=[[999], [999], [999]], + tensors=[ + torch.empty((4,), dtype=torch.float8_e4m3fn), + torch.empty((1,), dtype=torch.float32), + torch.empty((2,), dtype=torch.bfloat16), + ], + ) + + assert get_weight_chunk_metadata(chunk) == { + "names": ["w", "scale", "norm"], + "dtype_names": ["float8_e4m3fn", "float32", "bfloat16"], + "shapes": [[4], [1], [2]], + } + + +@pytest.mark.parametrize("uuid", ["GPU-123", b"GPU-123"]) +def test_cuda_uuid_to_str_normalizes_string_and_bytes(uuid): + assert cuda_uuid_to_str(uuid) == "GPU-123" diff --git a/tests/backends/skyrl_train/workers/megatron/test_bridge_mappings.py b/tests/backends/skyrl_train/workers/megatron/test_bridge_mappings.py new file mode 100644 index 0000000000..fb7d799f5d --- /dev/null +++ b/tests/backends/skyrl_train/workers/megatron/test_bridge_mappings.py @@ -0,0 +1,87 @@ +import sys +from types import ModuleType + +from skyrl.backends.skyrl_train.weight_sync.megatron_bridge_mappings import ( + get_packed_qwen3_moe_conversion_tasks, +) + + +def test_qwen3_weight_sync_replaces_per_expert_mappings_with_packed_mappings(monkeypatch): + class AutoMapping: + def __init__(self, megatron_param, hf_param): + self.megatron_param = megatron_param + self.hf_param = hf_param + + class GatedMLPMapping(AutoMapping): + def __init__(self, megatron_param, gate, up): + super().__init__(megatron_param, {"gate": gate, "up": up}) + + class FusedExpertMapping(AutoMapping): + pass + + class FusedGatedExpertMapping(AutoMapping): + pass + + class MegatronMappingRegistry: + def __init__(self, *mappings): + self.mappings = list(mappings) + + def __iter__(self): + return iter(self.mappings) + + def get_all_mappings(self): + return self.mappings.copy() + + module_names = ( + "megatron", + "megatron.bridge", + "megatron.bridge.models", + "megatron.bridge.models.conversion", + ) + for module_name in module_names: + module = ModuleType(module_name) + module.__path__ = [] + monkeypatch.setitem(sys.modules, module_name, module) + + mapping_registry_module = ModuleType("megatron.bridge.models.conversion.mapping_registry") + mapping_registry_module.MegatronMappingRegistry = MegatronMappingRegistry + monkeypatch.setitem(sys.modules, mapping_registry_module.__name__, mapping_registry_module) + param_mapping_module = ModuleType("megatron.bridge.models.conversion.param_mapping") + param_mapping_module.FusedExpertMapping = FusedExpertMapping + param_mapping_module.FusedGatedExpertMapping = FusedGatedExpertMapping + monkeypatch.setitem(sys.modules, param_mapping_module.__name__, param_mapping_module) + + standard_registry = MegatronMappingRegistry( + AutoMapping("embedding.word_embeddings.weight", "model.embed_tokens.weight"), + GatedMLPMapping( + "decoder.layers.*.mlp.experts.linear_fc1.weight*", + gate="model.layers.*.mlp.experts.*.gate_proj.weight", + up="model.layers.*.mlp.experts.*.up_proj.weight", + ), + AutoMapping( + "decoder.layers.*.mlp.experts.linear_fc2.weight*", + "model.layers.*.mlp.experts.*.down_proj.weight", + ), + ) + + class FakeModelBridge: + def mapping_registry(self): + return standard_registry + + def build_conversion_tasks(self, _hf_pretrained, _model): + return self.mapping_registry().get_all_mappings() + + class FakeAutoBridge: + _model_bridge = FakeModelBridge() + hf_pretrained = object() + + mappings = get_packed_qwen3_moe_conversion_tasks(FakeAutoBridge(), object()) + + assert any(isinstance(mapping, FusedGatedExpertMapping) for mapping in mappings) + assert any(isinstance(mapping, FusedExpertMapping) for mapping in mappings) + assert any(mapping.megatron_param == "embedding.word_embeddings.weight" for mapping in mappings) + assert not any( + isinstance(mapping, GatedMLPMapping) + and not isinstance(mapping, FusedGatedExpertMapping) + for mapping in mappings + ) diff --git a/tests/backends/skyrl_train/workers/megatron/test_fp8_block_amax_epsilon_patch.py b/tests/backends/skyrl_train/workers/megatron/test_fp8_block_amax_epsilon_patch.py new file mode 100644 index 0000000000..13a81aeafe --- /dev/null +++ b/tests/backends/skyrl_train/workers/megatron/test_fp8_block_amax_epsilon_patch.py @@ -0,0 +1,64 @@ +import sys +from dataclasses import dataclass +from types import ModuleType, SimpleNamespace + +import pytest + +from skyrl.backends.skyrl_train.workers.megatron._fp8_block_amax_epsilon_patch import ( + apply_fp8_block_amax_epsilon_patch, +) + + +def _install_fake_te_recipe(monkeypatch): + @dataclass(frozen=True) + class QParams: + power_2_scale: bool + amax_epsilon: float = 0.0 + random_hadamard_transform: bool = False + stochastic_rounding: bool = False + fp4_2d_quantization: bool = False + + class Float8BlockScaling: + fp8_quant_fwd_inp = QParams(power_2_scale=False) + fp8_quant_fwd_weight = QParams(power_2_scale=True) + fp8_quant_bwd_grad = QParams(power_2_scale=True) + + def __init__(self, fp8_format): + self.fp8_format = fp8_format + + transformer_engine = ModuleType("transformer_engine") + transformer_engine.__path__ = [] + common = ModuleType("transformer_engine.common") + common.__path__ = [] + recipe = ModuleType("transformer_engine.common.recipe") + recipe.Float8BlockScaling = Float8BlockScaling + recipe.Format = SimpleNamespace(E4M3="e4m3") + transformer_engine.common = common + common.recipe = recipe + + monkeypatch.setitem(sys.modules, "transformer_engine", transformer_engine) + monkeypatch.setitem(sys.modules, "transformer_engine.common", common) + monkeypatch.setitem(sys.modules, "transformer_engine.common.recipe", recipe) + return Float8BlockScaling + + +def test_fp8_block_amax_patch_updates_all_qparams_and_preserves_scale_modes(monkeypatch): + recipe_cls = _install_fake_te_recipe(monkeypatch) + monkeypatch.setenv("NVTE_FP8_BLOCK_AMAX_EPSILON", "1e-4") + + apply_fp8_block_amax_epsilon_patch() + apply_fp8_block_amax_epsilon_patch() + + assert recipe_cls.fp8_quant_fwd_inp.amax_epsilon == 1e-4 + assert recipe_cls.fp8_quant_fwd_weight.amax_epsilon == 1e-4 + assert recipe_cls.fp8_quant_bwd_grad.amax_epsilon == 1e-4 + assert recipe_cls.fp8_quant_fwd_inp.power_2_scale is False + assert recipe_cls.fp8_quant_fwd_weight.power_2_scale is True + + +@pytest.mark.parametrize("value", ["invalid", "-0.1", "nan", "inf"]) +def test_fp8_block_amax_patch_rejects_invalid_explicit_values(monkeypatch, value): + monkeypatch.setenv("NVTE_FP8_BLOCK_AMAX_EPSILON", value) + + with pytest.raises(ValueError, match="NVTE_FP8_BLOCK_AMAX_EPSILON"): + apply_fp8_block_amax_epsilon_patch() diff --git a/tests/backends/skyrl_train/workers/megatron/test_fp8_param.py b/tests/backends/skyrl_train/workers/megatron/test_fp8_param.py new file mode 100644 index 0000000000..00a9ceab10 --- /dev/null +++ b/tests/backends/skyrl_train/workers/megatron/test_fp8_param.py @@ -0,0 +1,171 @@ +import pytest +import torch + +from skyrl.backends.skyrl_train.workers.megatron.fp8_param import ( + initialize_fp8_param_optimizer_masters, + is_fp8_param_enabled, +) + + +class _FakeOptimizer: + def __init__(self): + self.calls = 0 + self.state_dict = None + + def _copy_model_params_to_main_params(self, state_dict=None): + self.calls += 1 + self.state_dict = state_dict + + +class _Range: + def __init__(self, start: int, end: int): + self.start = start + self.end = end + self.size = end - start + + +class _FakeQuantizedParam: + def __init__(self, values): + self._values = torch.tensor(values, dtype=torch.float32) + + def dequantize(self): + return self._values + + +class _FakeHybridDeviceOptimizer: + def __init__(self, cpu_masters, gpu_masters): + self.cpu_copys_map_gpu_param = dict(zip(cpu_masters, gpu_masters)) + self.param_to_fp32_param = {} + + +class _FakeHybridMegatronOptimizer: + def __init__(self): + self.model = _FakeQuantizedParam([2.0, 4.0, 6.0, 8.0]) + self.unquantized_model = torch.tensor([10.0, 20.0, 30.0, 40.0]) + self.gpu_master = torch.full((2,), -1.0) + self.gpu_master_unquantized = torch.full((2,), -3.0) + self.cpu_master = torch.full((2,), -2.0) + self.cpu_master_unquantized = torch.full((2,), -4.0) + self.optimizer = _FakeHybridDeviceOptimizer( + [self.cpu_master, self.cpu_master_unquantized], + [self.gpu_master, self.gpu_master_unquantized], + ) + self.model_float16_groups = [[self.model, self.unquantized_model]] + self.shard_fp32_from_float16_groups = [[self.gpu_master, self.gpu_master_unquantized]] + self.model_fp32_groups = [] + self.shard_fp32_groups = [] + self.loaded_quantized_model = torch.tensor([100.0, 200.0, 300.0, 400.0]) + self.loaded_unquantized_model = torch.tensor([50.0, 60.0, 70.0, 80.0]) + self.exact_state_dict = {"model": {"ignored": torch.tensor(0.0)}} + + def _is_distopt_quantized_param(self, param): + return param is self.model + + def _get_model_param_range_map(self, param): + assert param is self.model or param is self.unquantized_model + return {"param": _Range(1, 3)} + + def _build_model_param_to_state_dict_param_map(self, state_dict): + assert state_dict is self.exact_state_dict + return { + self.model: self.loaded_quantized_model, + self.unquantized_model: self.loaded_unquantized_model, + } + + def _copy_model_params_to_main_params(self): + raise AssertionError("Hybrid CPU-offload path must seed both master copies directly") + + +def test_fp8_param_enablement_reads_skyrl_transformer_config_mapping(): + assert is_fp8_param_enabled({"fp8_param": True}) + assert not is_fp8_param_enabled({"fp8_param": False}) + assert not is_fp8_param_enabled({"fp8_param": "false"}) + assert not is_fp8_param_enabled({}) + + +def test_fp8_param_master_initialization_reloads_each_chained_optimizer(): + first = _FakeOptimizer() + second = _FakeOptimizer() + chained = type("FakeChainedOptimizer", (), {"chained_optimizers": [first, second]})() + + initialized = initialize_fp8_param_optimizer_masters( + chained, + fp8_param=True, + fp8_param_gather=True, + state_dict={"model": {}}, + ) + + assert initialized == 2 + assert first.calls == 1 + assert second.calls == 1 + + +def test_fp8_param_cpu_offload_uses_exact_checkpoint_state_for_all_masters(): + optimizer = _FakeHybridMegatronOptimizer() + + initialized = initialize_fp8_param_optimizer_masters( + optimizer, + fp8_param=True, + fp8_param_gather=True, + state_dict=optimizer.exact_state_dict, + ) + + assert initialized == 1 + torch.testing.assert_close(optimizer.gpu_master, torch.tensor([200.0, 300.0])) + torch.testing.assert_close(optimizer.cpu_master, torch.tensor([200.0, 300.0])) + torch.testing.assert_close(optimizer.gpu_master_unquantized, torch.tensor([60.0, 70.0])) + torch.testing.assert_close(optimizer.cpu_master_unquantized, torch.tensor([60.0, 70.0])) + + +def test_fp8_param_master_initialization_passes_exact_state_to_normal_optimizer(): + optimizer = _FakeOptimizer() + state_dict = {"model": {"weight": torch.tensor([1.0])}} + + initialized = initialize_fp8_param_optimizer_masters( + optimizer, + fp8_param=True, + fp8_param_gather=True, + state_dict=state_dict, + ) + + assert initialized == 1 + assert optimizer.state_dict is state_dict + + +def test_fp8_param_master_initialization_requires_fp8_param_gather(): + optimizer = _FakeOptimizer() + + with pytest.raises(ValueError, match="fp8_param_gather=true"): + initialize_fp8_param_optimizer_masters( + optimizer, + fp8_param=True, + fp8_param_gather=False, + ) + + assert optimizer.calls == 0 + + +def test_fp8_param_master_initialization_requires_exact_checkpoint_state(): + optimizer = _FakeOptimizer() + + with pytest.raises(ValueError, match="exact unquantized checkpoint state"): + initialize_fp8_param_optimizer_masters( + optimizer, + fp8_param=True, + fp8_param_gather=True, + ) + + assert optimizer.calls == 0 + + +def test_non_persistent_fp8_does_not_touch_optimizer_masters(): + optimizer = _FakeOptimizer() + + initialized = initialize_fp8_param_optimizer_masters( + optimizer, + fp8_param=False, + fp8_param_gather=False, + ) + + assert initialized == 0 + assert optimizer.calls == 0 diff --git a/tests/train/test_config.py b/tests/train/test_config.py index 2838c8fa79..79bce252de 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -17,7 +17,12 @@ _resolve_class_type, build_nested_dataclass, ) -from skyrl.train.utils.utils import validate_cfg, validate_inference_engine_cfg +from skyrl.train.utils import utils as train_utils +from skyrl.train.utils.utils import ( + prepare_runtime_environment, + validate_cfg, + validate_inference_engine_cfg, +) from tests.train.util import example_dummy_config @@ -129,6 +134,11 @@ def test_cli_overrides_empty_args(): assert cfg.trainer.seed == 42 +def test_cli_overrides_fp8_param_gather(): + cfg = SkyRLTrainConfig.from_cli_overrides(["trainer.policy.megatron_config.ddp_config.fp8_param_gather=true"]) + assert cfg.trainer.policy.megatron_config.ddp_config.fp8_param_gather is True + + @pytest.mark.parametrize( ("field_name", "value"), [ @@ -143,6 +153,118 @@ def test_trainer_config_rejects_invalid_vocab_entropy_chunking(field_name, value TrainerConfig(**{field_name: value}) +def test_runtime_env_forwards_te_block_scale_mode(monkeypatch): + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1") + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + + env_vars = prepare_runtime_environment(example_dummy_config()) + + assert env_vars["NVTE_FP8_BLOCK_SCALING_FP32_SCALES"] == "1" + + +def test_runtime_env_supports_fsdp_without_megatron_configs(monkeypatch): + monkeypatch.delenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", raising=False) + monkeypatch.delenv("NVTE_FP8_BLOCK_AMAX_EPSILON", raising=False) + monkeypatch.delenv("VLLM_USE_DEEP_GEMM_E8M0", raising=False) + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + cfg = example_dummy_config() + cfg.trainer.strategy = "fsdp" + cfg.trainer.policy.megatron_config = None + cfg.trainer.ref.megatron_config = None + + env_vars = prepare_runtime_environment(cfg) + + assert "NVTE_FP8_BLOCK_SCALING_FP32_SCALES" not in env_vars + assert "VLLM_USE_DEEP_GEMM_E8M0" not in env_vars + + +def test_serialized_fp8_runtime_defaults_to_fp32_scales(monkeypatch): + monkeypatch.delenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", raising=False) + monkeypatch.delenv("VLLM_USE_DEEP_GEMM_E8M0", raising=False) + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + cfg = example_dummy_config() + cfg.generator.inference_engine.fp8_weight_sync_mode = "serialized_blockwise" + + env_vars = prepare_runtime_environment(cfg) + + assert env_vars["NVTE_FP8_BLOCK_SCALING_FP32_SCALES"] == "1" + assert env_vars["VLLM_USE_DEEP_GEMM_E8M0"] == "0" + + +@pytest.mark.parametrize("mode", ["serialized_blockwise", "serialized_mxfp8"]) +def test_serialized_fp8_weight_sync_requires_megatron(mode): + cfg = _make_validated_test_config() + cfg.trainer.strategy = "fsdp" + cfg.generator.inference_engine.fp8_weight_sync_mode = mode + + with pytest.raises(ValueError, match="requires trainer.strategy='megatron'"): + validate_inference_engine_cfg(cfg) + + +@pytest.mark.parametrize("mode", ["serialized_blockwise", "serialized_mxfp8"]) +def test_serialized_fp8_weight_sync_rejects_adapter_only_megatron_lora(mode): + cfg = _make_validated_test_config() + cfg.trainer.strategy = "megatron" + cfg.generator.inference_engine.fp8_weight_sync_mode = mode + cfg.trainer.policy.model.lora.rank = 8 + cfg.trainer.policy.megatron_config.lora_config.merge_lora = False + + with pytest.raises(ValueError, match="requires full-weight updates"): + validate_inference_engine_cfg(cfg) + + +def test_megatron_fp8_compute_defaults_to_fp32_scales_without_serialized_sync(monkeypatch): + monkeypatch.delenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", raising=False) + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + cfg = example_dummy_config() + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8"] = "hybrid" + + env_vars = prepare_runtime_environment(cfg) + + assert env_vars["NVTE_FP8_BLOCK_SCALING_FP32_SCALES"] == "1" + + +def test_power_2_mode_rejects_persistent_fp8_without_serialized_sync(monkeypatch): + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "0") + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + cfg = example_dummy_config() + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8_param"] = True + + with pytest.raises(ValueError, match="fp8_param=false on Blackwell"): + prepare_runtime_environment(cfg) + + +def test_megatron_validation_requires_fp8_param_gather_for_training(): + cfg = _make_validated_test_config() + cfg.trainer.strategy = "megatron" + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8_param"] = True + cfg.trainer.policy.megatron_config.ddp_config.fp8_param_gather = False + + with pytest.raises(ValueError, match="fp8_param_gather=true"): + train_utils.validate_megatron_cfg(cfg) + + +def test_megatron_validation_allows_inference_only_fp8_param_without_gather(): + cfg = _make_validated_test_config() + cfg.trainer.strategy = "megatron" + cfg.trainer.policy.inference_only_init = True + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8_param"] = True + cfg.trainer.policy.megatron_config.ddp_config.fp8_param_gather = False + + train_utils.validate_megatron_cfg(cfg) + + +def test_serialized_fp8_fp32_scales_reject_vllm_e8m0(monkeypatch): + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1") + monkeypatch.setenv("VLLM_USE_DEEP_GEMM_E8M0", "1") + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + cfg = example_dummy_config() + cfg.generator.inference_engine.fp8_weight_sync_mode = "serialized_blockwise" + + with pytest.raises(ValueError, match="VLLM_USE_DEEP_GEMM_E8M0=0"): + prepare_runtime_environment(cfg) + + def test_cli_overrides_plus_prefix_rejected(): with pytest.raises(ValueError, match="The '\\+' prefix"): SkyRLTrainConfig.from_cli_overrides(["+new_field=value"]) @@ -411,6 +533,110 @@ def test_fake_int4_qat_requires_unmerged_lora_sync(): ) +def test_expert_mxfp8_defaults(): + cfg = SkyRLTrainConfig.from_cli_overrides([]) + mxfp8 = cfg.trainer.policy.model.expert_mxfp8 + assert mxfp8.enabled is False + assert mxfp8.training is True + assert mxfp8.rollout is True + assert mxfp8.persistent is False + + +def test_expert_mxfp8_cli_overrides(): + cfg = SkyRLTrainConfig.from_cli_overrides( + [ + "trainer.strategy=megatron", + "trainer.policy.model.expert_mxfp8.enabled=true", + "trainer.policy.model.expert_mxfp8.rollout=false", + "trainer.policy.model.expert_mxfp8.persistent=true", + "trainer.policy.megatron_config.ddp_config.fp8_param_gather=true", + "trainer.policy.megatron_config.ddp_config.reuse_grad_buf_for_mxfp8_param_ag=true", + "trainer.policy.megatron_config.optimizer_config_kwargs.fp8_recipe=mxfp8", + "trainer.policy.megatron_config.optimizer_config_kwargs.reuse_grad_buf_for_mxfp8_param_ag=true", + ] + ) + assert cfg.trainer.policy.model.expert_mxfp8.enabled is True + assert cfg.trainer.policy.model.expert_mxfp8.rollout is False + assert cfg.trainer.policy.model.expert_mxfp8.persistent is True + assert cfg.trainer.policy.megatron_config.ddp_config.reuse_grad_buf_for_mxfp8_param_ag is True + assert cfg.trainer.policy.megatron_config.optimizer_config_kwargs["fp8_recipe"] == "mxfp8" + assert cfg.trainer.policy.megatron_config.optimizer_config_kwargs["reuse_grad_buf_for_mxfp8_param_ag"] is True + + +def test_persistent_expert_mxfp8_requires_training(): + with pytest.raises(ValueError, match="training=true"): + SkyRLTrainConfig.from_cli_overrides( + [ + "trainer.strategy=megatron", + "trainer.policy.model.expert_mxfp8.enabled=true", + "trainer.policy.model.expert_mxfp8.training=false", + "trainer.policy.model.expert_mxfp8.persistent=true", + "trainer.policy.megatron_config.ddp_config.fp8_param_gather=true", + ] + ) + + +def test_persistent_expert_mxfp8_requires_fp8_param_gather(): + with pytest.raises(ValueError, match="fp8_param_gather=true"): + SkyRLTrainConfig.from_cli_overrides( + [ + "trainer.strategy=megatron", + "trainer.policy.model.expert_mxfp8.enabled=true", + "trainer.policy.model.expert_mxfp8.persistent=true", + ] + ) + + +def test_expert_mxfp8_requires_megatron(): + with pytest.raises(ValueError, match="strategy=megatron"): + SkyRLTrainConfig.from_cli_overrides(["trainer.policy.model.expert_mxfp8.enabled=true"]) + + +def test_expert_mxfp8_rejects_unmerged_lora_rollout(): + with pytest.raises(ValueError, match="unmerged LoRA"): + SkyRLTrainConfig.from_cli_overrides( + [ + "trainer.strategy=megatron", + "trainer.policy.model.lora.rank=32", + "trainer.policy.megatron_config.lora_config.merge_lora=false", + "trainer.policy.model.expert_mxfp8.enabled=true", + ] + ) + + +def test_expert_mxfp8_conflicts_with_fake_int4(): + with pytest.raises(ValueError, match="mutually exclusive"): + SkyRLTrainConfig.from_cli_overrides( + [ + "trainer.strategy=megatron", + "trainer.policy.model.lora.rank=32", + "trainer.policy.megatron_config.lora_config.merge_lora=false", + "trainer.policy.model.fake_int4_qat.enabled=true", + "trainer.policy.model.expert_mxfp8.enabled=true", + "trainer.policy.model.expert_mxfp8.rollout=false", + ] + ) + + +@pytest.mark.parametrize( + "override,match", + [ + ("moe_grouped_gemm=false", "moe_grouped_gemm=true"), + ("moe_router_dtype=bf16", "moe_router_dtype=fp32"), + ("fp8_dot_product_attention=true", "transformer_config_kwargs"), + ], +) +def test_expert_mxfp8_rejects_transformer_overrides(override, match): + with pytest.raises(ValueError, match=match): + SkyRLTrainConfig.from_cli_overrides( + [ + "trainer.strategy=megatron", + "trainer.policy.model.expert_mxfp8.enabled=true", + f"trainer.policy.megatron_config.transformer_config_kwargs.{override}", + ] + ) + + class TestTrainerUseSamplePackingAlias: """`trainer.use_sample_packing` is a deprecated alias for `trainer.remove_microbatch_padding` on the RL entrypoint config (mirrors the ``fsdp2``->``fsdp`` alias).""" diff --git a/tests/train/test_packing_round_trip.py b/tests/train/test_packing_round_trip.py index 3eeb298729..073f5c1717 100644 --- a/tests/train/test_packing_round_trip.py +++ b/tests/train/test_packing_round_trip.py @@ -155,9 +155,8 @@ def test_mbs2_tp4_multisubseq_rows_pack_correctly(self): - ``mbs = 2`` so the THD offset stride matters per micro-batch row. - ``tp_size = 4`` so the TP-alignment padding inside rows kicks in. - Multi-subseq rows ``[(7, 5)]`` and ``[(3, 11)]`` to expose offset - mismatches in both directions (sub-seq 0 length 7 is NOT a multiple - of 4, sub-seq 1 length 5 is also not; row 1's sub-seq 0 length 3 is - shorter than its sub-seq 1 length 11). + mismatches in both directions (none of the lengths are alignment + multiples, and row 1's first sub-seq is shorter than its second). """ from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( preprocess_packed_seqs, @@ -177,17 +176,15 @@ def test_mbs2_tp4_multisubseq_rows_pack_correctly(self): ) # Sanity: collator layout matches what PackedDataCollator does. - # row 0: [101..107] + pad-to-16 + [201..205] + pad-to-16 => 32 slots - # row 1: [301..303] + pad-to-16 + [401..411] + pad-to-16 => 32 slots - assert sequences.shape == (2, 32) + assert sequences.shape == (2, 2 * align_size) assert sequences[0, :7].tolist() == row_0_sub_0 - assert sequences[0, 7:16].tolist() == [0] * 9 - assert sequences[0, 16:21].tolist() == row_0_sub_1 + assert sequences[0, 7:align_size].tolist() == [0] * (align_size - 7) + assert sequences[0, align_size : align_size + 5].tolist() == row_0_sub_1 assert sequences[1, :3].tolist() == row_1_sub_0 - assert sequences[1, 3:16].tolist() == [0] * 13 - assert sequences[1, 16:27].tolist() == row_1_sub_1 + assert sequences[1, 3:align_size].tolist() == [0] * (align_size - 3) + assert sequences[1, align_size : align_size + 11].tolist() == row_1_sub_1 # attention_mask is True ONLY at valid slots (NOT at TP-alignment gaps). - assert attention_mask[0, 7:16].any().item() is False + assert attention_mask[0, 7:align_size].any().item() is False assert attention_mask[0].sum().item() == 7 + 5 assert attention_mask[1].sum().item() == 3 + 11 @@ -207,19 +204,14 @@ def test_mbs2_tp4_multisubseq_rows_pack_correctly(self): ) # cu_seqlens_q (== cu_seqlens_q_padded for THD) enumerates 4 sub-seqs: - assert params.cu_seqlens_q.tolist() == [0, 16, 32, 48, 64] - # Packed slab is 64 tokens. Verify each sub-seq's *valid* tokens were - # read from the *correct intra-row offset* — i.e. preprocess respected - # the TP-alignment gap that the collator inserted. - assert packed.shape == (1, 64) - assert packed[0, :7].tolist() == row_0_sub_0 # row 0 sub 0 - assert packed[0, 7:16].tolist() == [0] * 9 # alignment pad - assert packed[0, 16:21].tolist() == row_0_sub_1 # row 0 sub 1 - assert packed[0, 21:32].tolist() == [0] * 11 # alignment pad - assert packed[0, 32:35].tolist() == row_1_sub_0 # row 1 sub 0 - assert packed[0, 35:48].tolist() == [0] * 13 # alignment pad - assert packed[0, 48:59].tolist() == row_1_sub_1 # row 1 sub 1 - assert packed[0, 59:64].tolist() == [0] * 5 # alignment pad + assert params.cu_seqlens_q.tolist() == [i * align_size for i in range(5)] + # Verify valid tokens are read from offsets produced by alignment padding. + assert packed.shape == (1, 4 * align_size) + assert packed[0, :7].tolist() == row_0_sub_0 + assert packed[0, 7:align_size].tolist() == [0] * (align_size - 7) + assert packed[0, align_size : align_size + 5].tolist() == row_0_sub_1 + assert packed[0, 2 * align_size : 2 * align_size + 3].tolist() == row_1_sub_0 + assert packed[0, 3 * align_size : 3 * align_size + 11].tolist() == row_1_sub_1 def test_mbs2_tp1_singlesubseq_rows_match_legacy_preprocess_path(self): """No regression in the legacy path: when each row has 1 sub-seq and tp_size=1.""" diff --git a/tests/train/test_sft_packing_collate.py b/tests/train/test_sft_packing_collate.py index b2a1a2e26b..7a5fbe1787 100644 --- a/tests/train/test_sft_packing_collate.py +++ b/tests/train/test_sft_packing_collate.py @@ -212,17 +212,18 @@ def test_tp_alignment_pads_each_sub_seq(self): subseq_lengths = batch["sub_seq_lengths"][0].tolist() assert sum(subseq_lengths) == 12 # raw, un-padded - def test_fp8_tp_alignment_pads_each_sub_seq_to_16(self): - """When FP8 is enabled, TP-only packed subseqs are also 16-aligned.""" + def test_fp8_tp_alignment_cost_prevents_overpacking(self): + """Packing uses each sequence's aligned FP8/TP footprint.""" collator = _make_collator(num_gpus=4, batch_size=2, max_length=128, tp=4, fp8="hybrid") examples = [ _make_example(7, 3), _make_example(5, 3), ] batch = collator(examples, batch_size=2) - assert batch["sequences"].shape[1] >= 32 - subseq_lengths = batch["sub_seq_lengths"][0].tolist() - assert sum(subseq_lengths) == 12 + # TP4 gives each sequence a 512-token footprint, so the 128-token budget + # expands to one footprint without packing both sequences together. + assert batch["sequences"].shape == (2, 512) + assert sorted(lengths.tolist() for lengths in batch["sub_seq_lengths"]) == [[5], [7]] def test_bf16_cp_alignment_does_not_apply_fp8_padding(self): """With CP>1 and BF16, keep only TP/CP layout padding.""" @@ -256,6 +257,38 @@ def test_fp8_cp_alignment_pads_each_sub_seq(self): padded = ((s + 31) // 32) * 32 assert (padded // 2) % 16 == 0 + def test_fp8_alignment_is_conditional(self): + examples = [ + _make_example(3, 1, base_token=100), + _make_example(3, 1, base_token=200), + ] + + bf16 = _make_collator(num_gpus=1, batch_size=2, max_length=64, fp8=False) + bf16_batch = bf16(examples, batch_size=2) + assert bf16_batch["sequences"].shape[1] == 6 + bf16_chunks = [bf16_batch["sequences"][0, :3].tolist(), bf16_batch["sequences"][0, 3:6].tolist()] + assert sorted(bf16_chunks) == [[100, 101, 102], [200, 201, 202]] + + fp8 = _make_collator(num_gpus=1, batch_size=2, max_length=64, fp8=True) + fp8_batch = fp8(examples, batch_size=2) + assert fp8_batch["sequences"].shape[1] == 32 + fp8_chunks = [fp8_batch["sequences"][0, :3].tolist(), fp8_batch["sequences"][0, 16:19].tolist()] + assert sorted(fp8_chunks) == [[100, 101, 102], [200, 201, 202]] + + def test_tp_gt_1_fp8_alignment_uses_local_128(self): + """TP>1 FP8 row layout must satisfy TE blockwise input all-gather.""" + examples = [ + _make_example(3, 1, base_token=100), + _make_example(3, 1, base_token=200), + ] + + fp8 = _make_collator(num_gpus=2, batch_size=2, max_length=512, tp=2, fp8=True) + fp8_batch = fp8(examples, batch_size=2) + + assert fp8_batch["sequences"].shape[1] == 512 + fp8_chunks = [fp8_batch["sequences"][0, :3].tolist(), fp8_batch["sequences"][0, 256:259].tolist()] + assert sorted(fp8_chunks) == [[100, 101, 102], [200, 201, 202]] + def test_eval_path_falls_back_to_super(self): """When batch_size != self.sft_cfg.batch_size (eval), no packing happens.""" collator = _make_collator(num_gpus=1, batch_size=4, max_length=64)