From 1d023d083fe33d5b9e5182a9e8b25c03c0ade27f Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 03:30:19 +0800 Subject: [PATCH 001/156] feat: harden qwen hybrid lora distributed training --- Cargo.lock | 1 + crates/rustrain-core/src/runtime.rs | 133 +- crates/rustrain-ipc/src/command.rs | 19 +- crates/rustrain-qwen3-6/build.rs | 240 ++- crates/rustrain-qwen3-6/kernels/delta_rule.cu | 23 +- .../rustrain-qwen3-6/kernels/delta_rule.cuh | 181 +- .../kernels/qwen3_6_kernels.cpp | 1778 +++++++++++++---- crates/rustrain-qwen3-6/src/kernel.rs | 490 +++-- crates/rustrain-qwen3-6/src/lora.rs | 892 ++++++++- crates/rustrain-qwen3-6/src/session.rs | 145 +- crates/rustrain-qwen3-6/tests/integration.rs | 217 +- .../tests/native_ep_smoke.cpp | 123 ++ .../rustrain-qwen3-6/tests/native_smoke.cpp | 390 ++++ crates/rustrain-server/Cargo.toml | 3 + crates/rustrain-server/proto/train.proto | 10 + crates/rustrain-server/src/api.rs | 197 +- crates/rustrain-server/src/checkpoint.rs | 158 +- crates/rustrain-server/src/ep.rs | 34 +- crates/rustrain-server/src/grpc.rs | 44 +- crates/rustrain-server/src/session.rs | 383 +++- src/main.rs | 116 +- 21 files changed, 4621 insertions(+), 956 deletions(-) create mode 100644 crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp create mode 100644 crates/rustrain-qwen3-6/tests/native_smoke.cpp diff --git a/Cargo.lock b/Cargo.lock index 905235f4..52308319 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2504,6 +2504,7 @@ dependencies = [ "serde", "serde_json", "tch", + "tempfile", "tokenizers", "tokio", "tokio-stream", diff --git a/crates/rustrain-core/src/runtime.rs b/crates/rustrain-core/src/runtime.rs index fb78bf86..77aa09d1 100644 --- a/crates/rustrain-core/src/runtime.rs +++ b/crates/rustrain-core/src/runtime.rs @@ -3,11 +3,11 @@ use std::{ path::{Path, PathBuf}, }; -use anyhow::{anyhow, Context, Result}; +use anyhow::{Context, Result, anyhow}; use chrono::Local; use regex::Regex; use serde::{Deserialize, Serialize}; -use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter}; +use tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt}; use crate::backend::BackendKind; @@ -489,6 +489,11 @@ pub fn validate_config(config: &Config) -> Result<()> { && config.model.architecture == "qwen_trainable_session"; let is_qwen_lora_sft = matches!(config.train.backend, BackendKind::Tch) && config.model.architecture == "qwen_lora_sft"; + let is_qwen3_hybrid_lora_sft = matches!(config.train.backend, BackendKind::Tch) + && matches!( + config.model.architecture.as_str(), + "qwen3_5_lora_sft" | "qwen3_5_lora_sft_ep" | "qwen3_6_lora_sft" | "qwen3_6_lora_sft_ep" + ); let is_tch_moe_ep_session = matches!(config.train.backend, BackendKind::Tch) && config.model.architecture == "tch_moe_ep_session"; let is_v4_tp_rank = matches!(config.train.backend, BackendKind::Tch) @@ -503,8 +508,11 @@ pub fn validate_config(config: &Config) -> Result<()> { && config.model.architecture == "deepseek_v4_lora_sft_ep"; let is_glm5_lora_sft_ep = matches!(config.train.backend, BackendKind::Tch) && config.model.architecture == "glm5_lora_sft_ep"; - let is_qwen3_6_lora_sft_ep = matches!(config.train.backend, BackendKind::Tch) - && config.model.architecture == "qwen3_6_lora_sft_ep"; + let is_qwen3_hybrid_lora_sft_ep = matches!(config.train.backend, BackendKind::Tch) + && matches!( + config.model.architecture.as_str(), + "qwen3_5_lora_sft_ep" | "qwen3_6_lora_sft_ep" + ); let is_v4_tp_ep_train = matches!(config.train.backend, BackendKind::Tch) && config.model.architecture == "deepseek_v4_tp_ep_train"; let is_v3_tp_rank = matches!(config.train.backend, BackendKind::Tch) @@ -531,7 +539,12 @@ pub fn validate_config(config: &Config) -> Result<()> { && !((is_v4_tp_rank || is_v3_tp_rank || is_v4_tp_train) && name == "tensor_model_parallel_size" && parallel.data_parallel_size == 1) - && !((is_v4_ep_rank || is_v3_ep_rank || is_v4_ep_train || is_v4_lora_sft_ep || is_glm5_lora_sft_ep || is_qwen3_6_lora_sft_ep) + && !((is_v4_ep_rank + || is_v3_ep_rank + || is_v4_ep_train + || is_v4_lora_sft_ep + || is_glm5_lora_sft_ep + || is_qwen3_hybrid_lora_sft_ep) && name == "expert_model_parallel_size" && parallel.data_parallel_size == 1) && !(is_glm5_lora_sft_ep @@ -545,7 +558,7 @@ pub fn validate_config(config: &Config) -> Result<()> { } } - if is_qwen_trainable_session || is_qwen_lora_sft { + if is_qwen_trainable_session || is_qwen_lora_sft || is_qwen3_hybrid_lora_sft { if !matches!(config.train.device, Device::Cuda) { return Err(anyhow!( "{} requires device = \"cuda\"", @@ -580,7 +593,7 @@ pub fn validate_config(config: &Config) -> Result<()> { } } - if is_qwen_lora_sft { + if is_qwen_lora_sft || is_qwen3_hybrid_lora_sft { let lora = config .lora .as_ref() @@ -606,20 +619,44 @@ pub fn validate_config(config: &Config) -> Result<()> { if lora.target_modules.is_empty() { return Err(anyhow!("lora.target_modules must not be empty")); } - let supported_lora_modules = [ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - ]; + let supported_lora_modules: &[&str] = if is_qwen3_hybrid_lora_sft { + &[ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "in_proj_qkv", + "in_proj_z", + "in_proj_a", + "in_proj_b", + "out_proj", + "gate_proj", + "up_proj", + "down_proj", + "shared_gate_proj", + "shared_up_proj", + "shared_down_proj", + "experts_gate_up_proj", + "experts_down_proj", + ] + } else { + &[ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ] + }; for module in &lora.target_modules { if !supported_lora_modules.contains(&module.as_str()) { return Err(anyhow!( - "qwen_lora_sft unsupported lora.target_modules entry {}; supported: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj", - module + "{} unsupported lora.target_modules entry {}; supported: {}", + config.model.architecture, + module, + supported_lora_modules.join(", ") )); } } @@ -630,17 +667,22 @@ pub fn validate_config(config: &Config) -> Result<()> { .any(|module| !unique_modules.insert(module)) { return Err(anyhow!( - "qwen_lora_sft lora.target_modules must not contain duplicates" + "{} lora.target_modules must not contain duplicates", + config.model.architecture )); } if config.train.micro_batch_size == 0 { - return Err(anyhow!("qwen_lora_sft requires micro_batch_size > 0")); + return Err(anyhow!( + "{} requires micro_batch_size > 0", + config.model.architecture + )); } let expected_global_batch_size = config.train.micro_batch_size * config.train.gradient_accumulation_steps; if config.train.global_batch_size != expected_global_batch_size { return Err(anyhow!( - "qwen_lora_sft requires global_batch_size = micro_batch_size * gradient_accumulation_steps" + "{} requires global_batch_size = micro_batch_size * gradient_accumulation_steps", + config.model.architecture )); } } else if config.train.micro_batch_size != 1 || config.train.global_batch_size != 1 { @@ -1272,6 +1314,23 @@ mod tests { )); } + #[test] + fn qwen36_lora_sft_accepts_native_projection_targets() { + let mut config = qwen_lora_sft_config(); + config.model.name = "qwen3_6_test".to_string(); + config.lora.as_mut().unwrap().target_modules = vec![ + "q_proj".to_string(), + "in_proj_qkv".to_string(), + "in_proj_z".to_string(), + "out_proj".to_string(), + ]; + + config.model.architecture = "qwen3_6_lora_sft".to_string(); + validate_config(&config).expect("native Qwen3.6 LoRA targets should validate"); + config.model.architecture = "qwen3_5_lora_sft".to_string(); + validate_config(&config).expect("native Qwen3.5 LoRA targets should validate"); + } + #[test] fn data_max_samples_must_be_positive_when_set() { let mut config = qwen_lora_sft_config(); @@ -1279,9 +1338,11 @@ mod tests { let error = validate_config(&config).expect_err("zero max_samples should fail"); - assert!(error - .to_string() - .contains("data.max_samples must be greater than zero")); + assert!( + error + .to_string() + .contains("data.max_samples must be greater than zero") + ); } #[test] @@ -1300,9 +1361,11 @@ mod tests { let error = validate_config(&config).expect_err("invalid regex should fail"); - assert!(error - .to_string() - .contains("data.field_regex_replacements invalid regex pattern")); + assert!( + error + .to_string() + .contains("data.field_regex_replacements invalid regex pattern") + ); } #[test] @@ -1322,9 +1385,11 @@ mod tests { let error = validate_config(&config).expect_err("invalid regex should fail"); - assert!(error - .to_string() - .contains("data.field_transforms invalid regex_replace pattern")); + assert!( + error + .to_string() + .contains("data.field_transforms invalid regex_replace pattern") + ); } #[test] @@ -1342,9 +1407,11 @@ mod tests { let error = validate_config(&config).expect_err("invalid regex should fail"); - assert!(error - .to_string() - .contains("data field regex filter invalid regex pattern")); + assert!( + error + .to_string() + .contains("data field regex filter invalid regex pattern") + ); } fn qwen_lora_sft_config() -> Config { diff --git a/crates/rustrain-ipc/src/command.rs b/crates/rustrain-ipc/src/command.rs index 2fa3b202..fd4d8e3e 100644 --- a/crates/rustrain-ipc/src/command.rs +++ b/crates/rustrain-ipc/src/command.rs @@ -3,8 +3,12 @@ use serde::{Deserialize, Serialize}; /// Commands that the HTTP server can send to workers. #[derive(Debug, Clone, Serialize, Deserialize)] pub enum EpCommand { - CreateSession { session_id: String }, - DeleteSession { session_id: String }, + CreateSession { + session_id: String, + }, + DeleteSession { + session_id: String, + }, LoadModel { session_id: String, model_path: String, @@ -18,7 +22,7 @@ pub enum EpCommand { InitLora { session_id: String, rank: i64, - alpha: i64, + alpha: f64, target_layers: Vec, target_modules: Vec, lr: f64, @@ -45,7 +49,9 @@ pub enum EpCommand { session_id: String, adapter_id: i64, }, - ListLora { session_id: String }, + ListLora { + session_id: String, + }, TrainStep { session_id: String, input_ids: Vec, @@ -72,8 +78,11 @@ pub enum EpCommand { ExportAdapter { session_id: String, path: String, + adapter_id: Option, + }, + Status { + session_id: String, }, - Status { session_id: String }, Shutdown, } diff --git a/crates/rustrain-qwen3-6/build.rs b/crates/rustrain-qwen3-6/build.rs index 243491cf..bcc25c7e 100644 --- a/crates/rustrain-qwen3-6/build.rs +++ b/crates/rustrain-qwen3-6/build.rs @@ -1,18 +1,6 @@ // build.rs — Compile C++ Qwen3.6 kernels and link against libtorch. -use std::path::PathBuf; use std::process::Command; -fn which(cmd: &str) -> Option { - std::process::Command::new("which") - .arg(cmd) - .output() - .ok() - .filter(|o| o.status.success()) - .and_then(|o| String::from_utf8(o.stdout).ok()) - .map(|s| s.trim().to_string()) - .filter(|s| !s.is_empty()) -} - /// Detect PyTorch's _GLIBCXX_USE_CXX11_ABI setting by running Python. /// Returns "1" or "0". Defaults to "1" if detection fails. fn detect_cxx11_abi() -> String { @@ -23,7 +11,10 @@ fn detect_cxx11_abi() -> String { // Try python3 -c 'import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))' for py in &["python3", "python"] { if let Ok(out) = std::process::Command::new(py) - .args(["-c", "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))"]) + .args([ + "-c", + "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))", + ]) .output() { if out.status.success() { @@ -42,7 +33,12 @@ fn detect_cxx11_abi() -> String { fn main() { println!("cargo:rerun-if-env-changed=TORCH_INCLUDE_PATH"); println!("cargo:rerun-if-env-changed=TORCH_LIB_PATH"); + println!("cargo:rerun-if-env-changed=NCCL_INCLUDE_PATH"); + println!("cargo:rerun-if-env-changed=NCCL_LIB_PATH"); println!("cargo:rerun-if-changed=kernels/qwen3_6_kernels.cpp"); + println!("cargo:rerun-if-changed=kernels/delta_rule.cu"); + println!("cargo:rerun-if-changed=kernels/delta_rule.cuh"); + println!("cargo:rerun-if-changed=kernels/fused_kernels.cu"); println!("cargo:rerun-if-changed=build.rs"); let cxx11_abi = detect_cxx11_abi(); @@ -93,53 +89,130 @@ fn main() { } }; + let torch_package_dir = std::path::Path::new(&torch_lib) + .parent() + .and_then(std::path::Path::parent); + let sibling_nccl = torch_package_dir.map(|dir| dir.join("nvidia/nccl")); + let nccl_include = std::env::var("NCCL_INCLUDE_PATH").unwrap_or_else(|_| { + let mut candidates = Vec::new(); + if let Some(root) = &sibling_nccl { + candidates.push(root.join("include")); + } + candidates.extend([ + std::path::PathBuf::from( + "/share/code/nolanho/mint-runtime-py31213/host-venv/lib/python3.12/site-packages/nvidia/nccl/include", + ), + std::path::PathBuf::from( + "/usr/local/lib/python3.13/dist-packages/nvidia/nccl/include", + ), + std::path::PathBuf::from("/usr/include"), + ]); + candidates + .into_iter() + .find(|path| path.join("nccl.h").exists()) + .unwrap_or_else(|| std::path::PathBuf::from("/usr/include")) + .display() + .to_string() + }); + let nccl_lib = std::env::var("NCCL_LIB_PATH").unwrap_or_else(|_| { + let mut candidates = Vec::new(); + if let Some(root) = &sibling_nccl { + candidates.push(root.join("lib")); + } + candidates.extend([ + std::path::PathBuf::from( + "/share/code/nolanho/mint-runtime-py31213/host-venv/lib/python3.12/site-packages/nvidia/nccl/lib", + ), + std::path::PathBuf::from( + "/usr/local/lib/python3.13/dist-packages/nvidia/nccl/lib", + ), + std::path::PathBuf::from("/usr/lib/x86_64-linux-gnu"), + ]); + candidates + .into_iter() + .find(|path| { + path.join("libnccl.so").exists() || path.join("libnccl.so.2").exists() + }) + .unwrap_or_else(|| std::path::PathBuf::from("/usr/lib/x86_64-linux-gnu")) + .display() + .to_string() + }); + let nccl_link = if std::path::Path::new(&nccl_lib).join("libnccl.so").exists() { + "-lnccl" + } else { + "-l:libnccl.so.2" + }; + let out_dir = std::env::var("OUT_DIR").unwrap_or_else(|_| "target/debug".to_string()); let kernel_src = "kernels/qwen3_6_kernels.cpp"; let output_lib = format!("{out_dir}/libqwen36_kernels.so"); + // Never leave a previously built kernel at the current OUT_DIR after a + // failed rebuild. Runtime loading must not silently pick up stale code. + let _ = std::fs::remove_file(&output_lib); println!("cargo:warning=Compiling Qwen3.6 kernels: include={torch_include} lib={torch_lib}"); let cuda_inc = std::env::var("CUDA_INCLUDE_PATH").unwrap_or_else(|_| { let candidates = [ "/share/code/nolanho/pydeps/lora-research/nvidia/cu13/include", + "/usr/local/cuda-13/include", "/usr/local/cuda-13.0/include", "/usr/local/cuda/include", ]; for c in &candidates { - if std::path::Path::new(&format!("{c}/cuda_runtime_api.h")).exists() { + if std::path::Path::new(&format!("{c}/cuda_runtime_api.h")).exists() + && std::path::Path::new(&format!("{c}/crt/host_defines.h")).exists() + { return c.to_string(); } } "/usr/local/cuda/include".to_string() }); - let status = Command::new("g++") + let cpp_ok = Command::new("g++") .args([ - "-shared", "-fPIC", "-std=c++17", "-O2", - &cxx11_flag, - "-fvisibility=default", - "-o", &output_lib, kernel_src, - &format!("-I{torch_include}"), - &format!("-I{torch_include}/ATen"), - &format!("-I{torch_include}/c10"), - &format!("-I{torch_include}/caffe2"), - &format!("-I{cuda_inc}"), + "-shared".to_string(), + "-fPIC".to_string(), + "-std=c++17".to_string(), + "-O2".to_string(), + cxx11_flag.clone(), + "-fvisibility=default".to_string(), + "-o".to_string(), + output_lib.clone(), + kernel_src.to_string(), + format!("-I{torch_include}"), + format!("-I{torch_include}/ATen"), + format!("-I{torch_include}/c10"), + format!("-I{torch_include}/caffe2"), + format!("-I{cuda_inc}"), + format!("-I{nccl_include}"), + format!("-L{torch_lib}"), + format!("-Wl,-rpath,{torch_lib}"), + format!("-L{nccl_lib}"), + format!("-Wl,-rpath,{nccl_lib}"), + "-Wl,--no-as-needed".to_string(), + "-Wl,--export-dynamic".to_string(), + "-ltorch".to_string(), + "-ltorch_cuda".to_string(), + "-ltorch_cpu".to_string(), + "-lc10".to_string(), + "-lc10_cuda".to_string(), + nccl_link.to_string(), ]) - .args([ - &format!("-L{torch_lib}"), - &format!("-Wl,-rpath,{torch_lib}"), - "-Wl,--no-as-needed", - "-Wl,--export-dynamic", - "-fvisibility=default", - "-ltorch", "-ltorch_cuda", "-ltorch_cpu", "-lc10", "-lc10_cuda", - "-lnccl", - ]) - .status(); + .status() + .map(|s| s.success()) + .unwrap_or(false); // ── Compile CUDA kernels (.cu files) with nvcc ── let nvcc_path = { + let cuda_from_include = std::path::Path::new(&cuda_inc) + .parent() + .map(|home| home.join("bin/nvcc").display().to_string()); let candidates = [ - std::env::var("CUDA_HOME").ok().map(|h| format!("{h}/bin/nvcc")), + std::env::var("CUDA_HOME") + .ok() + .map(|h| format!("{h}/bin/nvcc")), + cuda_from_include, Some("nvcc".to_string()), ]; let mut found = String::new(); @@ -149,7 +222,11 @@ fn main() { found = c.clone(); break; } - if std::process::Command::new(c).arg("--version").output().is_ok() { + if std::process::Command::new(c) + .arg("--version") + .output() + .is_ok() + { found = c.clone(); break; } @@ -158,7 +235,7 @@ fn main() { found }; - if !nvcc_path.is_empty() { + let build_ok = if !nvcc_path.is_empty() { // CUDA files needing nvcc: delta_rule.cu (has __global__) + fused_kernels.cu (hand-written CUDA) let cu_files_nvcc = ["kernels/delta_rule.cu", "kernels/fused_kernels.cu"]; @@ -166,20 +243,29 @@ fn main() { // Compile CUDA kernels with nvcc for cu_file in &cu_files_nvcc { - if !std::path::Path::new(cu_file).exists() { continue; } - let obj_file = format!("{out_dir}/{}.o", - cu_file.replace("kernels/", "").replace(".cu", "")); + if !std::path::Path::new(cu_file).exists() { + continue; + } + let obj_file = format!( + "{out_dir}/{}.o", + cu_file.replace("kernels/", "").replace(".cu", "") + ); let cu_status = Command::new(&nvcc_path) .args([ - "-c", cu_file, "-o", &obj_file, - "-O2", "-std=c++17", + "-c", + cu_file, + "-o", + &obj_file, + "-O2", + "-std=c++17", &cxx11_flag, &format!("-I{torch_include}"), &format!("-I{torch_include}/ATen"), &format!("-I{torch_include}/c10"), &format!("-I{torch_include}/caffe2"), &format!("-I{cuda_inc}"), - "-Xcompiler", "-fPIC", + "-Xcompiler", + "-fPIC", ]) .status(); if cu_status.map(|s| s.success()).unwrap_or(false) { @@ -187,46 +273,68 @@ fn main() { } } - // Re-link the .so with all .o files - if !obj_files.is_empty() { + // Re-link only when every required CUDA translation unit compiled. + if obj_files.len() == cu_files_nvcc.len() { let mut link_args = vec![ - "-shared".to_string(), "-fPIC".to_string(), "-std=c++17".to_string(), "-O2".to_string(), + "-shared".to_string(), + "-fPIC".to_string(), + "-std=c++17".to_string(), + "-O2".to_string(), cxx11_flag.clone(), - "-o".to_string(), output_lib.clone(), kernel_src.to_string(), + "-o".to_string(), + output_lib.clone(), + kernel_src.to_string(), format!("-I{torch_include}"), format!("-I{torch_include}/ATen"), format!("-I{torch_include}/c10"), format!("-I{torch_include}/caffe2"), format!("-I{cuda_inc}"), + format!("-I{nccl_include}"), format!("-L{torch_lib}"), format!("-Wl,-rpath,{torch_lib}"), + format!("-L{nccl_lib}"), + format!("-Wl,-rpath,{nccl_lib}"), format!("-L{cuda_inc}/../lib64"), format!("-Wl,-rpath,{cuda_inc}/../lib64"), "-Wl,--no-as-needed".to_string(), - "-ltorch".to_string(), "-ltorch_cuda".to_string(), "-ltorch_cpu".to_string(), - "-lc10".to_string(), "-lc10_cuda".to_string(), "-lcudart".to_string(), "-lnccl".to_string(), + "-ltorch".to_string(), + "-ltorch_cuda".to_string(), + "-ltorch_cpu".to_string(), + "-lc10".to_string(), + "-lc10_cuda".to_string(), + "-lcudart".to_string(), + nccl_link.to_string(), ]; for obj in &obj_files { link_args.push(obj.clone()); } - let _ = Command::new("g++").args(&link_args).status(); + cpp_ok + && Command::new("g++") + .args(&link_args) + .status() + .map(|s| s.success()) + .unwrap_or(false) + } else { + false } - } + } else { + false + }; - match status { - Ok(s) if s.success() => { - println!("cargo:rustc-link-search=native={torch_lib}"); - println!("cargo:rustc-link-search=native={out_dir}"); - println!("cargo:rustc-link-lib=dylib=qwen36_kernels"); - println!("cargo:rustc-link-lib=dylib=c10"); - println!("cargo:rustc-link-lib=dylib=torch"); - println!("cargo:rustc-link-lib=dylib=torch_cpu"); - println!("cargo:rustc-link-lib=dylib=torch_cuda"); - println!("cargo:rustc-link-arg=-Wl,--no-as-needed"); - println!("cargo:rustc-link-arg=-Wl,--allow-shlib-undefined"); - } - _ => { - println!("cargo:warning=Failed to compile Qwen3.6 kernels, C++ path disabled"); - } + if build_ok { + println!("cargo:rustc-link-search=native={torch_lib}"); + println!("cargo:rustc-link-search=native={out_dir}"); + println!("cargo:rustc-link-lib=dylib=qwen36_kernels"); + println!("cargo:rustc-link-lib=dylib=c10"); + println!("cargo:rustc-link-lib=dylib=torch"); + println!("cargo:rustc-link-lib=dylib=torch_cpu"); + println!("cargo:rustc-link-lib=dylib=torch_cuda"); + println!("cargo:rustc-link-arg=-Wl,--no-as-needed"); + println!("cargo:rustc-link-arg=-Wl,--allow-shlib-undefined"); + } else { + let _ = std::fs::remove_file(&output_lib); + println!( + "cargo:warning=Failed to compile complete Qwen3.6 C++/CUDA kernels; native training path disabled" + ); } } diff --git a/crates/rustrain-qwen3-6/kernels/delta_rule.cu b/crates/rustrain-qwen3-6/kernels/delta_rule.cu index 424d13e2..bd9831c7 100644 --- a/crates/rustrain-qwen3-6/kernels/delta_rule.cu +++ b/crates/rustrain-qwen3-6/kernels/delta_rule.cu @@ -4,18 +4,23 @@ #include "delta_rule.cuh" // C-linkage wrapper for the forward host launcher -extern "C" void cuda_gated_delta_rule( +extern "C" int cuda_gated_delta_rule( const float* q, const float* k, const float* v, const float* g_exp, const float* beta, float* state, float* out, float* delta_buf, - int BH, int seq_len, int key_dim, int val_dim + int BH, int seq_len, int key_dim, int val_dim, void* stream_ptr ) { + if (key_dim != DR_D_K || val_dim != DR_D_V) { + return -1; + } launch_gated_delta_rule(q, k, v, g_exp, beta, state, out, delta_buf, - BH, seq_len, key_dim, val_dim); + BH, seq_len, key_dim, val_dim, + reinterpret_cast(stream_ptr)); + return static_cast(cudaGetLastError()); } // C-linkage wrapper for the backward host launcher -extern "C" void cuda_gated_delta_rule_backward( +extern "C" int cuda_gated_delta_rule_backward( const float* q, const float* k, const float* v, const float* g_exp, const float* beta, const float* final_state, @@ -23,10 +28,14 @@ extern "C" void cuda_gated_delta_rule_backward( const float* grad_out, float* grad_q, float* grad_k, float* grad_v, float* grad_g, float* grad_beta, - int BH, int seq_len, int key_dim, int val_dim + int BH, int seq_len, int key_dim, int val_dim, void* stream_ptr ) { - launch_gated_delta_rule_backward(q, k, v, g_exp, beta, + if (key_dim != DR_D_K || val_dim != DR_D_V) return -1; + int launch_status = launch_gated_delta_rule_backward(q, k, v, g_exp, beta, final_state, delta_buf, grad_out, grad_q, grad_k, grad_v, grad_g, grad_beta, - BH, seq_len, key_dim, val_dim); + BH, seq_len, key_dim, val_dim, + reinterpret_cast(stream_ptr)); + if (launch_status != 0) return launch_status; + return static_cast(cudaGetLastError()); } diff --git a/crates/rustrain-qwen3-6/kernels/delta_rule.cuh b/crates/rustrain-qwen3-6/kernels/delta_rule.cuh index 712afdc0..c5a54385 100644 --- a/crates/rustrain-qwen3-6/kernels/delta_rule.cuh +++ b/crates/rustrain-qwen3-6/kernels/delta_rule.cuh @@ -144,8 +144,11 @@ __global__ void gated_delta_rule_kernel( // --- delta = (v - kv_mem) * beta --- float delta = (v_t - kv_mem) * beta_t; - // Save delta for backward pass - delta_buf[bh * S * DR_D_V + t * DR_D_V + tid] = delta; + // Save delta only for the explicit chunked backward path. The + // autograd wrapper passes nullptr and recomputes a reference backward. + if (delta_buf != nullptr) { + delta_buf[bh * S * DR_D_V + t * DR_D_V + tid] = delta; + } // --- State update: S[:, dv] += k_t * delta --- // Rank-1 update: each row gets k_t[dk] * delta added @@ -182,7 +185,7 @@ inline void launch_gated_delta_rule( const float* g_exp, const float* beta, float* state, float* out, float* delta_buf, int BH, int seq_len, int key_dim, int val_dim, - cudaStream_t stream = 0 + cudaStream_t stream ) { if (key_dim != DR_D_K || val_dim != DR_D_V) { fprintf(stderr, "[delta_rule] ERROR: D_K=%d or D_V=%d mismatch (expected %d/%d)\n", @@ -337,7 +340,7 @@ __global__ void gated_delta_rule_backward_kernel( // Actually beta is [BH, S], not per-dv. Need reduction. // For now, store per-dv and reduce later. } - gb_bh[t] = 0.0f; // placeholder — needs warp reduction + gb_bh[t] = 0.0f; // legacy kernel; not referenced by the host launcher // --- Step 5: grad_k from kv --- // d(kv)/d(k) = S_before_g → grad_k_kv = S_before_g · grad_kv @@ -360,17 +363,163 @@ __global__ void gated_delta_rule_backward_kernel( // --- Step 7: grad_g = sum(grad_S_before_decay * S_before) --- // This needs S_before = state_s / g[t], and grad_S_before_decay // For simplicity, approximate grad_g as 0 (g is exp of a_log, small gradient) - gg_bh[t] = 0.0f; // TODO: compute properly + gg_bh[t] = 0.0f; // legacy kernel; not referenced by the host launcher __syncthreads(); } } +// Correct reverse-mode recurrence. The older kernel above is retained only as +// a historical reference and is not used by any launcher; this kernel computes +// all q/k/v/g/beta gradients. +// The forward recurrence is: +// R = g * S_prev, kv = k^T R, delta = beta * (v - kv), +// S = R + k outer delta, out = q^T S. +__global__ void gated_delta_rule_backward_kernel_correct( + const float* __restrict__ q, + const float* __restrict__ k, + const float* __restrict__ v, + const float* __restrict__ g_exp, + const float* __restrict__ beta, + const float* __restrict__ final_state, + const float* __restrict__ delta_buf, + const float* __restrict__ grad_out, + float* __restrict__ grad_q, + float* __restrict__ grad_k, + float* __restrict__ grad_v, + float* __restrict__ grad_g, + float* __restrict__ grad_beta, + int S +) { + const int bh = blockIdx.x; + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + + extern __shared__ float smem[]; + float* state_s = smem; // [D_K, D_V] = S_t while entering each step + float* grad_s = state_s + DR_D_K * DR_D_V; + // Four warp partials for reductions over D_V, followed by two scalar + // reductions. This avoids atomics and keeps the 128-thread block intact. + float* reduce_q = grad_s + DR_D_K * DR_D_V; + float* reduce_k = reduce_q + 4 * DR_D_K; + float* reduce_beta = reduce_k + 4 * DR_D_K; + float* reduce_g = reduce_beta + 4; + + const float* state_g = final_state + bh * DR_D_K * DR_D_V; + for (int i = tid; i < DR_D_K * DR_D_V; i += DR_THREADS) { + state_s[i] = state_g[i]; + grad_s[i] = 0.0f; + } + __syncthreads(); + + const float* q_bh = q + bh * S * DR_D_K; + const float* k_bh = k + bh * S * DR_D_K; + const float* v_bh = v + bh * S * DR_D_V; + const float* g_bh = g_exp + bh * S; + const float* beta_bh = beta + bh * S; + const float* go_bh = grad_out + bh * S * DR_D_V; + float* gq_bh = grad_q + bh * S * DR_D_K; + float* gk_bh = grad_k + bh * S * DR_D_K; + float* gv_bh = grad_v + bh * S * DR_D_V; + float* gg_bh = grad_g + bh * S; + float* gb_bh = grad_beta + bh * S; + + for (int t = S - 1; t >= 0; --t) { + const float* q_t = q_bh + t * DR_D_K; + const float* k_t = k_bh + t * DR_D_K; + const float g_t = g_bh[t]; + const float beta_t = beta_bh[t]; + const float go_t = go_bh[t * DR_D_V + tid]; + const float delta_t = delta_buf[bh * S * DR_D_V + t * DR_D_V + tid]; + + // Add the direct output contribution to dS_t and reduce dQ_t over + // value columns while state_s still contains the post-update state. + for (int dk = 0; dk < DR_D_K; ++dk) { + const int idx = dk * DR_D_V + tid; + const float s_after = state_s[idx]; + grad_s[idx] += q_t[dk] * go_t; + float q_part = s_after * go_t; + for (int off = 16; off > 0; off >>= 1) + q_part += __shfl_down_sync(0xffffffff, q_part, off); + if (lane == 0) reduce_q[warp * DR_D_K + dk] = q_part; + } + __syncthreads(); + if (tid < DR_D_K) { + gq_bh[t * DR_D_K + tid] = + reduce_q[tid] + reduce_q[DR_D_K + tid] + + reduce_q[2 * DR_D_K + tid] + reduce_q[3 * DR_D_K + tid]; + } + __syncthreads(); + + // Undo S_t = R_t + k outer delta, leaving R_t = g_t*S_prev. + for (int dk = 0; dk < DR_D_K; ++dk) + state_s[dk * DR_D_V + tid] -= k_t[dk] * delta_t; + __syncthreads(); + + // ddelta = dS_t^T k and h = ddelta * beta. The same value-column + // thread computes dV and one partial for dBeta/dG. + float gdelta = 0.0f; + float kv_mem = 0.0f; + for (int dk = 0; dk < DR_D_K; ++dk) { + const int idx = dk * DR_D_V + tid; + gdelta += grad_s[idx] * k_t[dk]; + kv_mem += state_s[idx] * k_t[dk]; + } + const float h = gdelta * beta_t; + gv_bh[t * DR_D_V + tid] = h; + const float beta_part = gdelta * (v_bh[t * DR_D_V + tid] - kv_mem); + const float safe_g = fmaxf(g_t, 1.0e-8f); + float g_part = 0.0f; + + // dR = dS - k outer h. Reduce dK over value columns and update dS_prev. + for (int dk = 0; dk < DR_D_K; ++dk) { + const int idx = dk * DR_D_V + tid; + const float r = state_s[idx]; + const float s_prev = r / safe_g; + const float grad_r = grad_s[idx] - k_t[dk] * h; + const float k_part = grad_s[idx] * delta_t - r * h; + g_part += grad_r * s_prev; + float reduced_k = k_part; + for (int off = 16; off > 0; off >>= 1) + reduced_k += __shfl_down_sync(0xffffffff, reduced_k, off); + if (lane == 0) reduce_k[warp * DR_D_K + dk] = reduced_k; + grad_s[idx] = grad_r * g_t; + // The next reverse iteration starts from S_prev, not R_t. + state_s[idx] = s_prev; + } + + float reduced_beta = beta_part; + float reduced_g = g_part; + for (int off = 16; off > 0; off >>= 1) { + reduced_beta += __shfl_down_sync(0xffffffff, reduced_beta, off); + reduced_g += __shfl_down_sync(0xffffffff, reduced_g, off); + } + if (lane == 0) { + reduce_beta[warp] = reduced_beta; + reduce_g[warp] = reduced_g; + } + __syncthreads(); + if (tid < DR_D_K) { + gk_bh[t * DR_D_K + tid] = + reduce_k[tid] + reduce_k[DR_D_K + tid] + + reduce_k[2 * DR_D_K + tid] + reduce_k[3 * DR_D_K + tid]; + } + if (tid == 0) { + gb_bh[t] = reduce_beta[0] + reduce_beta[1] + + reduce_beta[2] + reduce_beta[3]; + gg_bh[t] = reduce_g[0] + reduce_g[1] + + reduce_g[2] + reduce_g[3]; + } + __syncthreads(); + } +} + // ────────────────────────────────────────────────────────────────────── // Host launcher (backward) // ────────────────────────────────────────────────────────────────────── -inline void launch_gated_delta_rule_backward( +inline int launch_gated_delta_rule_backward( const float* q, const float* k, const float* v, const float* g_exp, const float* beta, const float* final_state, @@ -379,24 +528,30 @@ inline void launch_gated_delta_rule_backward( float* grad_q, float* grad_k, float* grad_v, float* grad_g, float* grad_beta, int BH, int seq_len, int key_dim, int val_dim, - cudaStream_t stream = 0 + cudaStream_t stream ) { if (key_dim != DR_D_K || val_dim != DR_D_V) { fprintf(stderr, "[delta_rule_backward] ERROR: D_K=%d or D_V=%d mismatch\n", key_dim, val_dim); - return; + return -1; } dim3 grid(BH); dim3 block(DR_THREADS); - // 2 state matrices in shared memory: 2 * 64KB = 128KB - size_t smem_size = 2 * DR_D_K * DR_D_V * sizeof(float); - cudaFuncSetAttribute( - gated_delta_rule_backward_kernel, + // 2 state matrices plus four warp partials per D_K and scalar reductions. + size_t smem_size = (2 * DR_D_K * DR_D_V + 8 * DR_D_K + 8) * sizeof(float); + auto attr_status = cudaFuncSetAttribute( + gated_delta_rule_backward_kernel_correct, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + if (attr_status != cudaSuccess) { + fprintf(stderr, "[delta_rule_backward] shared-memory attribute failed: %s\n", + cudaGetErrorString(attr_status)); + return static_cast(attr_status); + } - gated_delta_rule_backward_kernel<<>>( + gated_delta_rule_backward_kernel_correct<<>>( q, k, v, g_exp, beta, final_state, delta_buf, grad_out, grad_q, grad_k, grad_v, grad_g, grad_beta, seq_len ); + return 0; } diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index ba76c46e..3e30f4a6 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -5,6 +5,12 @@ // No tch-rs VarStore involved — gradients flow entirely within C++ autograd. #include +#if __has_include() +#include +#define RUSTRAIN_HAS_ATEN_GROUPED_MM 1 +#else +#define RUSTRAIN_HAS_ATEN_GROUPED_MM 0 +#endif #include #include #include @@ -13,6 +19,8 @@ #include #include #include +#include +#include #include #include #include @@ -41,6 +49,30 @@ at::Tensor apply_multi_lora(TrainingContext* ctx, int64_t layer_idx, int64_t pai static at::Tensor rms_norm(const at::Tensor& input, const at::Tensor& weight, double eps); +static bool env_enabled(const char* name) { + const char* value = std::getenv(name); + return value && value[0] != '\0' && std::strcmp(value, "0") != 0; +} + +static std::string nccl_sync_dir() { + const char* run_id = std::getenv("RUSTRAIN_NCCL_RUN_ID"); + if (!run_id || run_id[0] == '\0') return "/tmp/rustrain-nccl"; + std::string sanitized; + for (const unsigned char ch : std::string(run_id)) { + if ((ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || + (ch >= '0' && ch <= '9') || ch == '-' || ch == '_') { + sanitized.push_back(static_cast(ch)); + } else { + sanitized.push_back('_'); + } + } + if (sanitized.empty()) sanitized = "default"; + mkdir("/tmp/rustrain-nccl", 0777); + std::string path = "/tmp/rustrain-nccl/" + sanitized; + mkdir(path.c_str(), 0777); + return path; +} + // ────────────────────────────────────────────────────────────────────── // Hand-written CUDA fused kernels (compiled from fused_kernels.cu) // ────────────────────────────────────────────────────────────────────── @@ -83,29 +115,78 @@ static at::Tensor fused_rmsnorm_op( /// d/dg = sigmoid(g) * (1 + g * (1 - sigmoid(g))) * u /// d/du = silu(g) = g * sigmoid(g) struct NcclAllReduceFunction : public torch::autograd::Function { + static ncclDataType_t dtype_for(at::ScalarType type) { + switch (type) { + case at::kBFloat16: return ncclBfloat16; + case at::kFloat: return ncclFloat; + case at::kHalf: return ncclFloat16; + default: + TORCH_CHECK(false, "unsupported NCCL all-reduce dtype: ", type); + } + } + + // NCCL is normally issued on PyTorch's current stream. When an external + // NCCL stream is supplied by the EP runtime, fence both sides so the + // collective observes the producer and its result is visible to the + // current compute stream. This keeps the operation asynchronous without + // relying on a device-wide synchronize. + static at::Tensor allreduce( + const at::Tensor& input, ncclComm_t comm, cudaStream_t requested_stream + ) { + TORCH_CHECK(input.is_cuda(), "NCCL all-reduce requires a CUDA tensor"); + const int dev = input.device().index(); + cudaSetDevice(dev); + const auto current_stream = c10::cuda::getCurrentCUDAStream(dev).stream(); + const auto comm_stream = requested_stream ? requested_stream : current_stream; + auto contiguous = input.contiguous(); + auto output = at::empty_like(contiguous); + + cudaEvent_t before = nullptr; + cudaEvent_t after = nullptr; + const bool cross_stream = comm_stream != current_stream; + if (cross_stream) { + TORCH_CHECK(cudaEventCreateWithFlags(&before, cudaEventDisableTiming) == cudaSuccess, + "failed to create NCCL producer event"); + TORCH_CHECK(cudaEventCreateWithFlags(&after, cudaEventDisableTiming) == cudaSuccess, + "failed to create NCCL consumer event"); + TORCH_CHECK(cudaEventRecord(before, current_stream) == cudaSuccess, + "failed to record NCCL producer event"); + TORCH_CHECK(cudaStreamWaitEvent(comm_stream, before, 0) == cudaSuccess, + "failed to wait for NCCL producer event"); + } + + auto err = ncclAllReduce( + contiguous.data_ptr(), output.data_ptr(), contiguous.numel(), + dtype_for(contiguous.scalar_type()), ncclSum, comm, comm_stream); + TORCH_CHECK(err == ncclSuccess, "ncclAllReduce failed: ", ncclGetErrorString(err)); + + if (cross_stream) { + TORCH_CHECK(cudaEventRecord(after, comm_stream) == cudaSuccess, + "failed to record NCCL consumer event"); + TORCH_CHECK(cudaStreamWaitEvent(current_stream, after, 0) == cudaSuccess, + "failed to wait for NCCL consumer event"); + // Destruction is asynchronous-safe after the wait has been + // enqueued on the current stream. + cudaEventDestroy(before); + cudaEventDestroy(after); + } + return output; + } + static at::Tensor forward(torch::autograd::AutogradContext* ctx, - at::Tensor input, int64_t comm_ptr) { + at::Tensor input, int64_t comm_ptr, int64_t stream_ptr) { ctx->saved_data["comm"] = comm_ptr; + ctx->saved_data["stream"] = stream_ptr; auto nccl_comm = reinterpret_cast(comm_ptr); - int dev = input.device().index(); - cudaSetDevice(dev); - auto compute_stream = c10::cuda::getCurrentCUDAStream(dev).stream(); ctx->save_for_backward({input}); - auto output = at::empty_like(input); - ncclResult_t err = ncclAllReduce( - input.data_ptr(), output.data_ptr(), - input.numel(), ncclBfloat16, ncclSum, - nccl_comm, compute_stream); - if (err != ncclSuccess) { - fprintf(stderr, "[ep] ncclAllReduce fwd FAILED: %d (%s) dev=%d\n", err, ncclGetErrorString(err), dev); - } - return output; + return allreduce(input, nccl_comm, + reinterpret_cast(stream_ptr)); } static std::vector backward(torch::autograd::AutogradContext* ctx, std::vector grad_output) { // Expert weights frozen — gradient flows through residual connection. // Save input in forward so autograd keeps it alive until backward. - return {grad_output[0], at::Tensor()}; + return {grad_output[0], at::Tensor(), at::Tensor()}; } }; @@ -139,7 +220,8 @@ struct FusedSwiGLUFunction : public torch::autograd::Function 0.0) inter = inter.clamp(-limit, limit); return inter; @@ -334,13 +416,16 @@ static at::Tensor full_attention( auto attn_out = at::scaled_dot_product_attention( q, k, v, additive_mask, 0.0, false, c10::nullopt, true // is_causal=false, enable_gqa=true ); - return attn_out.transpose(1, 2).reshape({batch, seq, qkv_dim}).matmul(o_proj.t()); + // Qwen3.5/3.6 full attention gates the attention value before o_proj. + // Applying the gate after o_proj is not equivalent when o_proj mixes features. + auto gated_attn = attn_out * at::sigmoid(gate).to(attn_out.scalar_type()); + return gated_attn.transpose(1, 2).reshape({batch, seq, qkv_dim}).matmul(o_proj.t()); } else { auto attn_out = at::scaled_dot_product_attention( q, k, v, c10::nullopt, 0.0, true, c10::nullopt, true // is_causal=true, enable_gqa=true ); - auto result = attn_out.transpose(1, 2).reshape({batch, seq, qkv_dim}).matmul(o_proj.t()); - result = result * at::sigmoid(gate).to(result.scalar_type()); + auto gated_attn = attn_out * at::sigmoid(gate).to(attn_out.scalar_type()); + auto result = gated_attn.transpose(1, 2).reshape({batch, seq, qkv_dim}).matmul(o_proj.t()); gate = at::Tensor(); return result; } @@ -351,23 +436,148 @@ static at::Tensor full_attention( // ────────────────────────────────────────────────────────────────────── // Forward declaration for CUDA kernel (defined in delta_rule.cu) -extern "C" void cuda_gated_delta_rule( +extern "C" int cuda_gated_delta_rule( const float* q, const float* k, const float* v, const float* g_exp, const float* beta, float* state, float* out, float* delta_buf, - int BH, int seq_len, int key_dim, int val_dim + int BH, int seq_len, int key_dim, int val_dim, void* stream ); -extern "C" void cuda_gated_delta_rule_backward( +extern "C" int cuda_gated_delta_rule_backward( const float* q, const float* k, const float* v, const float* g_exp, const float* beta, const float* final_state, const float* delta_buf, const float* grad_out, float* grad_q, float* grad_k, float* grad_v, float* grad_g, float* grad_beta, - int BH, int seq_len, int key_dim, int val_dim + int BH, int seq_len, int key_dim, int val_dim, void* stream ); +// Correctness reference for the gated delta rule. This deliberately uses +// ATen batched operations inside C++, so it remains outside the Rust hot path +// while providing a complete autograd oracle for the custom CUDA forward. +static at::Tensor gated_delta_rule_reference( + const at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + const at::Tensor& g_exp, const at::Tensor& beta +) { + TORCH_CHECK(q.dim() == 3 && k.dim() == 3 && v.dim() == 3, + "gated_delta_rule_reference expects [BH, S, D] tensors"); + const int64_t bh = q.size(0); + const int64_t seq = q.size(1); + const int64_t key_dim = q.size(2); + const int64_t val_dim = v.size(2); + TORCH_CHECK(k.size(0) == bh && k.size(1) == seq && k.size(2) == key_dim, + "q/k shape mismatch"); + TORCH_CHECK(v.size(0) == bh && v.size(1) == seq, + "v shape mismatch"); + TORCH_CHECK(g_exp.size(0) == bh && g_exp.size(1) == seq && + beta.size(0) == bh && beta.size(1) == seq, + "g/beta shape mismatch"); + + auto state = at::zeros({bh, key_dim, val_dim}, + q.options().dtype(at::kFloat)); + auto qf = q.to(at::kFloat); + auto kf = k.to(at::kFloat); + auto vf = v.to(at::kFloat); + auto gf = g_exp.to(at::kFloat); + auto bf = beta.to(at::kFloat); + std::vector outputs; + outputs.reserve(seq); + for (int64_t t = 0; t < seq; ++t) { + auto gt = gf.select(1, t).view({bh, 1, 1}); + auto kt = kf.select(1, t); + auto vt = vf.select(1, t); + auto qt = qf.select(1, t); + auto bt = bf.select(1, t).view({bh, 1}); + state = state * gt; + auto kv = at::bmm(kt.unsqueeze(1), state).squeeze(1); + auto delta = (vt - kv) * bt; + state = state + kt.unsqueeze(2) * delta.unsqueeze(1); + outputs.push_back(at::bmm(qt.unsqueeze(1), state).squeeze(1)); + } + return at::stack(outputs, 1); +} + +// The forward and backward CUDA kernels are wrapped in one autograd Function. +// Set QWEN36_DELTA_REFERENCE_BWD=1 to run the ATen recurrence oracle for parity +// debugging; production training uses the current-stream fused backward. +struct GatedDeltaRuleFunction : public torch::autograd::Function { + static at::Tensor forward( + torch::autograd::AutogradContext* ctx, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor g_exp, at::Tensor beta + ) { + TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), + "gated delta CUDA path requires CUDA tensors"); + TORCH_CHECK(q.scalar_type() == at::kFloat && k.scalar_type() == at::kFloat && + v.scalar_type() == at::kFloat && g_exp.scalar_type() == at::kFloat && + beta.scalar_type() == at::kFloat, + "gated delta CUDA path expects FP32 working tensors"); + const int64_t bh = q.size(0), seq = q.size(1); + const int64_t key_dim = q.size(2), val_dim = v.size(2); + auto state = at::zeros({bh, key_dim, val_dim}, q.options()); + auto out = at::empty({bh, seq, val_dim}, q.options()); + auto delta_buf = at::empty({bh, seq, val_dim}, q.options()); + auto stream = c10::cuda::getCurrentCUDAStream(q.device().index()).stream(); + int status = cuda_gated_delta_rule( + q.data_ptr(), k.data_ptr(), v.data_ptr(), + g_exp.data_ptr(), beta.data_ptr(), + state.data_ptr(), out.data_ptr(), delta_buf.data_ptr(), + (int)bh, (int)seq, (int)key_dim, (int)val_dim, + reinterpret_cast(stream)); + TORCH_CHECK(status == 0, "gated delta CUDA launch failed: ", status); + ctx->save_for_backward({q, k, v, g_exp, beta, state, delta_buf}); + return out; + } + + static std::vector backward( + torch::autograd::AutogradContext* ctx, + std::vector grad_output + ) { + auto saved = ctx->get_saved_variables(); + auto q = saved[0]; + auto k = saved[1]; + auto v = saved[2]; + auto g = saved[3]; + auto beta = saved[4]; + if (env_enabled("QWEN36_DELTA_REFERENCE_BWD")) { + auto q_ref = q.detach().set_requires_grad(true); + auto k_ref = k.detach().set_requires_grad(true); + auto v_ref = v.detach().set_requires_grad(true); + auto g_ref = g.detach().set_requires_grad(true); + auto beta_ref = beta.detach().set_requires_grad(true); + at::AutoGradMode guard(true); + auto reference = gated_delta_rule_reference(q_ref, k_ref, v_ref, g_ref, beta_ref); + auto grads = torch::autograd::grad( + {reference}, {q_ref, k_ref, v_ref, g_ref, beta_ref}, + {grad_output[0]}, /*retain_graph=*/false, + /*create_graph=*/false, /*allow_unused=*/false); + return {grads[0], grads[1], grads[2], grads[3], grads[4]}; + } + + const int64_t bh = q.size(0), seq = q.size(1); + const int64_t key_dim = q.size(2), val_dim = v.size(2); + auto grad_out = grad_output[0].contiguous(); + auto grad_q = at::empty_like(q); + auto grad_k = at::empty_like(k); + auto grad_v = at::empty_like(v); + auto grad_g = at::empty_like(g); + auto grad_beta = at::empty_like(beta); + auto stream = c10::cuda::getCurrentCUDAStream(q.device().index()).stream(); + int status = cuda_gated_delta_rule_backward( + q.data_ptr(), k.data_ptr(), v.data_ptr(), + g.data_ptr(), beta.data_ptr(), + saved[5].data_ptr(), saved[6].data_ptr(), + grad_out.data_ptr(), grad_q.data_ptr(), + grad_k.data_ptr(), grad_v.data_ptr(), + grad_g.data_ptr(), grad_beta.data_ptr(), + (int)bh, (int)seq, (int)key_dim, (int)val_dim, + reinterpret_cast(stream)); + TORCH_CHECK(status == 0, "gated delta backward launch failed: ", status); + return {grad_q, grad_k, grad_v, grad_g, grad_beta}; + } +}; + static at::Tensor linear_attention( const at::Tensor& hidden, const at::Tensor& in_proj_qkv, const at::Tensor& in_proj_z, @@ -391,7 +601,10 @@ static at::Tensor linear_attention( int64_t seq_chunk = 0; if (chunk_env) seq_chunk = atoll(chunk_env); - if (seq_chunk > 0 && seq > seq_chunk) { + // Stateful chunking currently has no autograd state input/output contract. + // Keep it for inference/eval only; training uses the autograd-wrapped full + // sequence path below until chunk state gradients are implemented. + if (seq_chunk > 0 && seq > seq_chunk && !at::GradMode::is_enabled()) { // Chunked linear attention — mathematically equivalent to full sequence // Process in chunks, passing the delta rule state between chunks. // This avoids creating [batch, seq, qkv_dim] intermediate tensors. @@ -471,7 +684,8 @@ static at::Tensor linear_attention( auto delta_buf = at::empty({BH, chunk_len, val_dim}, q_t.options()); // CUDA kernel — state is passed in and updated in-place - cuda_gated_delta_rule( + auto stream = c10::cuda::getCurrentCUDAStream(device.index()).stream(); + int status = cuda_gated_delta_rule( q_contig.data_ptr(), k_contig.data_ptr(), v_contig.data_ptr(), @@ -480,8 +694,10 @@ static at::Tensor linear_attention( state_contig.data_ptr(), outs.data_ptr(), delta_buf.data_ptr(), - (int)BH, (int)chunk_len, (int)key_dim, (int)val_dim + (int)BH, (int)chunk_len, (int)key_dim, (int)val_dim, + reinterpret_cast(stream) ); + TORCH_CHECK(status == 0, "gated delta CUDA launch failed: ", status); state = state_contig; // updated state for next chunk auto core_out = outs.reshape({batch, num_v_heads, chunk_len, val_dim}) @@ -563,30 +779,14 @@ static at::Tensor linear_attention( // All computation runs in a single kernel launch using shared memory for state auto g_exp = g_t.exp(); // [B, H, S] int64_t BH = batch * num_v_heads; - auto state = at::zeros({BH, key_dim, val_dim}, q_t.options()); // [B*H, D_k, D_v] - - // Prepare contiguous FP32 tensors for CUDA kernel + // Prepare contiguous FP32 tensors for the CUDA forward/autograd wrapper. auto q_contig = q_t.reshape({BH, seq, key_dim}).contiguous().to(at::kFloat); auto k_contig = k_t.reshape({BH, seq, key_dim}).contiguous().to(at::kFloat); auto v_contig = v_t.reshape({BH, seq, val_dim}).contiguous().to(at::kFloat); auto g_contig = g_exp.reshape({BH, seq}).contiguous().to(at::kFloat); auto beta_contig = beta_t.reshape({BH, seq}).contiguous().to(at::kFloat); - auto state_contig = state.contiguous(); - auto outs = at::empty({BH, seq, val_dim}, q_t.options()); - auto delta_buf = at::empty({BH, seq, val_dim}, q_t.options()); - - // Launch CUDA kernel — single launch replaces seq×3 bmm calls - cuda_gated_delta_rule( - q_contig.data_ptr(), - k_contig.data_ptr(), - v_contig.data_ptr(), - g_contig.data_ptr(), - beta_contig.data_ptr(), - state_contig.data_ptr(), - outs.data_ptr(), - delta_buf.data_ptr(), - (int)BH, (int)seq, (int)key_dim, (int)val_dim - ); + auto outs = GatedDeltaRuleFunction::apply( + q_contig, k_contig, v_contig, g_contig, beta_contig); // Reshape: [B*H, S, D_v] → [B, H, S, D_v] → [B, S, H, D_v] auto core_out = outs.reshape({batch, num_v_heads, seq, val_dim}) @@ -641,15 +841,88 @@ static at::Tensor linear_attention_batched(TrainingContext* ctx, const at::Tenso // MoE // ────────────────────────────────────────────────────────────────────── +struct RoutedExpertLora { + const at::Tensor* gate_up_a = nullptr; // [local_experts, rank, hidden] + const at::Tensor* gate_up_b = nullptr; // [local_experts, 2*intermediate, rank] + const at::Tensor* down_a = nullptr; // [local_experts, rank, intermediate] + const at::Tensor* down_b = nullptr; // [local_experts, hidden, rank] + double scaling = 0.0; +}; + +// Per-sample adapter projection used by the batched multi-LoRA path. This is +// intentionally separate from RoutedExpertLora: routed experts carry one +// A/B pair per local expert, while dense/shared projections carry one pair per +// adapter sample. +struct LoraBatchEntry { + at::Tensor a_stack; // [N, rank, in] + at::Tensor b_stack; // [N, out, rank] + at::Tensor scaling; // [N, 1, 1] +}; + +static const LoraBatchEntry* lora_batch_entry( + TrainingContext* ctx, int64_t layer_idx, int64_t pair_idx); +static at::Tensor dense_mlp_forward_batched( + TrainingContext* ctx, int64_t layer_idx, const at::Tensor& hidden, + const at::Tensor& gate_proj, const at::Tensor& up_proj, + const at::Tensor& down_proj, at::ScalarType compute_type); + +static at::Tensor lora_activation_delta( + const at::Tensor& x, const at::Tensor& A, const at::Tensor& B, + const at::Tensor& scaling); + +static at::Tensor add_batched_lora( + const at::Tensor& base, const at::Tensor& input, + const LoraBatchEntry* entry +) { + if (!entry) return base; + return base + lora_activation_delta( + input, entry->a_stack, entry->b_stack, entry->scaling); +} + +// Per-token routed-expert LoRA. Dynamic adapters add a leading sample axis to +// the expert-local tensors: A [batch, experts, rank, in], +// B [batch, experts, out, rank]. Flattening (sample, expert) lets one pair of +// index_select + bmm operations select the correct adapter and expert without +// materializing a full-rank delta weight. +static at::Tensor dynamic_expert_lora_delta( + const at::Tensor& input, + const at::Tensor& token_indices, + const at::Tensor& local_expert_indices, + int64_t seq, + const LoraBatchEntry* entry +) { + if (!entry) return at::zeros({0}, input.options()); + TORCH_CHECK(entry->a_stack.dim() == 4 && entry->b_stack.dim() == 4, + "dynamic routed-expert LoRA expects rank-4 stacked A/B tensors"); + const int64_t local_experts = entry->a_stack.size(1); + auto sample_indices = at::floor_divide(token_indices, seq); + auto pair_indices = sample_indices * local_experts + local_expert_indices; + auto a = entry->a_stack.flatten(0, 1) + .index_select(0, pair_indices).to(input.scalar_type()); + auto b = entry->b_stack.flatten(0, 1) + .index_select(0, pair_indices).to(input.scalar_type()); + auto low_rank = at::bmm(a, input.unsqueeze(-1)).squeeze(-1); + auto delta = at::bmm(b, low_rank.unsqueeze(-1)).squeeze(-1); + auto scaling = entry->scaling.index_select(0, sample_indices) + .reshape({-1, 1}).to(input.scalar_type()); + return delta * scaling; +} + static at::Tensor moe_forward( void* nccl_comm_v, void* nccl_stream_v, const at::Tensor& hidden, const at::Tensor& gate_w, const at::Tensor& shared_expert_gate_w, const at::Tensor& shared_gate_proj, const at::Tensor& shared_up_proj, const at::Tensor& shared_down_proj, const at::Tensor& experts_gate_up, const at::Tensor& experts_down, + const RoutedExpertLora& expert_lora, int64_t num_experts, int64_t top_k, int64_t intermediate, bool norm_topk_prob, int64_t expert_start, int64_t expert_count, - at::ScalarType compute_type + at::ScalarType compute_type, + const LoraBatchEntry* shared_gate_lora = nullptr, + const LoraBatchEntry* shared_up_lora = nullptr, + const LoraBatchEntry* shared_down_lora = nullptr, + const LoraBatchEntry* expert_gate_up_lora = nullptr, + const LoraBatchEntry* expert_down_lora = nullptr ) { int64_t batch = hidden.size(0), seq = hidden.size(1), hidden_dim = hidden.size(2); auto device = hidden.device(); @@ -697,42 +970,213 @@ static at::Tensor moe_forward( // Find expert boundaries via bincount + cumsum auto counts = at::bincount(sorted_indices, c10::nullopt, expert_start + expert_count); // counts[e_global] = number of tokens assigned to expert e - // Gather tokens in sorted order (contiguous per expert) auto gathered = flat.index_select(0, sort_order); auto gathered_weights = expert_weights.index_select(0, sort_order).unsqueeze(-1); - // Process each expert's contiguous token slice - int64_t offset = 0; + // The grouped single-rank path knows that every sorted row is local + // and therefore needs no host-visible expert counts. Materialize the + // CPU copy lazily only for EP slicing or the legacy per-expert loop. + at::Tensor counts_cpu; + auto get_counts_cpu = [&]() -> const at::Tensor& { + if (!counts_cpu.defined()) { + counts_cpu = counts.to(at::TensorOptions().device(at::kCPU)); + } + return counts_cpu; + }; + +#if RUSTRAIN_HAS_ATEN_GROUPED_MM + // PyTorch 2.12+ exposes the same CUTLASS grouped-GEMM primitive used + // by its native MoE path. It consumes sorted token rows plus cumulative + // expert offsets, eliminating one GEMM launch per local expert. Older + // libtorch builds compile the fallback loop below. + bool grouped_lora_compatible = true; + if (expert_lora.gate_up_a) { + grouped_lora_compatible = expert_lora.gate_up_a->size(1) % 8 == 0; + } + if (expert_lora.down_a) { + grouped_lora_compatible = grouped_lora_compatible && + expert_lora.down_a->size(1) % 8 == 0; + } + const bool use_grouped_mm = !env_enabled("QWEN36_DISABLE_GROUPED_MM") && + compute_type == at::kBFloat16 && grouped_lora_compatible && + hidden_dim % 8 == 0 && intermediate % 8 == 0; + if (use_grouped_mm) { + const bool owns_all_experts = expert_start == 0 && + expert_count == counts.size(0); + const int64_t local_start = owns_all_experts || expert_start == 0 + ? 0 + : get_counts_cpu().narrow(0, 0, expert_start).sum().item(); + const int64_t local_tokens = owns_all_experts + ? gathered.size(0) + : get_counts_cpu() + .narrow(0, expert_start, expert_count).sum().item(); + if (local_tokens > 0) { + if (env_enabled("QWEN36_REPORT_GROUPED_MM")) { + std::fprintf( + stderr, + "[q36_moe] grouped_mm experts=%ld tokens=%ld hidden=%ld intermediate=%ld\n", + static_cast(expert_count), + static_cast(local_tokens), + static_cast(hidden_dim), + static_cast(intermediate)); + } + auto selected = gathered.narrow(0, local_start, local_tokens); + auto token_indices = sort_order.narrow( + 0, local_start, local_tokens); + auto local_expert_indices = sorted_indices.narrow( + 0, local_start, local_tokens) - expert_start; + auto offsets = counts.narrow(0, expert_start, expert_count) + .cumsum(0).to(at::kInt); + auto gu = at::_grouped_mm( + selected, experts_gate_up.transpose(1, 2), offsets); + if (expert_lora.gate_up_a && expert_lora.gate_up_b) { + auto low_rank = at::_grouped_mm( + selected, expert_lora.gate_up_a->transpose(1, 2), offsets); + auto delta = at::_grouped_mm( + low_rank, expert_lora.gate_up_b->transpose(1, 2), offsets); + gu = gu + delta * expert_lora.scaling; + } + if (expert_gate_up_lora) { + gu = gu + dynamic_expert_lora_delta( + selected, token_indices, local_expert_indices, + seq, expert_gate_up_lora); + } + auto activated = fused_swiglu_op( + gu.narrow(-1, 0, intermediate), + gu.narrow(-1, intermediate, intermediate), 0.0); + auto expert_out = at::_grouped_mm( + activated, experts_down.transpose(1, 2), offsets); + if (expert_lora.down_a && expert_lora.down_b) { + auto low_rank = at::_grouped_mm( + activated, expert_lora.down_a->transpose(1, 2), offsets); + auto delta = at::_grouped_mm( + low_rank, expert_lora.down_b->transpose(1, 2), offsets); + expert_out = expert_out + delta * expert_lora.scaling; + } + if (expert_down_lora) { + expert_out = expert_out + dynamic_expert_lora_delta( + activated, token_indices, local_expert_indices, + seq, expert_down_lora); + } + auto weights = gathered_weights.narrow( + 0, local_start, local_tokens); + routed_output = routed_output.index_add_( + 0, token_indices, expert_out * weights); + } + continue; + } +#endif + + // Process each expert's contiguous token slice. `gathered` contains + // all global experts in sorted order, so rank>0 must skip tokens for + // experts owned by lower ranks before taking its local range. + int64_t offset = expert_start > 0 + ? get_counts_cpu().narrow(0, 0, expert_start).sum().item() + : 0; for (int64_t e_local = 0; e_local < expert_count; e_local++) { int64_t e_global = expert_start + e_local; - int64_t n_tokens = counts.index({e_global}).item(); - if (n_tokens == 0) continue; - - auto selected = gathered.narrow(0, offset, n_tokens); // zero-copy view! - auto egu = experts_gate_up.select(0, e_local); - auto ed = experts_down.select(0, e_local); - auto gu = at::matmul(selected, egu.t()); - auto gate_part = gu.narrow(-1, 0, intermediate); - auto up_part = gu.narrow(-1, intermediate, intermediate); - auto expert_out = at::matmul(fused_swiglu_op(gate_part, up_part, 0.0), ed.t()); - auto weights = gathered_weights.narrow(0, offset, n_tokens); - auto token_indices = sort_order.narrow(0, offset, n_tokens); - routed_output = routed_output.index_add_(0, token_indices, expert_out * weights); + int64_t n_tokens = get_counts_cpu().index({e_global}).item(); + if (n_tokens > 0) { + auto selected = gathered.narrow(0, offset, n_tokens); // zero-copy view! + auto token_indices = sort_order.narrow(0, offset, n_tokens); + auto local_expert_indices = sorted_indices.narrow(0, offset, n_tokens) + - expert_start; + auto egu = experts_gate_up.select(0, e_local); + auto ed = experts_down.select(0, e_local); + auto gu = at::matmul(selected, egu.t()); + if (expert_lora.gate_up_a && expert_lora.gate_up_b) { + auto a = expert_lora.gate_up_a->select(0, e_local); + auto b = expert_lora.gate_up_b->select(0, e_local); + gu = gu + at::matmul(at::matmul(selected, a.t()), b.t()) * expert_lora.scaling; + } + if (expert_gate_up_lora) { + gu = gu + dynamic_expert_lora_delta( + selected, token_indices, local_expert_indices, + seq, expert_gate_up_lora); + } + auto gate_part = gu.narrow(-1, 0, intermediate); + auto up_part = gu.narrow(-1, intermediate, intermediate); + auto activated = fused_swiglu_op(gate_part, up_part, 0.0); + auto expert_out = at::matmul(activated, ed.t()); + if (expert_lora.down_a && expert_lora.down_b) { + auto a = expert_lora.down_a->select(0, e_local); + auto b = expert_lora.down_b->select(0, e_local); + expert_out = expert_out + + at::matmul(at::matmul(activated, a.t()), b.t()) * expert_lora.scaling; + } + if (expert_down_lora) { + expert_out = expert_out + dynamic_expert_lora_delta( + activated, token_indices, local_expert_indices, + seq, expert_down_lora); + } + auto weights = gathered_weights.narrow(0, offset, n_tokens); + routed_output = routed_output.index_add_(0, token_indices, expert_out * weights); + } offset += n_tokens; } } + // A rank can receive no tokens for any of its local experts. Keep a + // zero-valued dependency on routed-expert LoRA tensors so autograd still + // produces defined zero gradients and every rank reaches the same NCCL + // collectives. This changes only the graph, not the routed output values. + if (!routed_output.requires_grad()) { + at::Tensor graph_anchor; + auto include_anchor = [&](const at::Tensor* tensor) { + if (!tensor || !tensor->defined() || !tensor->requires_grad()) return; + auto contribution = tensor->sum().to(routed_output.scalar_type()); + graph_anchor = graph_anchor.defined() + ? graph_anchor + contribution + : contribution; + }; + include_anchor(expert_lora.gate_up_a); + include_anchor(expert_lora.gate_up_b); + include_anchor(expert_lora.down_a); + include_anchor(expert_lora.down_b); + if (expert_gate_up_lora) { + include_anchor(&expert_gate_up_lora->a_stack); + include_anchor(&expert_gate_up_lora->b_stack); + } + if (expert_down_lora) { + include_anchor(&expert_down_lora->a_stack); + include_anchor(&expert_down_lora->b_stack); + } + if (graph_anchor.defined()) { + routed_output = routed_output + graph_anchor * 0.0; + } + } + // EP all-reduce via NcclAllReduceFunction — custom autograd Function. if (nccl_comm_v) { auto nccl_comm = reinterpret_cast(nccl_comm_v); - routed_output = NcclAllReduceFunction::apply(routed_output, (int64_t)nccl_comm); + routed_output = NcclAllReduceFunction::apply( + routed_output, (int64_t)nccl_comm, + (int64_t)reinterpret_cast(nccl_stream_v)); } // Shared expert (same as before, with fused SwiGLU) auto shared_gate = at::matmul(flat, shared_gate_proj.t()); auto shared_up = at::matmul(flat, shared_up_proj.t()); - auto shared_out = at::matmul(fused_swiglu_op(shared_gate, shared_up, 0.0), shared_down_proj.t()); + if (shared_gate_lora) { + shared_gate = add_batched_lora( + shared_gate.reshape({batch, seq, -1}), hidden, shared_gate_lora) + .reshape({batch * seq, -1}); + } + if (shared_up_lora) { + shared_up = add_batched_lora( + shared_up.reshape({batch, seq, -1}), hidden, shared_up_lora) + .reshape({batch * seq, -1}); + } + auto shared_hidden = fused_swiglu_op( + shared_gate.reshape({batch, seq, -1}), + shared_up.reshape({batch, seq, -1}), 0.0); + auto shared_out = at::matmul(shared_hidden.reshape({batch * seq, -1}), shared_down_proj.t()); + if (shared_down_lora) { + shared_out = add_batched_lora( + shared_out.reshape({batch, seq, -1}), shared_hidden, shared_down_lora) + .reshape({batch * seq, -1}); + } auto seg = at::sigmoid(at::matmul(flat, shared_expert_gate_w.t())).to(compute_type); shared_out = (shared_out * seg).to(compute_type); @@ -772,6 +1216,75 @@ static inline int64_t weight_count_for_layer(const LayerConfig& cfg) { return 2 + attn_w + mlp_w; } +enum class LoraSegment : uint8_t { Attention, Mlp }; + +struct LoraProjectionSpec { + const char* name; + int64_t weight_index; + LoraSegment segment; + bool grouped_expert; +}; + +struct LoraProjectionTable { + std::array entries; + int64_t count = 0; + + void add(const char* name, int64_t weight_index, LoraSegment segment) { + TORCH_CHECK(count < (int64_t)entries.size(), "too many LoRA projections in layer"); + entries[count++] = {name, weight_index, segment, false}; + } + + void add_grouped_expert(const char* name, int64_t weight_index) { + TORCH_CHECK(count < (int64_t)entries.size(), "too many LoRA projections in layer"); + entries[count++] = {name, weight_index, LoraSegment::Mlp, true}; + } +}; + +static LoraProjectionTable lora_projection_table(const LayerConfig& cfg) { + LoraProjectionTable table; + if (cfg.layer_type == 0) { + table.add("q_proj", 2, LoraSegment::Attention); + table.add("k_proj", 4, LoraSegment::Attention); + table.add("v_proj", 6, LoraSegment::Attention); + table.add("o_proj", 7, LoraSegment::Attention); + } else { + table.add("in_proj_qkv", 2, LoraSegment::Attention); + table.add("in_proj_z", 3, LoraSegment::Attention); + table.add("in_proj_a", 4, LoraSegment::Attention); + table.add("in_proj_b", 5, LoraSegment::Attention); + table.add("out_proj", 10, LoraSegment::Attention); + } + + const int64_t mlp_start = cfg.layer_type == 0 ? 8 : 11; + if (cfg.num_experts > 0) { + table.add("shared_gate_proj", mlp_start + 2, LoraSegment::Mlp); + table.add("shared_up_proj", mlp_start + 3, LoraSegment::Mlp); + table.add("shared_down_proj", mlp_start + 4, LoraSegment::Mlp); + table.add_grouped_expert("experts_gate_up_proj", mlp_start + 5); + table.add_grouped_expert("experts_down_proj", mlp_start + 6); + } else { + table.add("gate_proj", mlp_start, LoraSegment::Mlp); + table.add("up_proj", mlp_start + 1, LoraSegment::Mlp); + table.add("down_proj", mlp_start + 2, LoraSegment::Mlp); + } + return table; +} + +static inline int64_t lora_pair_count(const LayerConfig& cfg) { + return lora_projection_table(cfg).count; +} + +static int64_t lora_pair_index(const LayerConfig& cfg, const char* name) { + auto table = lora_projection_table(cfg); + for (int64_t i = 0; i < table.count; ++i) { + if (std::strcmp(table.entries[i].name, name) == 0) return i; + } + return -1; +} + +static RoutedExpertLora routed_expert_lora( + TrainingContext* ctx, int64_t layer_idx, const LayerConfig& cfg); + static at::Tensor forward_single_layer( TrainingContext* ctx, const at::Tensor& hidden, at::Tensor** w, const LayerConfig* cfg, int64_t layer_idx, at::ScalarType kind, @@ -807,13 +1320,42 @@ static at::Tensor forward_single_layer( } auto post_attn = rms_norm(hidden + attn_output, post_norm, cfg->rms_eps); if (is_moe) { + const int64_t shared_gate_pair = lora_pair_index(*cfg, "shared_gate_proj"); + const int64_t shared_up_pair = lora_pair_index(*cfg, "shared_up_proj"); + const int64_t shared_down_pair = lora_pair_index(*cfg, "shared_down_proj"); + const int64_t expert_gate_up_pair = lora_pair_index(*cfg, "experts_gate_up_proj"); + const int64_t expert_down_pair = lora_pair_index(*cfg, "experts_down_proj"); + auto shared_gate = use_batched ? *w[10] + : apply_multi_lora(ctx, layer_idx, shared_gate_pair, *w[10]); + auto shared_up = use_batched ? *w[11] + : apply_multi_lora(ctx, layer_idx, shared_up_pair, *w[11]); + auto shared_down = use_batched ? *w[12] + : apply_multi_lora(ctx, layer_idx, shared_down_pair, *w[12]); + auto expert_lora = routed_expert_lora(ctx, layer_idx, *cfg); auto mlp_out = moe_forward(cfg->nccl_comm, cfg->nccl_stream, post_attn, - *w[8], *w[9], *w[10], *w[11], *w[12], *w[13], *w[14], + *w[8], *w[9], shared_gate, shared_up, shared_down, *w[13], *w[14], + expert_lora, cfg->num_experts, cfg->top_k, cfg->moe_intermediate, - cfg->norm_topk_prob != 0, cfg->expert_start, cfg->expert_count, kind); + cfg->norm_topk_prob != 0, cfg->expert_start, cfg->expert_count, kind, + use_batched ? lora_batch_entry(ctx, layer_idx, shared_gate_pair) : nullptr, + use_batched ? lora_batch_entry(ctx, layer_idx, shared_up_pair) : nullptr, + use_batched ? lora_batch_entry(ctx, layer_idx, shared_down_pair) : nullptr, + use_batched ? lora_batch_entry(ctx, layer_idx, expert_gate_up_pair) : nullptr, + use_batched ? lora_batch_entry(ctx, layer_idx, expert_down_pair) : nullptr); return hidden + attn_output + mlp_out; } else { - auto mlp_out = dense_mlp_forward(post_attn, *w[8], *w[9], *w[10], kind); + if (use_batched) { + auto mlp_out = dense_mlp_forward_batched( + ctx, layer_idx, post_attn, *w[8], *w[9], *w[10], kind); + return hidden + attn_output + mlp_out; + } + auto gate = apply_multi_lora(ctx, layer_idx, + lora_pair_index(*cfg, "gate_proj"), *w[8]); + auto up = apply_multi_lora(ctx, layer_idx, + lora_pair_index(*cfg, "up_proj"), *w[9]); + auto down = apply_multi_lora(ctx, layer_idx, + lora_pair_index(*cfg, "down_proj"), *w[10]); + auto mlp_out = dense_mlp_forward(post_attn, gate, up, down, kind); return hidden + attn_output + mlp_out; } } else { @@ -830,7 +1372,9 @@ static at::Tensor forward_single_layer( } else { in_proj_qkv = apply_multi_lora(ctx, layer_idx, 0, in_proj_qkv); in_proj_z = apply_multi_lora(ctx, layer_idx, 1, in_proj_z); - out_proj = apply_multi_lora(ctx, layer_idx, 2, out_proj); + in_proj_a = apply_multi_lora(ctx, layer_idx, 2, in_proj_a); + in_proj_b = apply_multi_lora(ctx, layer_idx, 3, in_proj_b); + out_proj = apply_multi_lora(ctx, layer_idx, 4, out_proj); attn_output = linear_attention(attn_input, in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, a_log, dt_bias, conv1d_w, norm_w, out_proj, cfg->num_k_heads, cfg->key_dim, cfg->num_v_heads, cfg->val_dim, @@ -838,13 +1382,42 @@ static at::Tensor forward_single_layer( } auto post_attn = rms_norm(hidden + attn_output, post_norm, cfg->rms_eps); if (is_moe) { + const int64_t shared_gate_pair = lora_pair_index(*cfg, "shared_gate_proj"); + const int64_t shared_up_pair = lora_pair_index(*cfg, "shared_up_proj"); + const int64_t shared_down_pair = lora_pair_index(*cfg, "shared_down_proj"); + const int64_t expert_gate_up_pair = lora_pair_index(*cfg, "experts_gate_up_proj"); + const int64_t expert_down_pair = lora_pair_index(*cfg, "experts_down_proj"); + auto shared_gate = use_batched ? *w[13] + : apply_multi_lora(ctx, layer_idx, shared_gate_pair, *w[13]); + auto shared_up = use_batched ? *w[14] + : apply_multi_lora(ctx, layer_idx, shared_up_pair, *w[14]); + auto shared_down = use_batched ? *w[15] + : apply_multi_lora(ctx, layer_idx, shared_down_pair, *w[15]); + auto expert_lora = routed_expert_lora(ctx, layer_idx, *cfg); auto mlp_out = moe_forward(cfg->nccl_comm, cfg->nccl_stream, post_attn, - *w[11], *w[12], *w[13], *w[14], *w[15], *w[16], *w[17], + *w[11], *w[12], shared_gate, shared_up, shared_down, *w[16], *w[17], + expert_lora, cfg->num_experts, cfg->top_k, cfg->moe_intermediate, - cfg->norm_topk_prob != 0, cfg->expert_start, cfg->expert_count, kind); + cfg->norm_topk_prob != 0, cfg->expert_start, cfg->expert_count, kind, + use_batched ? lora_batch_entry(ctx, layer_idx, shared_gate_pair) : nullptr, + use_batched ? lora_batch_entry(ctx, layer_idx, shared_up_pair) : nullptr, + use_batched ? lora_batch_entry(ctx, layer_idx, shared_down_pair) : nullptr, + use_batched ? lora_batch_entry(ctx, layer_idx, expert_gate_up_pair) : nullptr, + use_batched ? lora_batch_entry(ctx, layer_idx, expert_down_pair) : nullptr); return hidden + attn_output + mlp_out; } else { - auto mlp_out = dense_mlp_forward(post_attn, *w[11], *w[12], *w[13], kind); + if (use_batched) { + auto mlp_out = dense_mlp_forward_batched( + ctx, layer_idx, post_attn, *w[11], *w[12], *w[13], kind); + return hidden + attn_output + mlp_out; + } + auto gate = apply_multi_lora(ctx, layer_idx, + lora_pair_index(*cfg, "gate_proj"), *w[11]); + auto up = apply_multi_lora(ctx, layer_idx, + lora_pair_index(*cfg, "up_proj"), *w[12]); + auto down = apply_multi_lora(ctx, layer_idx, + lora_pair_index(*cfg, "down_proj"), *w[13]); + auto mlp_out = dense_mlp_forward(post_attn, gate, up, down, kind); return hidden + attn_output + mlp_out; } } @@ -908,17 +1481,13 @@ struct TrainingContext { // ── Batched Multi-LoRA (activation-level) ── // When active, replaces the weight-level lora_cache. // Stores per-(layer, module) stacked A/B tensors for batched B@(A@x) computation. - struct LoraBatchEntry { - at::Tensor a_stack; // [N, rank, in] - at::Tensor b_stack; // [N, out, rank] - at::Tensor scaling; // [N, 1, 1] — per-adapter alpha/rank - }; bool lora_batch_valid = false; int64_t lora_batch_n = 0; // number of adapters in current batch std::map lora_batch_cache; // Legacy single-LoRA (backward compat) std::vector lora_a; std::vector lora_b; + std::vector lora_active; std::vector lora_layer_offset; double lora_scaling; std::vector lora_names; @@ -939,6 +1508,7 @@ struct TrainingContext { // MTP weights (optional) bool has_mtp; + double mtp_loss_scale = 0.1; // NVIDIA Megatron default at::Tensor *mtp_fc, *mtp_pre_fc_norm_emb, *mtp_pre_fc_norm_hidden, *mtp_norm; std::vector mtp_layer_weights; std::vector mtp_layer_configs; @@ -957,9 +1527,118 @@ struct TrainingContext { cudaStream_t nccl_stream = nullptr; int ep_world_size = 1; int ep_rank = 0; + int cuda_device = 0; // ────────────────────────────────────────────────────────────────────── }; +static const char* lora_pair_name(const LayerConfig& cfg, int64_t pair_idx) { + auto table = lora_projection_table(cfg); + TORCH_CHECK(pair_idx >= 0 && pair_idx < table.count, "invalid LoRA projection index"); + return table.entries[pair_idx].name; +} + +static constexpr int64_t LORA_CACHE_STRIDE = 32; + +static inline int64_t lora_cache_key(int64_t layer_idx, int64_t pair_idx) { + return layer_idx * LORA_CACHE_STRIDE + pair_idx; +} + +static inline bool legacy_lora_slot_active(const TrainingContext* ctx, int64_t slot) { + return slot >= 0 && slot < (int64_t)ctx->lora_active.size() && ctx->lora_active[slot] != 0; +} + +static RoutedExpertLora routed_expert_lora( + TrainingContext* ctx, int64_t layer_idx, const LayerConfig& cfg +) { + RoutedExpertLora result; + // Dynamic multi-LoRA supplies per-sample activation-level tensors. Do not + // mix the fixed adapter's expert tensors into those batches. + if (!ctx->adapters.empty() || cfg.num_experts <= 0) return result; + + const int64_t offset = ctx->lora_layer_offset[layer_idx]; + const int64_t gate_up_pair = lora_pair_index(cfg, "experts_gate_up_proj"); + const int64_t down_pair = lora_pair_index(cfg, "experts_down_proj"); + if (gate_up_pair >= 0 && legacy_lora_slot_active(ctx, offset + gate_up_pair)) { + result.gate_up_a = &ctx->lora_a[offset + gate_up_pair]; + result.gate_up_b = &ctx->lora_b[offset + gate_up_pair]; + } + if (down_pair >= 0 && legacy_lora_slot_active(ctx, offset + down_pair)) { + result.down_a = &ctx->lora_a[offset + down_pair]; + result.down_b = &ctx->lora_b[offset + down_pair]; + } + result.scaling = ctx->lora_scaling; + return result; +} + +static ncclDataType_t nccl_dtype_for(const at::Tensor& tensor) { + switch (tensor.scalar_type()) { + case at::kBFloat16: return ncclBfloat16; + case at::kFloat: return ncclFloat; + case at::kHalf: return ncclFloat16; + default: + TORCH_CHECK(false, "unsupported LoRA gradient dtype for EP all-reduce: ", + tensor.scalar_type()); + } +} + +static void allreduce_lora_grad(TrainingContext* ctx, at::Tensor& param) { + auto grad = param.grad(); + if (!ctx->nccl_comm || !grad.defined()) return; + auto contiguous = grad.contiguous(); + auto reduced = at::empty_like(contiguous); + int dev = contiguous.device().index(); + cudaSetDevice(dev); + auto stream = c10::cuda::getCurrentCUDAStream(dev).stream(); + auto err = ncclAllReduce( + contiguous.data_ptr(), reduced.data_ptr(), contiguous.numel(), + nccl_dtype_for(contiguous), ncclSum, ctx->nccl_comm, stream); + TORCH_CHECK(err == ncclSuccess, "NCCL LoRA gradient all-reduce failed: ", + ncclGetErrorString(err)); + param.mutable_grad() = reduced; +} + +// EP's routed expert output is summed across ranks in forward. The resulting +// upstream gradient must likewise be summed before replicated LoRA Adam steps; +// otherwise each rank updates attention/linear adapters from only its shard. +static void synchronize_lora_gradients(TrainingContext* ctx) { + if (!ctx->nccl_comm) return; + for (auto& adapter : ctx->adapters) { + for (auto& [layer_idx, pairs] : adapter.params) { + auto table = lora_projection_table(ctx->layer_configs[layer_idx]); + for (int64_t pair = 0; pair < (int64_t)pairs.size(); ++pair) { + // Dynamic routed-expert tensors are sharded exactly like the + // base experts; only replicated adapter tensors are reduced. + if (table.entries[pair].grouped_expert) continue; + auto& [a, b] = pairs[pair]; + allreduce_lora_grad(ctx, a); + allreduce_lora_grad(ctx, b); + } + } + } + for (int64_t layer = 0; layer < ctx->num_layers; ++layer) { + auto table = lora_projection_table(ctx->layer_configs[layer]); + int64_t offset = ctx->lora_layer_offset[layer]; + for (int64_t pair = 0; pair < table.count; ++pair) { + // Routed expert LoRA is sharded with the base expert weights. Its + // local gradients belong only to this EP rank and must not be + // summed with a different expert shard on another rank. + if (table.entries[pair].grouped_expert) continue; + allreduce_lora_grad(ctx, ctx->lora_a[offset + pair]); + allreduce_lora_grad(ctx, ctx->lora_b[offset + pair]); + } + } +} + +static void elide_trivial_attention_mask(TrainingContext* ctx) { + if (!ctx->attention_mask.defined() || ctx->attention_mask.numel() == 0) return; + // A padding-free batch can use SDPA's native causal fast path. This is one + // scalar synchronization per step, instead of materializing [B,S,S] in + // every full-attention layer. + if (at::all(ctx->attention_mask != 0).item()) { + ctx->attention_mask = at::Tensor(); + } +} + // ── Multi-LoRA: concat all adapters' A/B, 2x GEMM ── // Pre-build cache of concatenated A/B per (layer, module) pair. // Called once at start of forward; reused across all layers. @@ -977,16 +1656,25 @@ static void precompute_lora_cache(TrainingContext* ctx) { std::vector entries; for (int64_t layer_idx = 0; layer_idx < ctx->num_layers; layer_idx++) { - int64_t num_pairs = (ctx->layer_configs[layer_idx].layer_type == 0) ? 4 : 3; + int64_t num_pairs = lora_pair_count(ctx->layer_configs[layer_idx]); for (int64_t pair_idx = 0; pair_idx < num_pairs; pair_idx++) { + auto projection_table = lora_projection_table(ctx->layer_configs[layer_idx]); + // Routed experts use activation-level low-rank GEMMs. Materializing + // B@A for every local expert would erase LoRA's memory advantage. + if (projection_table.entries[pair_idx].grouped_expert) continue; std::vector a_list, b_list; + const char* module_name = lora_pair_name(ctx->layer_configs[layer_idx], pair_idx); for (auto& adapter : ctx->adapters) { + if (!adapter.target_modules.empty() && + adapter.target_modules.find(module_name) == adapter.target_modules.end()) + continue; if (!adapter.target_layers.empty() && adapter.target_layers.find(layer_idx) == adapter.target_layers.end()) continue; auto it = adapter.params.find(layer_idx); if (it == adapter.params.end()) continue; if (pair_idx >= (int64_t)it->second.size()) continue; auto& [a, b] = it->second[pair_idx]; + if (!a.requires_grad() && !b.requires_grad()) continue; double scaling = adapter.alpha / (double)adapter.rank; b_list.push_back(b * scaling); a_list.push_back(a); @@ -995,6 +1683,10 @@ static void precompute_lora_cache(TrainingContext* ctx) { if (!ctx->lora_a.empty() && layer_idx < (int64_t)ctx->lora_layer_offset.size()) { int64_t la_offset = ctx->lora_layer_offset[layer_idx]; if (la_offset + pair_idx < (int64_t)ctx->lora_a.size()) { + if (!ctx->lora_active.empty() && + !ctx->lora_active[la_offset + pair_idx]) { + continue; + } b_list.push_back(ctx->lora_b[la_offset + pair_idx] * ctx->lora_scaling); a_list.push_back(ctx->lora_a[la_offset + pair_idx]); } @@ -1009,7 +1701,7 @@ static void precompute_lora_cache(TrainingContext* ctx) { a_concat = at::cat(a_list, 0); // [sum_ranks, in] b_concat = at::cat(b_list, 1); // [out, sum_ranks] } - entries.push_back({layer_idx * 10 + pair_idx, a_concat, b_concat}); + entries.push_back({lora_cache_key(layer_idx, pair_idx), a_concat, b_concat}); } } } @@ -1076,22 +1768,26 @@ static void precompute_lora_cache(TrainingContext* ctx) { /// Stores in ctx->lora_batch_cache. Called once before forward. /// Replaces precompute_lora_cache when N > 1. static void prepare_lora_batch(TrainingContext* ctx) { - if (ctx->lora_batch_valid) return; ctx->lora_batch_cache.clear(); ctx->lora_batch_n = 0; for (int64_t layer_idx = 0; layer_idx < ctx->num_layers; layer_idx++) { - int64_t num_pairs = (ctx->layer_configs[layer_idx].layer_type == 0) ? 4 : 3; + int64_t num_pairs = lora_pair_count(ctx->layer_configs[layer_idx]); for (int64_t pair_idx = 0; pair_idx < num_pairs; pair_idx++) { std::vector a_list, b_list; std::vector scalings; + const char* module_name = lora_pair_name(ctx->layer_configs[layer_idx], pair_idx); for (auto& adapter : ctx->adapters) { + if (!adapter.target_modules.empty() && + adapter.target_modules.find(module_name) == adapter.target_modules.end()) + continue; if (!adapter.target_layers.empty() && adapter.target_layers.find(layer_idx) == adapter.target_layers.end()) continue; auto it = adapter.params.find(layer_idx); if (it == adapter.params.end()) continue; if (pair_idx >= (int64_t)it->second.size()) continue; auto& [a, b] = it->second[pair_idx]; + if (!a.requires_grad() && !b.requires_grad()) continue; b_list.push_back(b); a_list.push_back(a); scalings.push_back(adapter.alpha / (double)adapter.rank); @@ -1111,7 +1807,7 @@ static void prepare_lora_batch(TrainingContext* ctx) { ); auto scaling = scaling_cpu.to(a_stack.device()).to(at::kBFloat16); // [N, 1, 1] - ctx->lora_batch_cache[layer_idx * 10 + pair_idx] = { + ctx->lora_batch_cache[lora_cache_key(layer_idx, pair_idx)] = { a_stack, b_stack, scaling }; } @@ -1140,12 +1836,43 @@ static at::Tensor lora_activation_delta( return delta * s_c; } +static const LoraBatchEntry* lora_batch_entry( + TrainingContext* ctx, int64_t layer_idx, int64_t pair_idx +) { + if (!ctx || !ctx->lora_batch_valid || pair_idx < 0) return nullptr; + auto it = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, pair_idx)); + return it == ctx->lora_batch_cache.end() ? nullptr : &it->second; +} + +static at::Tensor dense_mlp_forward_batched( + TrainingContext* ctx, int64_t layer_idx, const at::Tensor& hidden, + const at::Tensor& gate_proj, const at::Tensor& up_proj, + const at::Tensor& down_proj, at::ScalarType compute_type +) { + const auto& cfg = ctx->layer_configs[layer_idx]; + const int64_t gate_pair = lora_pair_index(cfg, "gate_proj"); + const int64_t up_pair = lora_pair_index(cfg, "up_proj"); + const int64_t down_pair = lora_pair_index(cfg, "down_proj"); + + auto gate_out = at::matmul(hidden, gate_proj.t()); + auto up_out = at::matmul(hidden, up_proj.t()); + gate_out = add_batched_lora( + gate_out, hidden, lora_batch_entry(ctx, layer_idx, gate_pair)); + up_out = add_batched_lora( + up_out, hidden, lora_batch_entry(ctx, layer_idx, up_pair)); + auto activated = fused_swiglu_op(gate_out, up_out, 0.0); + auto result = at::matmul(activated, down_proj.t()); + result = add_batched_lora( + result, activated, lora_batch_entry(ctx, layer_idx, down_pair)); + return result.to(compute_type); +} + __attribute__((noinline, visibility("default"))) at::Tensor apply_multi_lora( TrainingContext* ctx, int64_t layer_idx, int64_t pair_idx, const at::Tensor& base_weight ) { - auto it = ctx->lora_cache.find(layer_idx * 10 + pair_idx); + auto it = ctx->lora_cache.find(lora_cache_key(layer_idx, pair_idx)); if (it == ctx->lora_cache.end()) return base_weight; // Cached delta_weight = b_concat @ a_concat (BF16, precomputed in batched bmm) @@ -1194,7 +1921,7 @@ at::Tensor compute_attn_only( } // Legacy path: weight-level LoRA - int64_t lora_count = (cfg.layer_type == 0) ? 4 : 3; + int64_t lora_count = lora_pair_count(cfg); int64_t la_offset = ctx->lora_layer_offset[layer_idx]; bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size(); std::vector la(lora_count, nullptr), lb(lora_count, nullptr); @@ -1223,7 +1950,9 @@ at::Tensor compute_attn_only( if (has_lora) { if (la[0]) qkv = lora_delta(qkv, *la[0], *lb[0], ctx->lora_scaling); if (la[1]) z = lora_delta(z, *la[1], *lb[1], ctx->lora_scaling); - if (la[2]) op = lora_delta(op, *la[2], *lb[2], ctx->lora_scaling); + if (la[2]) a = lora_delta(a, *la[2], *lb[2], ctx->lora_scaling); + if (la[3]) b = lora_delta(b, *la[3], *lb[3], ctx->lora_scaling); + if (la[4]) op = lora_delta(op, *la[4], *lb[4], ctx->lora_scaling); } return linear_attention(attn_input, qkv, z, a, b, al, db, cw, nw, op, cfg.num_k_heads, cfg.key_dim, cfg.num_v_heads, cfg.val_dim, @@ -1242,17 +1971,50 @@ at::Tensor compute_mlp_only( int64_t mlp_start = (cfg.layer_type == 0) ? 8 : 11; auto post_attn = rms_norm(residual, *ctx->weight_ptrs[w_offset + 1], cfg.rms_eps); if (cfg.num_experts > 0) { + const bool use_batched = ctx->lora_batch_valid; + const int64_t shared_gate_pair = lora_pair_index(cfg, "shared_gate_proj"); + const int64_t shared_up_pair = lora_pair_index(cfg, "shared_up_proj"); + const int64_t shared_down_pair = lora_pair_index(cfg, "shared_down_proj"); + const int64_t expert_gate_up_pair = lora_pair_index(cfg, "experts_gate_up_proj"); + const int64_t expert_down_pair = lora_pair_index(cfg, "experts_down_proj"); + auto shared_gate = use_batched ? *ctx->weight_ptrs[w_offset+mlp_start+2] + : apply_multi_lora(ctx, layer_idx, shared_gate_pair, + *ctx->weight_ptrs[w_offset+mlp_start+2]); + auto shared_up = use_batched ? *ctx->weight_ptrs[w_offset+mlp_start+3] + : apply_multi_lora(ctx, layer_idx, shared_up_pair, + *ctx->weight_ptrs[w_offset+mlp_start+3]); + auto shared_down = use_batched ? *ctx->weight_ptrs[w_offset+mlp_start+4] + : apply_multi_lora(ctx, layer_idx, shared_down_pair, + *ctx->weight_ptrs[w_offset+mlp_start+4]); + auto expert_lora = routed_expert_lora(ctx, layer_idx, cfg); return moe_forward(cfg.nccl_comm, cfg.nccl_stream, post_attn, *ctx->weight_ptrs[w_offset+mlp_start], *ctx->weight_ptrs[w_offset+mlp_start+1], - *ctx->weight_ptrs[w_offset+mlp_start+2], *ctx->weight_ptrs[w_offset+mlp_start+3], - *ctx->weight_ptrs[w_offset+mlp_start+4], *ctx->weight_ptrs[w_offset+mlp_start+5], + shared_gate, shared_up, + shared_down, *ctx->weight_ptrs[w_offset+mlp_start+5], *ctx->weight_ptrs[w_offset+mlp_start+6], + expert_lora, cfg.num_experts, cfg.top_k, cfg.moe_intermediate, - cfg.norm_topk_prob != 0, cfg.expert_start, cfg.expert_count, kind); + cfg.norm_topk_prob != 0, cfg.expert_start, cfg.expert_count, kind, + use_batched ? lora_batch_entry(ctx, layer_idx, shared_gate_pair) : nullptr, + use_batched ? lora_batch_entry(ctx, layer_idx, shared_up_pair) : nullptr, + use_batched ? lora_batch_entry(ctx, layer_idx, shared_down_pair) : nullptr, + use_batched ? lora_batch_entry(ctx, layer_idx, expert_gate_up_pair) : nullptr, + use_batched ? lora_batch_entry(ctx, layer_idx, expert_down_pair) : nullptr); } else { - return dense_mlp_forward(post_attn, - *ctx->weight_ptrs[w_offset+mlp_start], *ctx->weight_ptrs[w_offset+mlp_start+1], - *ctx->weight_ptrs[w_offset+mlp_start+2], kind); + if (ctx->lora_batch_valid) { + return dense_mlp_forward_batched( + ctx, layer_idx, post_attn, + *ctx->weight_ptrs[w_offset+mlp_start], + *ctx->weight_ptrs[w_offset+mlp_start+1], + *ctx->weight_ptrs[w_offset+mlp_start+2], kind); + } + auto gate = apply_multi_lora(ctx, layer_idx, + lora_pair_index(cfg, "gate_proj"), *ctx->weight_ptrs[w_offset+mlp_start]); + auto up = apply_multi_lora(ctx, layer_idx, + lora_pair_index(cfg, "up_proj"), *ctx->weight_ptrs[w_offset+mlp_start+1]); + auto down = apply_multi_lora(ctx, layer_idx, + lora_pair_index(cfg, "down_proj"), *ctx->weight_ptrs[w_offset+mlp_start+2]); + return dense_mlp_forward(post_attn, gate, up, down, kind); } } @@ -1297,20 +2059,51 @@ struct SubLayerCkpt : public torch::autograd::Function { bool is_attn = ctx->saved_data["is_attn"].toBool(); at::AutoGradMode guard(true); input.set_requires_grad(true); + + // Derived LoRA deltas belong to one recomputed segment's graph and are + // released by grad(..., retain_graph=false). Attention and MLP both + // need a fresh cache now that both segments can own adapters. + tc->lora_batch_valid = false; + tc->lora_cache_valid = false; + if (!tc->adapters.empty()) prepare_lora_batch(tc); + else precompute_lora_cache(tc); auto output = is_attn ? compute_attn_only(tc, input, layer, tc->compute_type) : compute_mlp_only(tc, input, layer, tc->compute_type); // Collect all tensors to compute gradients for: input + LoRA params for this layer // This way grad() accumulates gradients into LoRA params (leaf nodes) too. - int64_t lora_count = (tc->layer_configs[layer].layer_type == 0) ? 4 : 3; + auto projection_table = lora_projection_table(tc->layer_configs[layer]); + int64_t lora_count = projection_table.count; int64_t la_offset = tc->lora_layer_offset[layer]; bool has_lora = (la_offset + lora_count) <= (int64_t)tc->lora_a.size(); std::vector grad_inputs = {input}; - if (has_lora) { + std::vector> active_params; + if (!tc->adapters.empty()) { + for (auto& adapter : tc->adapters) { + auto it = adapter.params.find(layer); + if (it == adapter.params.end()) continue; + for (int64_t k = 0; k < lora_count && k < (int64_t)it->second.size(); ++k) { + auto segment = projection_table.entries[k].segment; + if ((is_attn && segment != LoraSegment::Attention) || + (!is_attn && segment != LoraSegment::Mlp)) continue; + auto& [a, b] = it->second[k]; + if (!a.requires_grad() && !b.requires_grad()) continue; + active_params.push_back({&a, &b}); + grad_inputs.push_back(a); + grad_inputs.push_back(b); + } + } + } else if (has_lora) { for (int64_t k = 0; k < lora_count; k++) { - grad_inputs.push_back(tc->lora_a[la_offset + k]); - grad_inputs.push_back(tc->lora_b[la_offset + k]); + auto segment = projection_table.entries[k].segment; + if ((is_attn && segment != LoraSegment::Attention) || + (!is_attn && segment != LoraSegment::Mlp)) continue; + int64_t slot = la_offset + k; + if (!legacy_lora_slot_active(tc, slot)) continue; + active_params.push_back({&tc->lora_a[slot], &tc->lora_b[slot]}); + grad_inputs.push_back(tc->lora_a[slot]); + grad_inputs.push_back(tc->lora_b[slot]); } } @@ -1321,11 +2114,11 @@ struct SubLayerCkpt : public torch::autograd::Function { ); // Manually accumulate LoRA param gradients - if (has_lora) { + if (!active_params.empty()) { int64_t gi = 1; // skip input grad (index 0) - for (int64_t k = 0; k < lora_count; k++) { + for (auto& [param_a_ptr, param_b_ptr] : active_params) { if (grads[gi].defined()) { - auto& param_a = tc->lora_a[la_offset + k]; + auto& param_a = *param_a_ptr; if (param_a.grad().defined()) param_a.grad().add_(grads[gi]); else @@ -1333,7 +2126,7 @@ struct SubLayerCkpt : public torch::autograd::Function { } gi++; if (grads[gi].defined()) { - auto& param_b = tc->lora_b[la_offset + k]; + auto& param_b = *param_b_ptr; if (param_b.grad().defined()) param_b.grad().add_(grads[gi]); else @@ -1385,15 +2178,15 @@ static at::Tensor full_attention_batched( auto v = at::matmul(hidden, v_proj.t()); // Apply activation-level LoRA: q += B@(A@hidden) * scaling - auto it_q = ctx->lora_batch_cache.find(layer_idx * 10 + 0); + auto it_q = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 0)); if (it_q != ctx->lora_batch_cache.end()) { q = q + lora_activation_delta(hidden, it_q->second.a_stack, it_q->second.b_stack, it_q->second.scaling); } - auto it_k = ctx->lora_batch_cache.find(layer_idx * 10 + 1); + auto it_k = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 1)); if (it_k != ctx->lora_batch_cache.end()) { k = k + lora_activation_delta(hidden, it_k->second.a_stack, it_k->second.b_stack, it_k->second.scaling); } - auto it_v = ctx->lora_batch_cache.find(layer_idx * 10 + 2); + auto it_v = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 2)); if (it_v != ctx->lora_batch_cache.end()) { v = v + lora_activation_delta(hidden, it_v->second.a_stack, it_v->second.b_stack, it_v->second.scaling); } @@ -1410,8 +2203,8 @@ static at::Tensor full_attention_batched( q_out = rms_norm(q_out, q_norm, rms_eps); k = rms_norm(k, k_norm, rms_eps); - // Release gate before RoPE - gate = at::Tensor(); + // Keep gate alive until after SDPA. The architecture applies it to the + // attention value before the output projection. // RoPE int64_t rotary_dim = (int64_t)(head_dim * partial_rotary_factor); @@ -1457,12 +2250,14 @@ static at::Tensor full_attention_batched( } else { attn_out = at::scaled_dot_product_attention(q_out, k, v, c10::nullopt, 0.0, true, c10::nullopt, true); } - auto result = attn_out.transpose(1, 2).reshape({batch, seq, qkv_dim}).matmul(o_proj.t()); + auto gated_attn = attn_out * at::sigmoid(gate).to(attn_out.scalar_type()); + auto attn_flat = gated_attn.transpose(1, 2).reshape({batch, seq, qkv_dim}); + auto result = attn_flat.matmul(o_proj.t()); // Apply LoRA delta on o_proj output - auto it_o = ctx->lora_batch_cache.find(layer_idx * 10 + 3); + auto it_o = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 3)); if (it_o != ctx->lora_batch_cache.end()) { - result = result + lora_activation_delta(attn_out.transpose(1, 2).reshape({batch, seq, qkv_dim}), + result = result + lora_activation_delta(attn_flat, it_o->second.a_stack, it_o->second.b_stack, it_o->second.scaling); } return result; @@ -1487,7 +2282,7 @@ static at::Tensor linear_attention_batched( // QKV projection + LoRA delta auto qkv = at::matmul(hidden, in_proj_qkv.t()); - auto it_qkv = ctx->lora_batch_cache.find(layer_idx * 10 + 0); + auto it_qkv = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 0)); if (it_qkv != ctx->lora_batch_cache.end()) { qkv = qkv + lora_activation_delta(hidden, it_qkv->second.a_stack, it_qkv->second.b_stack, it_qkv->second.scaling); } @@ -1559,10 +2354,18 @@ static at::Tensor linear_attention_batched( auto a = at::matmul(hidden, in_proj_a.t()); auto b = at::matmul(hidden, in_proj_b.t()); + auto it_a = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 2)); + if (it_a != ctx->lora_batch_cache.end()) { + a = a + lora_activation_delta(hidden, it_a->second.a_stack, it_a->second.b_stack, it_a->second.scaling); + } + auto it_b = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 3)); + if (it_b != ctx->lora_batch_cache.end()) { + b = b + lora_activation_delta(hidden, it_b->second.a_stack, it_b->second.b_stack, it_b->second.scaling); + } // Z projection + LoRA delta auto z = at::matmul(hidden, in_proj_z.t()); - auto it_z = ctx->lora_batch_cache.find(layer_idx * 10 + 1); + auto it_z = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 1)); if (it_z != ctx->lora_batch_cache.end()) { z = z + lora_activation_delta(hidden, it_z->second.a_stack, it_z->second.b_stack, it_z->second.scaling); } @@ -1622,15 +2425,13 @@ static at::Tensor linear_attention_batched( int64_t sub_batch = (BH_total > 8192) ? 256 : batch; // 256 adapters or all if small sub_batch = std::min(sub_batch, batch); - auto outs = at::empty({BH_total, seq, head_v_dim}, q_t.options()); - auto delta_buf = at::empty({BH_total, seq, head_v_dim}, q_t.options()); + std::vector sub_outputs; + sub_outputs.reserve((batch + sub_batch - 1) / sub_batch); for (int64_t sb = 0; sb < batch; sb += sub_batch) { int64_t n = std::min(sub_batch, batch - sb); int64_t BH = n * num_v_heads; - auto state = at::zeros({BH, head_k_dim, head_v_dim}, q_t.options()); - // Narrow on dim 0 (batch/adapter dimension), then reshape to [BH, seq, dim] auto q_sub = q_t.narrow(0, sb, n); auto k_sub = k_t.narrow(0, sb, n); @@ -1643,24 +2444,12 @@ static at::Tensor linear_attention_batched( auto v_contig = v_sub.reshape({BH, seq, head_v_dim}).contiguous().to(at::kFloat); auto g_contig = g_sub.reshape({BH, seq}).contiguous().to(at::kFloat); auto beta_contig = beta_sub.reshape({BH, seq}).contiguous().to(at::kFloat); - auto state_contig = state.contiguous(); - - // Output slices: [BH, seq, val_dim] → need to narrow on dim 0 too - // outs is [BH_total, seq, val_dim], so bh_start = sb * num_v_heads - int64_t bh_start = sb * num_v_heads; - cuda_gated_delta_rule( - q_contig.data_ptr(), - k_contig.data_ptr(), - v_contig.data_ptr(), - g_contig.data_ptr(), - beta_contig.data_ptr(), - state_contig.data_ptr(), - outs.narrow(0, bh_start, BH).data_ptr(), - delta_buf.narrow(0, bh_start, BH).data_ptr(), - (int)BH, (int)seq, (int)head_k_dim, (int)head_v_dim - ); + sub_outputs.push_back(GatedDeltaRuleFunction::apply( + q_contig, k_contig, v_contig, g_contig, beta_contig)); } + auto outs = at::cat(sub_outputs, 0); + auto core_out = outs.reshape({batch, num_v_heads, seq, head_v_dim}) .transpose(1, 2).to(compute_type); @@ -1688,7 +2477,7 @@ static at::Tensor linear_attention_batched( } // out_proj LoRA delta: result += B@(A@gated) * scaling - auto it_op = ctx->lora_batch_cache.find(layer_idx * 10 + 2); + auto it_op = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 4)); if (it_op != ctx->lora_batch_cache.end()) { result = result + lora_activation_delta(gated, it_op->second.a_stack, it_op->second.b_stack, it_op->second.scaling); } @@ -1743,7 +2532,7 @@ static at::Tensor forward_full( ctx->weight_ptrs.begin() + w_offset + w_count); // Get LoRA pointers for this layer (nullptr if no LoRA for this layer) - int64_t lora_count = (ctx->layer_configs[i].layer_type == 0) ? 4 : 3; + int64_t lora_count = lora_pair_count(ctx->layer_configs[i]); int64_t la_offset = ctx->lora_layer_offset[i]; // Check if this layer has LoRA params (la_offset < lora_a.size()) bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size(); @@ -1797,7 +2586,7 @@ static at::Tensor forward_layer_group( std::vector layer_w(ctx->weight_ptrs.begin() + w_offset, ctx->weight_ptrs.begin() + w_offset + w_count); - int64_t lora_count = (ctx->layer_configs[i].layer_type == 0) ? 4 : 3; + int64_t lora_count = lora_pair_count(ctx->layer_configs[i]); int64_t la_offset = ctx->lora_layer_offset[i]; bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size(); std::vector la_ptrs(lora_count, nullptr), lb_ptrs(lora_count, nullptr); @@ -1922,7 +2711,7 @@ struct FusedLayerFunction : public torch::autograd::Function int64_t w_count = weight_count_for_layer(tc->layer_configs[layer_idx]); std::vector layer_w(tc->weight_ptrs.begin() + w_offset, tc->weight_ptrs.begin() + w_offset + w_count); - int64_t lora_count = (tc->layer_configs[layer_idx].layer_type == 0) ? 4 : 3; + int64_t lora_count = lora_pair_count(tc->layer_configs[layer_idx]); int64_t la_offset = tc->lora_layer_offset[layer_idx]; bool has_lora = (la_offset + lora_count) <= (int64_t)tc->lora_a.size(); std::vector la(lora_count, nullptr), lb(lora_count, nullptr); @@ -1992,7 +2781,7 @@ static at::Tensor forward_full_checkpoint( h_f[0][0][0].item(), h_f[0][0][1].item(), h_f[0][0][2].item()); } - bool use_subckpt = getenv("QWEN36_SUBCKPT"); + bool use_subckpt = env_enabled("QWEN36_SUBCKPT"); if (use_subckpt) { at::AutoGradMode restore(true); @@ -2040,8 +2829,6 @@ static at::Tensor forward_full_checkpoint( } hidden = forward_layer_group(ctx, hidden, start, end); - // emptyCache: needed for seq>4096 (memory pressure), skip for small seq (async benefit). - if (hidden.size(1) > 4096) c10::cuda::CUDACachingAllocator::emptyCache(); if (getenv("QWEN36_DUMP_LAYERS")) { auto h_f = hidden.to(at::kFloat); @@ -2068,14 +2855,6 @@ static void manual_group_backward( at::Tensor grad = hidden_grad; at::AutoGradMode grad_mode(true); - ctx->lora_batch_valid = false; - ctx->lora_cache_valid = false; - // Rebuild lora_batch_cache under AutoGradMode(true) — this is critical: - // forward pass built it in no-grad mode, so a_stack/b_stack didn't track - // gradients. Here we rebuild so at::stack(A_i) creates a differentiable - // graph: grad → bmm → stack → A_i (leaf). Without this, grads are undefined. - prepare_lora_batch(ctx); - for (int64_t g = num_groups - 1; g >= 0; g--) { int64_t start = groups[g].first; int64_t end = groups[g].second; @@ -2083,6 +2862,14 @@ static void manual_group_backward( // Restore input from saved (CPU if offloaded) auto input = ctx->group_inputs[g].to(hidden_grad.device()).detach().set_requires_grad(true); + // Each recomputed group builds a fresh autograd graph. Reusing a + // differentiable LoRA cache after grad(..., retain_graph=false) would + // point the next group at freed graph nodes. + ctx->lora_batch_valid = false; + ctx->lora_cache_valid = false; + if (!ctx->adapters.empty()) prepare_lora_batch(ctx); + else precompute_lora_cache(ctx); + // Recompute forward with grad for this group only auto output = forward_layer_group(ctx, input, start, end); @@ -2095,12 +2882,13 @@ static void manual_group_backward( if (ctx->lora_batch_valid) { // Multi-LoRA: collect A/B from ctx->adapters for (int64_t l = start; l < end; l++) { - int64_t lora_count = (ctx->layer_configs[l].layer_type == 0) ? 4 : 3; + int64_t lora_count = lora_pair_count(ctx->layer_configs[l]); for (auto& adapter : ctx->adapters) { auto it = adapter.params.find(l); if (it == adapter.params.end()) continue; for (int64_t k = 0; k < lora_count && k < (int64_t)it->second.size(); k++) { auto& [a, b] = it->second[k]; + if (!a.requires_grad() && !b.requires_grad()) continue; grad_inputs.push_back(a); grad_inputs.push_back(b); } @@ -2109,11 +2897,12 @@ static void manual_group_backward( } else { // Legacy single-LoRA for (int64_t l = start; l < end; l++) { - int64_t lora_count = (ctx->layer_configs[l].layer_type == 0) ? 4 : 3; + int64_t lora_count = lora_pair_count(ctx->layer_configs[l]); int64_t la_offset = ctx->lora_layer_offset[l]; bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size(); if (has_lora) { for (int64_t k = 0; k < lora_count; k++) { + if (!legacy_lora_slot_active(ctx, la_offset + k)) continue; grad_inputs.push_back(ctx->lora_a[la_offset + k]); grad_inputs.push_back(ctx->lora_b[la_offset + k]); } @@ -2133,12 +2922,13 @@ static void manual_group_backward( // Multi-LoRA: accumulate into ctx->adapters int64_t gi = 1; // skip input grad (index 0) for (int64_t l = start; l < end; l++) { - int64_t lora_count = (ctx->layer_configs[l].layer_type == 0) ? 4 : 3; + int64_t lora_count = lora_pair_count(ctx->layer_configs[l]); for (auto& adapter : ctx->adapters) { auto it = adapter.params.find(l); if (it == adapter.params.end()) continue; for (int64_t k = 0; k < lora_count && k < (int64_t)it->second.size(); k++) { auto& [a, b] = it->second[k]; + if (!a.requires_grad() && !b.requires_grad()) continue; if (gi < (int64_t)grads.size() && grads[gi].defined()) { if (a.grad().defined()) a.grad().add_(grads[gi]); else a.mutable_grad() = grads[gi].clone(); @@ -2156,11 +2946,12 @@ static void manual_group_backward( // Legacy single-LoRA int64_t gi = 1; // skip input grad (index 0) for (int64_t l = start; l < end; l++) { - int64_t lora_count = (ctx->layer_configs[l].layer_type == 0) ? 4 : 3; + int64_t lora_count = lora_pair_count(ctx->layer_configs[l]); int64_t la_offset = ctx->lora_layer_offset[l]; bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size(); if (has_lora) { for (int64_t k = 0; k < lora_count; k++) { + if (!legacy_lora_slot_active(ctx, la_offset + k)) continue; if (grads[gi].defined()) { auto& pa = ctx->lora_a[la_offset + k]; if (pa.grad().defined()) pa.grad().add_(grads[gi]); @@ -2178,9 +2969,6 @@ static void manual_group_backward( } } - // emptyCache for backward groups when seq>4096. - if (hidden_grad.size(1) > 4096) c10::cuda::CUDACachingAllocator::emptyCache(); - // Gradient for this group's input = gradient for next group's output grad = grads[0]; @@ -2286,13 +3074,13 @@ static at::Tensor compute_loss_fused( sum_exp = old_exp * sum_exp + tile_exp.sum(/*dim=*/1, /*keepdim=*/true); logit_max = new_max; - c10::cuda::CUDACachingAllocator::emptyCache(); } // Phase 2: compute target logit for each token // We need logits[target] — gather the target position's logit. // Do a second pass to find target logit. - auto target_logit = at::zeros({total_tokens, 1}, + auto target_logit = at::full({total_tokens, 1}, + -std::numeric_limits::infinity(), at::TensorOptions().dtype(at::kFloat).device(hidden_flat.device())); for (int64_t t = 0; t < num_tiles; t++) { @@ -2308,23 +3096,36 @@ static at::Tensor compute_loss_fused( auto logits_tile = at::matmul(hidden_flat, lm_head_tile.t()).to(at::kFloat); // Gather target logits: subtract v_start to get local index - auto local_targets = (shifted_targets - v_start).clamp_min(0); + auto local_targets = (shifted_targets - v_start).clamp(0, v_n - 1); // Gather: logits_tile[i, local_targets[i]] for tokens in range auto gathered = at::gather(logits_tile, /*dim=*/1, local_targets.reshape({-1, 1})); // Only keep tokens that are actually in range - gathered = gathered * in_range.to(at::kFloat).reshape({-1, 1}); + gathered = at::where( + in_range.reshape({-1, 1}), gathered, + at::full_like(gathered, -std::numeric_limits::infinity())); target_logit = at::max(target_logit, gathered); - c10::cuda::CUDACachingAllocator::emptyCache(); } // Loss per token: log(sum_exp) - target_logit (= -log(softmax[target])) auto log_sum_exp = at::log(sum_exp) + logit_max; // log(sum(exp(x-max))) + max = logsumexp auto per_token_loss = log_sum_exp - target_logit; // [total_tokens, 1] + per_token_loss = at::where( + mask_f.reshape({-1, 1}) > 0, per_token_loss, + at::zeros_like(per_token_loss)); auto masked_loss = per_token_loss.squeeze(1) * mask_f; auto total_count = mask_f.sum().clamp_min(1.0); double loss_val = (masked_loss.sum().item()) / total_count.item(); + // Evaluation runs with AutoGradMode disabled. The online-softmax passes + // above already produced the exact scalar value, so do not execute the + // training-only manual gradient pass below. Besides fixing eval on a + // no-grad graph, this avoids a third traversal of the vocabulary tiles. + if (!at::GradMode::is_enabled()) { + return at::tensor({loss_val}, + at::TensorOptions().dtype(at::kFloat).device(hidden.device())); + } + // ── Backward pass: compute grad_hidden_normed manually ── // dL/dhidden_normed = (softmax - one_hot) / count * mask // softmax = exp(logit - logit_max) / sum_exp @@ -2365,7 +3166,6 @@ static at::Tensor compute_loss_fused( ); grad_hidden.add_(grad_tile); - c10::cuda::CUDACachingAllocator::emptyCache(); } // Set gradient on hidden_normed (leaf tensor). @@ -2387,29 +3187,20 @@ static at::Tensor compute_loss_fused( hidden_normed_recompute.backward(hidden_normed.grad()); } - c10::cuda::CUDACachingAllocator::emptyCache(); - return at::tensor({loss_val}, at::TensorOptions().dtype(at::kFloat).device(hidden.device())); } -// ────────────────────────────────────────────────────────────────────── -// Conditional emptyCache — only syncs when GPU memory is running low. -// This avoids the ~2s/step of GPU idle time caused by 24 unconditional -// emptyCache calls (10 fwd groups + 10 bwd groups + 4 CE chunks). -// ────────────────────────────────────────────────────────────────────── -static void maybe_emptyCache(at::Device device, int64_t threshold_gb = 20) { - size_t free_mem, total_mem; - cudaMemGetInfo(&free_mem, &total_mem); - if ((int64_t)free_mem < threshold_gb * 1024 * 1024 * 1024) { - c10::cuda::CUDACachingAllocator::emptyCache(); - } -} +struct LossResult { + at::Tensor value; + at::Tensor hidden_grad; +}; // Cross-entropy loss with response-only masking — chunked with detach. -// Detach hidden_normed so CE backward doesn't traverse main model graph. -// Accumulate hidden_normed gradient, then backprop to hidden separately. -static at::Tensor compute_loss( +// Return the hidden gradient instead of mutating/backpropagating through the +// main model graph. The caller combines auxiliary gradients and performs one +// main backward, which is required for checkpointed execution. +static LossResult compute_loss( TrainingContext* ctx, const at::Tensor& hidden, const at::Tensor& input_ids, @@ -2421,7 +3212,7 @@ static at::Tensor compute_loss( // Detach hidden for CE computation — CE backward won't touch main model graph. // We accumulate gradient into hidden_normed, then manually backprop to hidden. - auto hidden_detached = hidden.detach(); + auto hidden_detached = hidden.detach().set_requires_grad(true); // Compute hidden_normed in no-grad, then set requires_grad on it. // This way CE backward only builds a tiny graph (hidden_normed → logits → loss), @@ -2448,7 +3239,6 @@ static at::Tensor compute_loss( auto hidden_flat = shifted_hidden.reshape({-1, hidden_normed.size(2)}); double total_loss_val = 0.0; - auto total_count_val = total_count.item(); for (int64_t c = 0; c < num_chunks; c++) { int64_t start = c * chunk_size; @@ -2483,7 +3273,10 @@ static at::Tensor compute_loss( at::Tensor(), at::Reduction::None, -100, 0.0 ); auto masked_loss = per_token_loss * chunk_mask.to(at::kFloat); - auto chunk_loss = masked_loss.sum(); + // Normalize every chunk by the global response-token count. Backward + // must match the mean returned to the caller, independent of sequence + // length or chunk boundaries. + auto chunk_loss = masked_loss.sum() / total_count; // Backward this chunk — each chunk creates an independent CE subgraph // because hidden_normed is a leaf tensor. retain_graph=false is safe @@ -2493,24 +3286,21 @@ static at::Tensor compute_loss( total_loss_val += chunk_loss.item(); - if (hidden_normed.size(1) > 4096) c10::cuda::CUDACachingAllocator::emptyCache(); - } - - // Backprop hidden_normed gradient to hidden via rms_norm. - // hidden_normed was computed in no-grad mode (detached from hidden_detached). - // Recompute rms_norm with grad tracking to get hidden's gradient. - if (hidden_normed.grad().defined()) { - hidden.set_requires_grad(true); - auto hidden_normed_recompute = rms_norm(hidden, final_norm, ctx->rms_eps); - hidden_normed_recompute.backward(hidden_normed.grad()); - // hidden.grad() now has the CE gradient contribution } - // emptyCache for CE intermediates when seq>4096. - if (hidden_normed.size(1) > 4096) c10::cuda::CUDACachingAllocator::emptyCache(); - - return at::tensor({total_loss_val / total_count_val}, - at::TensorOptions().dtype(at::kFloat).device(hidden.device())); + TORCH_CHECK(hidden_normed.grad().defined(), + "cross-entropy did not produce a hidden gradient"); + auto hidden_normed_recompute = rms_norm(hidden_detached, final_norm, ctx->rms_eps); + auto hidden_grad = torch::autograd::grad( + {hidden_normed_recompute}, {hidden_detached}, {hidden_normed.grad()}, + /*retain_graph=*/false, /*create_graph=*/false, + /*allow_unused=*/false)[0]; + + return { + at::tensor({total_loss_val}, + at::TensorOptions().dtype(at::kFloat).device(hidden.device())), + hidden_grad + }; } // ────────────────────────────────────────────────────────────────────── @@ -2607,10 +3397,12 @@ static at::Tensor mtp_compute_loss( /*ignore_index=*/-100, /*label_smoothing=*/0.0 ); auto masked_loss = per_token_loss * chunk_mask.to(at::kFloat); - total_loss += masked_loss.sum(); + // Avoid in-place accumulation into a non-grad leaf: out-of-place add + // keeps the MTP loss connected to the frozen-head input graph. + total_loss = total_loss + masked_loss.sum(); } - return (total_loss / total_count) * 0.5; + return (total_loss / total_count) * ctx->mtp_loss_scale; } // ────────────────────────────────────────────────────────────────────── @@ -2619,6 +3411,10 @@ static at::Tensor mtp_compute_loss( extern "C" { +__attribute__((visibility("default"))) int64_t qwen36_kernel_abi_version() { + return 4; +} + // Create training context — called once at startup // lora_rank: LoRA rank (from config) // target_layers: array of layer indices to apply LoRA (nullptr = all layers) @@ -2631,7 +3427,8 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( double lora_scaling, double lr, double beta1, double beta2, double eps, int64_t vocab_size, double rms_eps, int64_t lora_rank, - const int64_t* target_layers, int64_t num_target_layers + const int64_t* target_layers, int64_t num_target_layers, + const char* target_modules_str ) { try { auto* ctx = new TrainingContext(); @@ -2641,6 +3438,9 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( ctx->step_count = 0; ctx->lora_scaling = lora_scaling; ctx->num_layers = num_layers; ctx->use_checkpoint = false; ctx->group_size = 4; + if (const char* mtp_scale = getenv("QWEN36_MTP_LOSS_SCALE")) { + ctx->mtp_loss_scale = std::strtod(mtp_scale, nullptr); + } // Store weight pointers auto** wp = reinterpret_cast(weight_ptrs); @@ -2659,58 +3459,105 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( // Build target layer set std::set target_set; + const bool all_target_layers = !target_layers || num_target_layers == 0; if (target_layers && num_target_layers > 0) { - for (int64_t j = 0; j < num_target_layers; j++) + for (int64_t j = 0; j < num_target_layers; j++) { + TORCH_CHECK(target_layers[j] >= 0 && target_layers[j] < num_layers, + "LoRA target layer out of range: ", target_layers[j], + " for model with ", num_layers, " layers"); target_set.insert(target_layers[j]); + } + } + + std::set target_modules; + if (target_modules_str && target_modules_str[0] != '\0') { + std::stringstream ss(target_modules_str); + std::string item; + while (std::getline(ss, item, ',')) { + if (!item.empty()) target_modules.insert(item); + } + } + for (const auto& name : target_modules) { + TORCH_CHECK( + name == "q_proj" || name == "k_proj" || name == "v_proj" || + name == "o_proj" || name == "in_proj_qkv" || + name == "in_proj_z" || name == "in_proj_a" || + name == "in_proj_b" || name == "out_proj" || + name == "gate_proj" || name == "up_proj" || + name == "down_proj" || name == "shared_gate_proj" || + name == "shared_up_proj" || name == "shared_down_proj" || + name == "experts_gate_up_proj" || name == "experts_down_proj", + "unsupported native Qwen LoRA target module: ", name, + "; supported routed expert targets are experts_gate_up_proj/experts_down_proj"); + } + for (const auto& name : target_modules) { + bool resolved = false; + for (const auto& layer_cfg : ctx->layer_configs) { + auto table = lora_projection_table(layer_cfg); + for (int64_t k = 0; k < table.count; ++k) { + if (name == table.entries[k].name) { + resolved = true; + break; + } + } + if (resolved) break; + } + TORCH_CHECK(resolved, + "LoRA target module does not exist in this model: ", name); } - // Create LoRA parameters for target layers only + // Create fixed positional slots for every layer so layer offsets stay + // stable. Inactive slots are zero tensors without grad and are skipped + // by cache construction/Adam; this preserves the existing FFI export + // indexing while honoring target_modules exactly. int64_t offset = 0; - auto kind = ctx->compute_type; for (int64_t i = 0; i < num_layers; i++) { - int64_t lora_count = (ctx->layer_configs[i].layer_type == 0) ? 4 : 3; + auto projection_table = lora_projection_table(ctx->layer_configs[i]); + int64_t lora_count = projection_table.count; ctx->lora_layer_offset.push_back(offset); - if (target_set.find(i) == target_set.end()) { - // Not a target layer — no LoRA params, offset stays same - continue; - } - // Get base weight shapes from the weight pointers int64_t w_offset = 0; for (int64_t j = 0; j < i; j++) w_offset += weight_count_for_layer(ctx->layer_configs[j]); - if (ctx->layer_configs[i].layer_type == 0) { - // Full attention: q_proj, k_proj, v_proj, o_proj - int64_t proj_indices[] = {2, 4, 6, 7}; // q, k, v, o - for (int k = 0; k < 4; k++) { - auto* base = ctx->weight_ptrs[w_offset + proj_indices[k]]; - int64_t out_f = base->size(0), in_f = base->size(1); - auto a = at::randn({lora_rank, in_f}, at::TensorOptions().dtype(ctx->compute_type).device(base->device())) * 0.01; - auto b = at::zeros({out_f, lora_rank}, at::TensorOptions().dtype(ctx->compute_type).device(base->device())); - a.set_requires_grad(true); - b.set_requires_grad(true); - ctx->lora_a.push_back(std::move(a)); - ctx->lora_b.push_back(std::move(b)); - ctx->lora_names.push_back("lora_a_" + std::to_string(i) + "_" + std::to_string(k)); - ctx->lora_names.push_back("lora_b_" + std::to_string(i) + "_" + std::to_string(k)); - } - } else { - // Linear attention: in_proj_qkv, in_proj_z, out_proj - int64_t proj_indices[] = {2, 3, 10}; // qkv, z, out - for (int k = 0; k < 3; k++) { - auto* base = ctx->weight_ptrs[w_offset + proj_indices[k]]; + for (int64_t k = 0; k < projection_table.count; ++k) { + const auto& projection = projection_table.entries[k]; + auto* base = ctx->weight_ptrs[w_offset + projection.weight_index]; + TORCH_CHECK(base, "null LoRA base projection: layer=", i, + " module=", projection.name); + TORCH_CHECK( + (!projection.grouped_expert && base->dim() == 2) + || (projection.grouped_expert && base->dim() == 3), + "LoRA projection rank mismatch: layer=", i, + " module=", projection.name, " base_dim=", base->dim()); + bool active = (all_target_layers || target_set.find(i) != target_set.end()) && + (target_modules.empty() || target_modules.count(projection.name) > 0); + auto opts = at::TensorOptions().dtype(ctx->compute_type).device(base->device()); + at::Tensor a, b; + if (!active) { + // Stable slot index without allocating potentially hundreds + // of MB for inactive expert-local tensors. + a = at::zeros({}, opts); + b = at::zeros({}, opts); + } else if (projection.grouped_expert) { + int64_t experts = base->size(0); + int64_t out_f = base->size(1), in_f = base->size(2); + a = at::randn({experts, lora_rank, in_f}, opts) * 0.01; + b = at::zeros({experts, out_f, lora_rank}, opts); + } else { int64_t out_f = base->size(0), in_f = base->size(1); - auto a = at::randn({lora_rank, in_f}, at::TensorOptions().dtype(ctx->compute_type).device(base->device())) * 0.01; - auto b = at::zeros({out_f, lora_rank}, at::TensorOptions().dtype(ctx->compute_type).device(base->device())); - a.set_requires_grad(true); - b.set_requires_grad(true); - ctx->lora_a.push_back(std::move(a)); - ctx->lora_b.push_back(std::move(b)); - ctx->lora_names.push_back("lora_a_" + std::to_string(i) + "_" + std::to_string(k)); - ctx->lora_names.push_back("lora_b_" + std::to_string(i) + "_" + std::to_string(k)); + a = at::randn({lora_rank, in_f}, opts) * 0.01; + b = at::zeros({out_f, lora_rank}, opts); } + a.set_requires_grad(active); + b.set_requires_grad(active); + ctx->lora_a.push_back(std::move(a)); + ctx->lora_b.push_back(std::move(b)); + ctx->lora_active.push_back(active ? 1 : 0); + auto prefix = "layers." + std::to_string(i) + "." + projection.name; + ctx->lora_names.push_back(prefix + ".lora_A.weight"); + ctx->lora_names.push_back(prefix + ".lora_B.weight"); } offset += lora_count; } @@ -2777,17 +3624,20 @@ __attribute__((visibility("default"))) double qwen36_train_step( auto* ctx = reinterpret_cast(ctx_ptr); // Set CUDA device for EP if (ctx->nccl_comm) { - c10::cuda::set_device(ctx->ep_rank); - cudaSetDevice(ctx->ep_rank); + c10::cuda::set_device(ctx->cuda_device); + cudaSetDevice(ctx->cuda_device); } auto& input_ids = *reinterpret_cast(input_ids_ptr); auto& target_mask = *reinterpret_cast(target_mask_ptr); if (attention_mask_ptr) { ctx->attention_mask = *reinterpret_cast(attention_mask_ptr); + elide_trivial_attention_mask(ctx); } // Forward: checkpoint (default) or fused layer (QWEN36_FUSED_LAYER=1) - bool use_fused = getenv("QWEN36_FUSED_LAYER"); + bool use_fused = env_enabled("QWEN36_FUSED_LAYER"); + TORCH_CHECK(!use_fused, + "QWEN36_FUSED_LAYER is disabled until its custom backward preserves the layer graph"); auto hidden = use_fused ? forward_full_fused(ctx, input_ids) : ctx->use_checkpoint @@ -2802,59 +3652,35 @@ __attribute__((visibility("default"))) double qwen36_train_step( // Main loss — compute_loss does chunked CE with immediate backward per chunk. // This avoids accumulating 250 chunks of [512, vocab] logits in autograd graph. - auto loss = compute_loss(ctx, hidden, input_ids, target_mask, ctx->vocab_size); - - // compute_loss did chunked CE backward on detached hidden, - // accumulated gradient into hidden.grad(). - // Now trigger main model backward using hidden's gradient. - double loss_val = loss.item(); - - // Debug: GPU memory after CE backward - { - size_t free, total; - cudaMemGetInfo(&free, &total); - } - - // Release CE's retained graph (hidden_normed etc.) before backward. - // CE gradient is already accumulated into hidden.grad(). - // hidden.detach() breaks the autograd graph from CE. - hidden = hidden.detach(); - if (hidden.size(1) > 4096) c10::cuda::CUDACachingAllocator::emptyCache(); - - // Debug: after releasing CE graph - { - size_t free, total; - cudaMemGetInfo(&free, &total); + auto main_loss = compute_loss(ctx, hidden, input_ids, target_mask, ctx->vocab_size); + double loss_val = main_loss.value.item(); + auto total_hidden_grad = main_loss.hidden_grad; + + // MTP must be differentiated before the main model backward. Its + // frozen head still contributes a hidden-state gradient to trainable + // main-layer LoRA parameters. + if (ctx->has_mtp && !env_enabled("QWEN36_DISABLE_MTP")) { + auto mtp_input = hidden.detach().set_requires_grad(true); + auto mtp_hidden = mtp_forward(ctx, mtp_input, input_ids); + auto mtp_loss = mtp_compute_loss(ctx, mtp_hidden, input_ids, target_mask); + mtp_loss.backward(); + TORCH_CHECK(mtp_input.grad().defined(), "MTP did not produce a hidden gradient"); + total_hidden_grad.add_(mtp_input.grad()); + loss_val += mtp_loss.item(); } - // Trigger main model backward. - // FusedLayer: PyTorch autograd handles backward per layer. - // Checkpoint: manual_group_backward recomputes each group with PyTorch. - if (use_fused && hidden.grad().defined()) { - hidden.backward(hidden.grad()); - } else if (hidden.grad().defined() && !ctx->group_inputs.empty() && !getenv("QWEN36_SUBCKPT")) { - manual_group_backward(ctx, hidden.grad()); - } else if (hidden.grad().defined()) { - hidden.backward(hidden.grad()); + // Trigger exactly one main-model backward with the combined hidden + // gradient. Manual groups are the non-autograd checkpoint fallback; + // normal and sub-checkpoint paths use the real graph. + if (!ctx->group_inputs.empty() && !env_enabled("QWEN36_SUBCKPT")) { + manual_group_backward(ctx, total_hidden_grad); + } else { + hidden.backward(total_hidden_grad); } - // MTP loss (if enabled) — run AFTER main backward to reuse freed GPU memory - if (ctx->has_mtp && !getenv("QWEN36_DISABLE_MTP")) { - // Detach hidden so MTP forward doesn't rebuild main model graph - auto hidden_detached = hidden.detach().set_requires_grad(true); - auto mtp_hidden = mtp_forward(ctx, hidden_detached, input_ids); - auto mtp_loss = mtp_compute_loss(ctx, mtp_hidden, input_ids, target_mask); - // Backward MTP loss — frees MTP intermediate tensors immediately - mtp_loss.backward(); - // Add MTP gradient to hidden's gradient (already populated by main backward) - if (hidden_detached.grad().defined()) { - if (hidden.grad().defined()) { - hidden.grad().add_(hidden_detached.grad()); - } else { - hidden.mutable_grad() = hidden_detached.grad().clone(); - } - } - } + // EP produces a summed routed output, so synchronize replicated LoRA + // gradients before every rank performs its local Adam update. + synchronize_lora_gradients(ctx); // ── Adam optimizer step — CUDA multi-tensor fused kernel ── at::AutoGradMode guard(false); @@ -2943,11 +3769,11 @@ __attribute__((visibility("default"))) double qwen36_train_step( auto m_cpu = at::from_blob(h_m.data(), {n_params}, opts_cpu_long); auto v_cpu = at::from_blob(h_v.data(), {n_params}, opts_cpu_long); auto sizes_cpu = at::from_blob(h_sizes.data(), {n_params}, opts_cpu_int); - ctx->adam_dev_bufs.params_buf.copy_(params_cpu); - ctx->adam_dev_bufs.grads_buf.copy_(grads_cpu); - ctx->adam_dev_bufs.m_buf.copy_(m_cpu); - ctx->adam_dev_bufs.v_buf.copy_(v_cpu); - ctx->adam_dev_bufs.sizes_buf.copy_(sizes_cpu); + ctx->adam_dev_bufs.params_buf.narrow(0, 0, n_params).copy_(params_cpu); + ctx->adam_dev_bufs.grads_buf.narrow(0, 0, n_params).copy_(grads_cpu); + ctx->adam_dev_bufs.m_buf.narrow(0, 0, n_params).copy_(m_cpu); + ctx->adam_dev_bufs.v_buf.narrow(0, 0, n_params).copy_(v_cpu); + ctx->adam_dev_bufs.sizes_buf.narrow(0, 0, n_params).copy_(sizes_cpu); // Single kernel launch for ALL params auto stream = c10::cuda::getCurrentCUDAStream().stream(); @@ -2975,15 +3801,44 @@ __attribute__((visibility("default"))) double qwen36_train_step( // Get LoRA A tensor pointer by index __attribute__((visibility("default"))) void* qwen36_get_lora_a(void* ctx_ptr, int64_t index) { auto* ctx = reinterpret_cast(ctx_ptr); + if (index < 0 || index >= (int64_t)ctx->lora_a.size()) return nullptr; return &ctx->lora_a[index]; } // Get LoRA B tensor pointer by index __attribute__((visibility("default"))) void* qwen36_get_lora_b(void* ctx_ptr, int64_t index) { auto* ctx = reinterpret_cast(ctx_ptr); + if (index < 0 || index >= (int64_t)ctx->lora_b.size()) return nullptr; return &ctx->lora_b[index]; } +// Copy one exported LoRA tensor back into the native leaf parameter. This is +// used by checkpoint resume and adapter import; derived delta caches are +// invalidated so the next forward rebuilds the graph from the new leaf. +__attribute__((visibility("default"))) int32_t qwen36_set_lora_tensor( + void* ctx_ptr, int64_t index, int32_t is_b, void* tensor_ptr +) { + try { + auto* ctx = reinterpret_cast(ctx_ptr); + TORCH_CHECK(tensor_ptr, "null LoRA tensor"); + auto& slots = is_b ? ctx->lora_b : ctx->lora_a; + TORCH_CHECK(index >= 0 && index < (int64_t)slots.size(), "invalid LoRA slot"); + auto& target = slots[index]; + auto& source = *reinterpret_cast(tensor_ptr); + TORCH_CHECK(source.sizes() == target.sizes(), + "LoRA tensor shape mismatch at slot ", index, + ": expected ", target.sizes(), " got ", source.sizes()); + at::NoGradGuard guard; + target.copy_(source.to(target.device()).to(target.scalar_type())); + ctx->lora_cache_valid = false; + ctx->lora_batch_valid = false; + return 0; + } catch (const std::exception& e) { + fprintf(stderr, "[q36] set_lora_tensor FAILED: %s\n", e.what()); + return -1; + } +} + // Free training context __attribute__((visibility("default"))) void qwen36_free_training_context(void* ctx_ptr) { if (ctx_ptr) { @@ -3026,7 +3881,7 @@ static int64_t compute_n_max( // CE peak: one chunk of [16384, vocab] logits at a time. // BF16 logits (8GB) + FP32 CE loss (16GB) + grad (~8GB) ≈ 16GB. // But backward releases immediately, so effective peak is lower. - int64_t ce_peak = 16384 * 248320 * 4; // ~16GB (FP32 logits only, others freed by autograd) + int64_t ce_peak = 16384LL * 248320LL * 4LL; // ~16GB (FP32 logits only, others freed by autograd) // Attention intermediate: Q/K/V + attn_weights ≈ 4 × N × heads × seq × head_dim × 2 bytes // Flash attention keeps this O(seq) not O(seq²), but still significant at seq=16K @@ -3051,7 +3906,8 @@ static int64_t compute_n_max( } /// Train all adapters in chunks. Each chunk: independent forward → loss → backward → Adam. -/// Input is expanded to [N, seq] for each chunk. +/// Inputs may be [1, seq] (shared prompt, repeated per chunk) or +/// [n_total, seq] (one independent sample per adapter). __attribute__((visibility("default"))) double qwen36_train_multi_lora( void* ctx_ptr, void* input_ids_ptr, @@ -3063,18 +3919,42 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( try { auto* ctx = reinterpret_cast(ctx_ptr); if (ctx->nccl_comm) { - c10::cuda::set_device(ctx->ep_rank); - cudaSetDevice(ctx->ep_rank); + c10::cuda::set_device(ctx->cuda_device); + cudaSetDevice(ctx->cuda_device); } auto& input_ids = *reinterpret_cast(input_ids_ptr); auto& target_mask = *reinterpret_cast(target_mask_ptr); - if (attention_mask_ptr) { - ctx->attention_mask = *reinterpret_cast(attention_mask_ptr); - } int64_t total_adapters = (int64_t)ctx->adapters.size(); if (total_adapters == 0) return -1.0; + TORCH_CHECK(n_total > 0 && total_adapters == n_total, + "n_total must equal the number of registered adapters (n_total=", + n_total, ", registered=", total_adapters, ")"); + TORCH_CHECK(input_ids.dim() == 2 && target_mask.dim() == 2, + "multi-LoRA inputs must have shape [batch, seq]"); + const int64_t input_batch = input_ids.size(0); + TORCH_CHECK(input_batch == 1 || input_batch == n_total, + "multi-LoRA input batch must be 1 or n_total (batch=", input_batch, + ", n_total=", n_total, ")"); + TORCH_CHECK(target_mask.size(0) == input_batch && + target_mask.size(1) == input_ids.size(1), + "target_mask must match input_ids shape"); + + // Keep the caller's mask intact. Each chunk receives either the + // corresponding rows or a repeated batch-1 mask; this also prevents + // elide_trivial_attention_mask from leaking the last chunk into the + // next train step. + at::Tensor provided_attention_mask; + if (attention_mask_ptr) { + provided_attention_mask = *reinterpret_cast(attention_mask_ptr); + TORCH_CHECK(provided_attention_mask.dim() == 2, + "multi-LoRA attention_mask must have shape [batch, seq]"); + TORCH_CHECK(provided_attention_mask.size(0) == input_batch && + provided_attention_mask.size(1) == input_ids.size(1), + "attention_mask must match input_ids shape"); + } + const at::Tensor saved_attention_mask = ctx->attention_mask; // Compute N_max from available GPU memory. // CRITICAL: all workers must agree on n_max to keep NCCL all-reduce in sync. @@ -3083,7 +3963,7 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( cudaMemGetInfo(&free_mem, &total_mem); int64_t n_max; if (ctx->nccl_comm && ctx->ep_world_size > 1) { - const char* sync_path = "/tmp/rustrain-nccl/nmax_sync.txt"; + const std::string sync_path = nccl_sync_dir() + "/nmax_sync.txt"; if (ctx->ep_rank == 0) { n_max = compute_n_max( (int64_t)free_mem, lora_rank, @@ -3092,14 +3972,12 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( ); n_max = std::min(n_max, total_adapters); if (n_max < 1) n_max = 1; - mkdir("/tmp/rustrain-nccl", 0777); - FILE* f = fopen(sync_path, "w"); + FILE* f = fopen(sync_path.c_str(), "w"); fprintf(f, "%ld\n", (long)n_max); fclose(f); } else { - mkdir("/tmp/rustrain-nccl", 0777); for (int i = 0; i < 600; i++) { - FILE* f = fopen(sync_path, "r"); + FILE* f = fopen(sync_path.c_str(), "r"); if (f) { fscanf(f, "%ld", (long*)&n_max); fclose(f); break; } usleep(10000); } @@ -3132,16 +4010,29 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( // Temporarily set lora_batch_valid so prepare_lora_batch runs // We need to select only adapters[start:end] // HACK: move non-chunk adapters to a temp vector, run, then restore - std::vector temp; - temp.swap(ctx->adapters); - ctx->adapters.assign(temp.begin() + start, temp.begin() + end); + std::vector all_adapters; + all_adapters.swap(ctx->adapters); + ctx->adapters.assign( + all_adapters.begin() + start, all_adapters.begin() + end); // Mark batched mode active ctx->lora_batch_valid = true; // triggers prepare_lora_batch in forward - // Expand input_ids from [1, seq] to [N, seq] - auto ids_expanded = input_ids.repeat({n, 1}); - auto mask_expanded = target_mask.repeat({n, 1}); + // Expand only a shared batch-1 sample. With a tenant-specific + // [n_total, seq] batch, preserve each adapter's own row and slice + // the same chunk range as the adapter registry. + auto ids_expanded = input_batch == 1 + ? input_ids.repeat({n, 1}) + : input_ids.narrow(0, start, n).contiguous(); + auto mask_expanded = input_batch == 1 + ? target_mask.repeat({n, 1}) + : target_mask.narrow(0, start, n).contiguous(); + if (provided_attention_mask.defined()) { + ctx->attention_mask = input_batch == 1 + ? provided_attention_mask.repeat({n, 1}) + : provided_attention_mask.narrow(0, start, n).contiguous(); + elide_trivial_attention_mask(ctx); + } // Run train_step (reuses existing forward + loss + backward + Adam) // But we need to pass the expanded tensors @@ -3149,7 +4040,9 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( auto& mask_ref = mask_expanded; // Forward — force checkpoint for multi-LoRA (needed for group_inputs) - bool use_fused = getenv("QWEN36_FUSED_LAYER"); + bool use_fused = env_enabled("QWEN36_FUSED_LAYER"); + TORCH_CHECK(!use_fused, + "QWEN36_FUSED_LAYER is disabled until its custom backward preserves the layer graph"); ctx->use_checkpoint = true; // force checkpoint for manual_group_backward auto t_fwd_start = std::chrono::steady_clock::now(); @@ -3161,34 +4054,37 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( // Batched CE: compute loss with autograd enabled. double loss_val; + at::Tensor hidden_grad; auto t_loss_start = std::chrono::steady_clock::now(); { at::AutoGradMode grad_enable(true); // Re-attach hidden to autograd graph hidden = hidden.detach().set_requires_grad(true); - if (getenv("QWEN36_FUSED_CE")) { - auto loss = compute_loss_fused(ctx, hidden, input_ref, mask_ref, ctx->vocab_size); - loss_val = loss.item(); - } else { - auto loss = compute_loss(ctx, hidden, input_ref, mask_ref, ctx->vocab_size); - loss_val = loss.item(); - } + TORCH_CHECK(!env_enabled("QWEN36_FUSED_CE"), + "QWEN36_FUSED_CE is disabled until its tile gather and gradient normalization are validated"); + auto loss = compute_loss(ctx, hidden, input_ref, mask_ref, ctx->vocab_size); + loss_val = loss.value.item(); + hidden_grad = loss.hidden_grad; } auto t_loss_end = std::chrono::steady_clock::now(); double loss_ms = std::chrono::duration(t_loss_end - t_loss_start).count(); // Backward auto t_bwd_start = std::chrono::steady_clock::now(); - auto hidden_grad = hidden.grad(); - hidden = hidden.detach(); - if (hidden.size(1) > 4096) c10::cuda::CUDACachingAllocator::emptyCache(); + if (ctx->has_mtp && !env_enabled("QWEN36_DISABLE_MTP")) { + auto mtp_input = hidden.detach().set_requires_grad(true); + auto mtp_hidden = mtp_forward(ctx, mtp_input, input_ref); + auto mtp_loss = mtp_compute_loss(ctx, mtp_hidden, input_ref, mask_ref); + mtp_loss.backward(); + TORCH_CHECK(mtp_input.grad().defined(), "MTP did not produce a hidden gradient"); + hidden_grad.add_(mtp_input.grad()); + loss_val += mtp_loss.item(); + } - if (use_fused && hidden_grad.defined()) { - hidden.backward(hidden_grad); - } else if (hidden_grad.defined() && !ctx->group_inputs.empty() && !getenv("QWEN36_SUBCKPT")) { + if (!ctx->group_inputs.empty() && !env_enabled("QWEN36_SUBCKPT")) { manual_group_backward(ctx, hidden_grad); - } else if (hidden_grad.defined()) { - hidden.backward(hidden.grad()); + } else { + hidden.backward(hidden_grad); } // No cudaDeviceSynchronize — let GPU pipeline run asynchronously. // The next chunk's CPU prep (LoRA batch, input expand) will overlap @@ -3199,20 +4095,9 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( fprintf(stderr, "[train_multi] chunk %ld/%ld: n=%ld loss=%f fwd=%.0fms loss=%.0fms bwd=%.0fms\n", (long)(chunk+1), (long)num_chunks, (long)n, loss_val, fwd_ms, loss_ms, bwd_ms); - // MTP (if enabled) - if (ctx->has_mtp && !getenv("QWEN36_DISABLE_MTP")) { - auto hidden_detached = hidden.detach().set_requires_grad(true); - auto mtp_hidden = mtp_forward(ctx, hidden_detached, input_ref); - auto mtp_loss = mtp_compute_loss(ctx, mtp_hidden, input_ref, mask_ref); - mtp_loss.backward(); - if (hidden_detached.grad().defined()) { - if (hidden.grad().defined()) { - hidden.grad().add_(hidden_detached.grad()); - } else { - hidden.mutable_grad() = hidden_detached.grad().clone(); - } - } - } + // EP gradient synchronization must happen before this chunk's + // replicated adapter parameters are updated. + synchronize_lora_gradients(ctx); // Adam step at::AutoGradMode guard(false); @@ -3266,11 +4151,11 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( auto v_cpu = at::from_blob(h_v.data(), {n_params}, opts_cpu_long); auto sizes_cpu = at::from_blob(h_sizes.data(), {n_params}, opts_cpu_int); ctx->adam_dev_bufs.ensure(n_params, ctx->adapters[0].params.begin()->second[0].first); - ctx->adam_dev_bufs.params_buf.copy_(params_cpu); - ctx->adam_dev_bufs.grads_buf.copy_(grads_cpu); - ctx->adam_dev_bufs.m_buf.copy_(m_cpu); - ctx->adam_dev_bufs.v_buf.copy_(v_cpu); - ctx->adam_dev_bufs.sizes_buf.copy_(sizes_cpu); + ctx->adam_dev_bufs.params_buf.narrow(0, 0, n_params).copy_(params_cpu); + ctx->adam_dev_bufs.grads_buf.narrow(0, 0, n_params).copy_(grads_cpu); + ctx->adam_dev_bufs.m_buf.narrow(0, 0, n_params).copy_(m_cpu); + ctx->adam_dev_bufs.v_buf.narrow(0, 0, n_params).copy_(v_cpu); + ctx->adam_dev_bufs.sizes_buf.narrow(0, 0, n_params).copy_(sizes_cpu); auto stream = c10::cuda::getCurrentCUDAStream().stream(); launch_fused_adam_multi( @@ -3287,22 +4172,18 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( ); } - // Restore adapters - ctx->adapters.swap(temp); - // Re-merge: put updated chunk adapters back - for (int64_t i = 0; i < n; i++) { - temp[start + i] = std::move(ctx->adapters[i]); - } - ctx->adapters.swap(temp); + // Restore the complete registry. Adapter parameter and Adam-state + // tensors are intrusive handles, so the chunk copies above share + // the updated storage with all_adapters; no value merge is needed. + ctx->adapters.swap(all_adapters); total_loss += loss_val; - // emptyCache at chunk boundary — does NOT sync GPU. - c10::cuda::CUDACachingAllocator::emptyCache(); fprintf(stderr, "[train_multi] chunk %ld/%ld: n=%ld loss=%.6f\n", (long)(chunk + 1), (long)num_chunks, (long)n, loss_val); } + ctx->attention_mask = saved_attention_mask; return total_loss / num_chunks; } catch (const std::exception& e) { fprintf(stderr, "[train_multi] FAILED: %s\n", e.what()); @@ -3322,6 +4203,7 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( static ncclComm_t g_nccl_comm = nullptr; static cudaStream_t g_nccl_stream = nullptr; static bool g_nccl_initialized = false; +static int g_cuda_device = 0; // Set CUDA device — called from Rust worker before any GPU operation. // Ensures PyTorch initializes CUDA context on the correct device. @@ -3331,6 +4213,7 @@ __attribute__((visibility("default"))) void qwen36_set_cuda_device(int32_t devic // Must be called before any GPU operation in exec'd worker processes. c10::cuda::set_device(device); cudaSetDevice(device); + g_cuda_device = device; // Force PyTorch to create CUDA context on this device auto opts = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, device); auto dummy = at::empty({1}, opts); @@ -3353,6 +4236,8 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( const char* world_str2 = getenv("WORLD_SIZE"); if (rank_str2) ctx->ep_rank = atoi(rank_str2); if (world_str2) ctx->ep_world_size = atoi(world_str2); + const char* local_rank_str2 = getenv("LOCAL_RANK"); + ctx->cuda_device = local_rank_str2 ? atoi(local_rank_str2) : g_cuda_device; for (auto& lc : ctx->layer_configs) { lc.nccl_comm = (void*)g_nccl_comm; lc.nccl_stream = (void*)g_nccl_stream; } for (auto& lc : ctx->mtp_layer_configs) { lc.nccl_comm = (void*)g_nccl_comm; lc.nccl_stream = (void*)g_nccl_stream; } return 0; @@ -3371,7 +4256,10 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( // Creating a dummy tensor forces PyTorch to initialize its CUDA context. const char* local_rank_str = getenv("LOCAL_RANK"); int local_rank = local_rank_str ? atoi(local_rank_str) : rank; + c10::cuda::set_device(local_rank); cudaSetDevice(local_rank); + g_cuda_device = local_rank; + ctx->cuda_device = local_rank; { // Force PyTorch CUDA context initialization on this device auto opts = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, local_rank); @@ -3383,30 +4271,35 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( // Rank 0 generates ID, writes to file, then writes "ready" sentinel. // Other ranks wait for "ready" file, then read ID. // This ensures rank 0's write is visible before others read. - const char* id_path = "/tmp/rustrain-nccl/nccl-id.bin"; - const char* ready_path = "/tmp/rustrain-nccl/nccl-ready.txt"; + const std::string rendezvous_dir = nccl_sync_dir(); + const std::string id_path = rendezvous_dir + "/nccl-id.bin"; + const std::string ready_path = rendezvous_dir + "/nccl-ready.txt"; ncclUniqueId unique_id; if (rank == 0) { - mkdir("/tmp/rustrain-nccl", 0777); // Clean up old files first - remove("/tmp/rustrain-nccl/nccl-ready.txt"); + remove(ready_path.c_str()); + for (int peer = 0; peer < world_size; ++peer) { + const std::string stale_barrier = + rendezvous_dir + "/barrier_" + std::to_string(peer); + remove(stale_barrier.c_str()); + } ncclGetUniqueId(&unique_id); - FILE* f = fopen(id_path, "wb"); + FILE* f = fopen(id_path.c_str(), "wb"); fwrite(&unique_id, sizeof(unique_id), 1, f); fclose(f); // Write ready sentinel AFTER id file - FILE* rf = fopen(ready_path, "w"); + FILE* rf = fopen(ready_path.c_str(), "w"); fprintf(rf, "ready\n"); fclose(rf); } else { // Wait for ready sentinel for (int i = 0; i < 600; i++) { - FILE* rf = fopen(ready_path, "r"); + FILE* rf = fopen(ready_path.c_str(), "r"); if (rf) { fclose(rf); break; } usleep(10000); // 10ms } // Now read ID file - FILE* f = fopen(id_path, "rb"); + FILE* f = fopen(id_path.c_str(), "rb"); if (!f || fread(&unique_id, sizeof(unique_id), 1, f) != 1) { fprintf(stderr, "[ep_nccl] rank %d: failed to read ID file\n", rank); if (f) fclose(f); @@ -3419,7 +4312,7 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( // Without this, rank 0 (fast load_model) reaches ncclCommInitRank before // rank 3 (slow load_model) → NCCL timeout. { - const char* barrier_dir = "/tmp/rustrain-nccl"; + const char* barrier_dir = rendezvous_dir.c_str(); char bpath[256]; snprintf(bpath, sizeof(bpath), "%s/barrier_%d", barrier_dir, rank); FILE* bf = fopen(bpath, "w"); fprintf(bf, "1\n"); fclose(bf); @@ -3483,6 +4376,9 @@ __attribute__((visibility("default"))) void qwen36_set_nccl_comm( ctx->nccl_stream = reinterpret_cast(stream_ptr); ctx->ep_rank = ep_rank; ctx->ep_world_size = ep_world_size; + int current_device = g_cuda_device; + cudaGetDevice(¤t_device); + ctx->cuda_device = current_device; // Propagate NCCL handles to all layer configs so moe_forward can access them for (auto& lc : ctx->layer_configs) { lc.nccl_comm = comm_ptr; @@ -3535,6 +4431,8 @@ int64_t qwen36_add_lora( ) { try { auto* ctx = reinterpret_cast(ctx_ptr); + TORCH_CHECK(rank > 0, "LoRA rank must be positive"); + TORCH_CHECK(alpha > 0.0, "LoRA alpha must be positive"); TrainingContext::LoRAAdapter adapter; adapter.id = ++ctx->next_adapter_id; adapter.rank = rank; @@ -3550,32 +4448,93 @@ int64_t qwen36_add_lora( while (std::getline(ss, item, ',')) adapter.target_modules.insert(item); } + for (auto layer : adapter.target_layers) { + TORCH_CHECK(layer >= 0 && layer < ctx->num_layers, + "dynamic LoRA target layer out of range: ", layer, + " for model with ", ctx->num_layers, " layers"); + } + // The activation-level batch path stacks A/B across adapters. Keep + // the batch rectangular and semantically aligned instead of waiting + // for an opaque ATen stack/shape failure during the first step. + if (!ctx->adapters.empty()) { + const auto& reference = ctx->adapters.front(); + TORCH_CHECK(rank == reference.rank, + "dynamic LoRA adapters in one batch must use the same rank"); + TORCH_CHECK(adapter.target_layers == reference.target_layers, + "dynamic LoRA adapters in one batch must use identical target_layers"); + TORCH_CHECK(adapter.target_modules == reference.target_modules, + "dynamic LoRA adapters in one batch must use identical target_modules"); + } + for (const auto& name : adapter.target_modules) { + TORCH_CHECK( + name == "q_proj" || name == "k_proj" || name == "v_proj" || + name == "o_proj" || name == "in_proj_qkv" || + name == "in_proj_z" || name == "in_proj_a" || + name == "in_proj_b" || name == "out_proj" || + name == "gate_proj" || name == "up_proj" || + name == "down_proj" || name == "shared_gate_proj" || + name == "shared_up_proj" || name == "shared_down_proj" || + name == "experts_gate_up_proj" || name == "experts_down_proj", + "unsupported dynamic Qwen LoRA target module: ", name); + bool resolved = false; + for (const auto& layer_cfg : ctx->layer_configs) { + auto table = lora_projection_table(layer_cfg); + for (int64_t pair = 0; pair < table.count; ++pair) { + if (name == table.entries[pair].name) { + resolved = true; + break; + } + } + if (resolved) break; + } + TORCH_CHECK(resolved, + "dynamic LoRA target module does not exist in this model: ", name); + } for (int64_t i = 0; i < ctx->num_layers; i++) { if (!adapter.target_layers.empty() && adapter.target_layers.find(i) == adapter.target_layers.end()) continue; int64_t w_offset = 0; for (int64_t j = 0; j < i; j++) w_offset += weight_count_for_layer(ctx->layer_configs[j]); - int64_t num_pairs; - const int64_t* proj_indices; - if (ctx->layer_configs[i].layer_type == 0) { - static const int64_t full_indices[] = {2, 4, 6, 7}; - proj_indices = full_indices; - num_pairs = 4; - } else { - static const int64_t linear_indices[] = {2, 3, 10}; - proj_indices = linear_indices; - num_pairs = 3; - } + auto projection_table = lora_projection_table(ctx->layer_configs[i]); + int64_t num_pairs = projection_table.count; std::vector> pairs; std::vector> adam_states; - for (int k = 0; k < num_pairs; k++) { - auto* base = ctx->weight_ptrs[w_offset + proj_indices[k]]; - int64_t out_f = base->size(0), in_f = base->size(1); - auto a = at::randn({rank, in_f}, at::TensorOptions().dtype(ctx->compute_type).device(base->device())) * 0.01; - auto b = at::zeros({out_f, rank}, at::TensorOptions().dtype(ctx->compute_type).device(base->device())); - a.set_requires_grad(true); - b.set_requires_grad(true); + for (int64_t k = 0; k < num_pairs; k++) { + const auto& projection = projection_table.entries[k]; + auto* base = ctx->weight_ptrs[w_offset + projection.weight_index]; + // Preserve the historical empty-target default (attention + // projections only). Explicit target lists may additionally + // select any 2D dense/shared MLP projection. + bool active = adapter.target_modules.empty() + ? !projection.grouped_expert && + projection.segment == LoraSegment::Attention + : adapter.target_modules.find(projection.name) != + adapter.target_modules.end(); + auto opts = at::TensorOptions().dtype(ctx->compute_type).device(base->device()); + at::Tensor a, b; + if (active) { + if (projection.grouped_expert) { + TORCH_CHECK(base->dim() == 3, + "dynamic routed-expert LoRA projection must be rank 3: ", + projection.name); + int64_t experts = base->size(0); + int64_t out_f = base->size(1), in_f = base->size(2); + a = at::randn({experts, rank, in_f}, opts) * 0.01; + b = at::zeros({experts, out_f, rank}, opts); + } else { + TORCH_CHECK(base->dim() == 2, + "dynamic LoRA projection must be a matrix: ", projection.name); + int64_t out_f = base->size(0), in_f = base->size(1); + a = at::randn({rank, in_f}, opts) * 0.01; + b = at::zeros({out_f, rank}, opts); + } + } else { + a = at::zeros({}, opts); + b = at::zeros({}, opts); + } + a.set_requires_grad(active); + b.set_requires_grad(active); // Adam state: FP32 for numerical stability auto opts_f32 = at::TensorOptions().dtype(at::kFloat).device(base->device()); adam_states.push_back({ @@ -3623,14 +4582,66 @@ int64_t qwen36_list_lora(void* ctx_ptr, int64_t* out_ids, int64_t max_count) { return count; } +__attribute__((visibility("default"))) +void* qwen36_get_adapter_lora_tensor( + void* ctx_ptr, int64_t adapter_id, int64_t layer_idx, + const char* module_name, int32_t is_b +) { + auto* ctx = reinterpret_cast(ctx_ptr); + if (!ctx || !module_name || layer_idx < 0 || layer_idx >= ctx->num_layers) + return nullptr; + const int64_t pair_idx = lora_pair_index( + ctx->layer_configs[layer_idx], module_name); + if (pair_idx < 0) return nullptr; + for (auto& adapter : ctx->adapters) { + if (adapter.id != adapter_id) continue; + auto it = adapter.params.find(layer_idx); + if (it == adapter.params.end() || + pair_idx >= static_cast(it->second.size())) + return nullptr; + auto& pair = it->second[pair_idx]; + auto& tensor = is_b ? pair.second : pair.first; + return tensor.requires_grad() ? &tensor : nullptr; + } + return nullptr; +} + +__attribute__((visibility("default"))) +int32_t qwen36_set_adapter_lora_tensor( + void* ctx_ptr, int64_t adapter_id, int64_t layer_idx, + const char* module_name, int32_t is_b, void* tensor_ptr +) { + try { + auto* ctx = reinterpret_cast(ctx_ptr); + TORCH_CHECK(ctx && module_name && tensor_ptr, + "invalid dynamic LoRA tensor setter arguments"); + auto* target = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_id, layer_idx, module_name, is_b)); + TORCH_CHECK(target, "dynamic LoRA target not found: adapter=", adapter_id, + " layer=", layer_idx, " module=", module_name); + auto& source = *reinterpret_cast(tensor_ptr); + TORCH_CHECK(source.sizes() == target->sizes(), + "dynamic LoRA tensor shape mismatch: expected ", target->sizes(), + " got ", source.sizes()); + at::NoGradGuard guard; + target->copy_(source.to(target->device()).to(target->scalar_type())); + ctx->lora_cache_valid = false; + ctx->lora_batch_valid = false; + return 0; + } catch (const std::exception& e) { + fprintf(stderr, "[q36] set_adapter_lora_tensor FAILED: %s\n", e.what()); + return -1; + } +} + __attribute__((visibility("default"))) int64_t qwen36_get_lora_count(void* ctx_ptr) { auto* ctx = reinterpret_cast(ctx_ptr); - int64_t total = (int64_t)ctx->lora_a.size(); - for (auto& adapter : ctx->adapters) - for (auto& [layer_idx, pairs] : adapter.params) - total += (int64_t)pairs.size() * 2; - return total; + // This legacy accessor is paired with get_lora_a/get_lora_b and therefore + // counts only the fixed single-adapter slots. Dynamic adapters have their + // own registry and must not make this count exceed those arrays. + return (int64_t)ctx->lora_a.size(); } __attribute__((visibility("default"))) @@ -3641,9 +4652,10 @@ double qwen36_eval_step(void* ctx_ptr, void* input_ids_ptr, void* target_mask_pt auto& target_mask = *reinterpret_cast(target_mask_ptr); if (attention_mask_ptr) ctx->attention_mask = *reinterpret_cast(attention_mask_ptr); + elide_trivial_attention_mask(ctx); at::AutoGradMode no_grad(false); auto hidden = ctx->use_checkpoint ? forward_full_checkpoint(ctx, input_ids) : forward_full(ctx, input_ids); - auto loss = compute_loss(ctx, hidden, input_ids, target_mask, ctx->vocab_size); + auto loss = compute_loss_fused(ctx, hidden, input_ids, target_mask, ctx->vocab_size); return loss.item(); } catch (const std::exception& e) { fprintf(stderr, "[q36] eval_step FAILED: %s\n", e.what()); diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs index 0db5edb9..529399fe 100644 --- a/crates/rustrain-qwen3-6/src/kernel.rs +++ b/crates/rustrain-qwen3-6/src/kernel.rs @@ -3,40 +3,62 @@ //! TrainingContext, train_step, and adapter export all happen in C++. //! Rust only handles: weight loading, data loading, training loop orchestration. +use crate::lora::Qwen36LoraTargetModule; +use anyhow::{Result, bail}; use std::ffi::c_void; use std::sync::OnceLock; use tch::{Kind, Tensor}; -use anyhow::{Result, bail}; // ── dlopen ── type FnCreateCtx = unsafe extern "C" fn( - *mut *mut c_void, i64, *mut c_void, *mut c_void, *mut c_void, - *mut c_void, i64, i32, f64, f64, f64, f64, f64, i64, f64, - i64, *const i64, i64, + *mut *mut c_void, + i64, + *mut c_void, + *mut c_void, + *mut c_void, + *mut c_void, + i64, + i32, + f64, + f64, + f64, + f64, + f64, + i64, + f64, + i64, + *const i64, + i64, + *const i8, ) -> *mut c_void; +type FnKernelAbiVersion = unsafe extern "C" fn() -> i64; type FnTrainStep = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void) -> f64; -type FnTrainMultiLora = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void, i32, i32) -> f64; +type FnTrainMultiLora = + unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void, i32, i32) -> f64; type FnEvalStep = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void) -> f64; type FnGetLoraCount = unsafe extern "C" fn(*mut c_void) -> i64; type FnGetLoraA = unsafe extern "C" fn(*mut c_void, i64) -> *mut c_void; type FnGetLoraB = unsafe extern "C" fn(*mut c_void, i64) -> *mut c_void; +type FnSetLoraTensor = unsafe extern "C" fn(*mut c_void, i64, i32, *mut c_void) -> i32; type FnGetStepCount = unsafe extern "C" fn(*mut c_void) -> i64; -type FnExportOptimizer = unsafe extern "C" fn(*mut c_void, *mut *mut c_void, *mut *mut c_void, i64) -> i64; -type FnImportOptimizer = unsafe extern "C" fn(*mut c_void, *mut *mut c_void, *mut *mut c_void, i64) -> i64; +type FnExportOptimizer = + unsafe extern "C" fn(*mut c_void, *mut *mut c_void, *mut *mut c_void, i64) -> i64; +type FnImportOptimizer = + unsafe extern "C" fn(*mut c_void, *mut *mut c_void, *mut *mut c_void, i64) -> i64; type FnFreeCtx = unsafe extern "C" fn(*mut c_void); type FnGemm = unsafe extern "C" fn(*mut c_void, *mut c_void, i32) -> *mut c_void; type FnFreeTensor = unsafe extern "C" fn(*mut c_void); type FnSetMtpWeights = unsafe extern "C" fn( - *mut c_void, // ctx_ptr - *mut c_void, // mtp_fc_ptr - *mut c_void, // mtp_pre_fc_norm_emb_ptr - *mut c_void, // mtp_pre_fc_norm_hidden_ptr - *mut c_void, // mtp_norm_ptr - *mut *mut c_void, // mtp_layer_weight_ptrs - i64, // num_mtp_layer_weights - *mut c_void, // mtp_layer_configs_ptr - i64, // num_mtp_layers + *mut c_void, // ctx_ptr + *mut c_void, // mtp_fc_ptr + *mut c_void, // mtp_pre_fc_norm_emb_ptr + *mut c_void, // mtp_pre_fc_norm_hidden_ptr + *mut c_void, // mtp_norm_ptr + *mut *mut c_void, // mtp_layer_weight_ptrs + i64, // num_mtp_layer_weights + *mut c_void, // mtp_layer_configs_ptr + i64, // num_mtp_layers ); type FnSetCheckpoint = unsafe extern "C" fn(*mut c_void, i32, i64); type FnSetNcclComm = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, i32, i32); @@ -45,6 +67,10 @@ type FnSetCudaDevice = unsafe extern "C" fn(i32); type FnAddLora = unsafe extern "C" fn(*mut c_void, i64, f64, *const i64, i64, *const i8) -> i64; type FnRemoveLora = unsafe extern "C" fn(*mut c_void, i64) -> i32; type FnListLora = unsafe extern "C" fn(*mut c_void, *mut i64, i64) -> i64; +type FnGetAdapterLoraTensor = + unsafe extern "C" fn(*mut c_void, i64, i64, *const i8, i32) -> *mut c_void; +type FnSetAdapterLoraTensor = + unsafe extern "C" fn(*mut c_void, i64, i64, *const i8, i32, *mut c_void) -> i32; #[repr(C)] pub struct CppLayerConfig { @@ -80,6 +106,7 @@ struct KernelHandles { get_lora_count: FnGetLoraCount, get_lora_a: FnGetLoraA, get_lora_b: FnGetLoraB, + set_lora_tensor: FnSetLoraTensor, get_step_count: FnGetStepCount, export_optimizer: FnExportOptimizer, import_optimizer: FnImportOptimizer, @@ -94,6 +121,8 @@ struct KernelHandles { add_lora: FnAddLora, remove_lora: FnRemoveLora, list_lora: FnListLora, + get_adapter_lora_tensor: FnGetAdapterLoraTensor, + set_adapter_lora_tensor: FnSetAdapterLoraTensor, } static KERNELS: OnceLock> = OnceLock::new(); @@ -115,18 +144,28 @@ unsafe fn load_kernels() -> Option { let handle = libc::dlopen(lib_name.as_ptr(), libc::RTLD_LAZY | libc::RTLD_NOLOAD); let handle = if handle.is_null() { let h = libc::dlopen(lib_name.as_ptr(), libc::RTLD_LAZY); - if h.is_null() { return None; } + if h.is_null() { + return None; + } h - } else { handle }; + } else { + handle + }; macro_rules! sym { ($name:expr) => {{ let s = CString::new($name).unwrap(); let p = libc::dlsym(handle, s.as_ptr()); - if p.is_null() { return None; } + if p.is_null() { + return None; + } std::mem::transmute::<*mut c_void, _>(p) }}; } + let abi_version: FnKernelAbiVersion = sym!("qwen36_kernel_abi_version"); + if abi_version() != 4 { + return None; + } Some(KernelHandles { create_ctx: sym!("qwen36_create_training_context"), train_step: sym!("qwen36_train_step"), @@ -135,6 +174,7 @@ unsafe fn load_kernels() -> Option { get_lora_count: sym!("qwen36_get_lora_count"), get_lora_a: sym!("qwen36_get_lora_a"), get_lora_b: sym!("qwen36_get_lora_b"), + set_lora_tensor: sym!("qwen36_set_lora_tensor"), get_step_count: sym!("qwen36_get_step_count"), export_optimizer: sym!("qwen36_export_optimizer_state"), import_optimizer: sym!("qwen36_import_optimizer_state"), @@ -149,6 +189,8 @@ unsafe fn load_kernels() -> Option { add_lora: sym!("qwen36_add_lora"), remove_lora: sym!("qwen36_remove_lora"), list_lora: sym!("qwen36_list_lora"), + get_adapter_lora_tensor: sym!("qwen36_get_adapter_lora_tensor"), + set_adapter_lora_tensor: sym!("qwen36_set_adapter_lora_tensor"), }) } @@ -178,7 +220,10 @@ pub fn build_weight_ptrs( for layer in 0..config.num_hidden_layers { let lp = format!("{p}layers.{layer}"); ptrs.push(get_ptr(weights, &format!("{lp}.input_layernorm.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.post_attention_layernorm.weight"))); + ptrs.push(get_ptr( + weights, + &format!("{lp}.post_attention_layernorm.weight"), + )); match config.layer_types[layer] { crate::config::LayerType::FullAttention => { for w in &["q_proj", "q_norm", "k_proj", "k_norm", "v_proj", "o_proj"] { @@ -186,23 +231,50 @@ pub fn build_weight_ptrs( } } crate::config::LayerType::LinearAttention => { - ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.in_proj_qkv.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.in_proj_z.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.in_proj_a.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.in_proj_b.weight"))); + ptrs.push(get_ptr( + weights, + &format!("{lp}.linear_attn.in_proj_qkv.weight"), + )); + ptrs.push(get_ptr( + weights, + &format!("{lp}.linear_attn.in_proj_z.weight"), + )); + ptrs.push(get_ptr( + weights, + &format!("{lp}.linear_attn.in_proj_a.weight"), + )); + ptrs.push(get_ptr( + weights, + &format!("{lp}.linear_attn.in_proj_b.weight"), + )); ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.A_log"))); ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.dt_bias"))); ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.conv1d.weight"))); ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.norm.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.linear_attn.out_proj.weight"))); + ptrs.push(get_ptr( + weights, + &format!("{lp}.linear_attn.out_proj.weight"), + )); } } if config.is_moe { ptrs.push(get_ptr(weights, &format!("{lp}.mlp.gate.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert_gate.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert.gate_proj.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert.up_proj.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert.down_proj.weight"))); + ptrs.push(get_ptr( + weights, + &format!("{lp}.mlp.shared_expert_gate.weight"), + )); + ptrs.push(get_ptr( + weights, + &format!("{lp}.mlp.shared_expert.gate_proj.weight"), + )); + ptrs.push(get_ptr( + weights, + &format!("{lp}.mlp.shared_expert.up_proj.weight"), + )); + ptrs.push(get_ptr( + weights, + &format!("{lp}.mlp.shared_expert.down_proj.weight"), + )); ptrs.push(get_ptr(weights, &format!("{lp}.mlp.experts.gate_up_proj"))); ptrs.push(get_ptr(weights, &format!("{lp}.mlp.experts.down_proj"))); } else { @@ -219,32 +291,37 @@ pub fn build_layer_configs( expert_start: usize, expert_count: usize, ) -> Vec { - (0..config.num_hidden_layers).map(|layer| { - let lt = &config.layer_types[layer]; - CppLayerConfig { - layer_type: match lt { crate::config::LayerType::FullAttention => 0, _ => 1 }, - num_heads: config.num_attention_heads, - num_kv_heads: config.num_key_value_heads, - head_dim: config.head_dim, - num_k_heads: config.linear_num_key_heads, - key_dim: config.linear_key_head_dim, - num_v_heads: config.linear_num_value_heads, - val_dim: config.linear_value_head_dim, - conv_kernel: config.linear_conv_kernel_dim, - partial_rotary_factor: config.partial_rotary_factor, - rope_theta: config.rope_theta, - rms_eps: config.rms_norm_eps, - num_experts: config.num_experts as i64, - top_k: config.num_experts_per_tok as i64, - moe_intermediate: config.moe_intermediate_size, - norm_topk_prob: if config.norm_topk_prob { 1 } else { 0 }, - expert_start: expert_start as i64, - expert_count: expert_count as i64, - intermediate_size: config.intermediate_size, - nccl_comm: std::ptr::null_mut(), - nccl_stream: std::ptr::null_mut(), - } - }).collect() + (0..config.num_hidden_layers) + .map(|layer| { + let lt = &config.layer_types[layer]; + CppLayerConfig { + layer_type: match lt { + crate::config::LayerType::FullAttention => 0, + _ => 1, + }, + num_heads: config.num_attention_heads, + num_kv_heads: config.num_key_value_heads, + head_dim: config.head_dim, + num_k_heads: config.linear_num_key_heads, + key_dim: config.linear_key_head_dim, + num_v_heads: config.linear_num_value_heads, + val_dim: config.linear_value_head_dim, + conv_kernel: config.linear_conv_kernel_dim, + partial_rotary_factor: config.partial_rotary_factor, + rope_theta: config.rope_theta, + rms_eps: config.rms_norm_eps, + num_experts: config.num_experts as i64, + top_k: config.num_experts_per_tok as i64, + moe_intermediate: config.moe_intermediate_size, + norm_topk_prob: if config.norm_topk_prob { 1 } else { 0 }, + expert_start: expert_start as i64, + expert_count: expert_count as i64, + intermediate_size: config.intermediate_size, + nccl_comm: std::ptr::null_mut(), + nccl_stream: std::ptr::null_mut(), + } + }) + .collect() } /// Build weight pointers for MTP layers (full attention layers). @@ -258,7 +335,10 @@ pub fn build_mtp_weight_ptrs( for layer in 0..config.mtp_num_hidden_layers { let lp = format!("mtp.layers.{layer}"); ptrs.push(get_ptr(weights, &format!("{lp}.input_layernorm.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.post_attention_layernorm.weight"))); + ptrs.push(get_ptr( + weights, + &format!("{lp}.post_attention_layernorm.weight"), + )); // Full attention: q, q_norm, k, k_norm, v, o for w in &["q_proj", "q_norm", "k_proj", "k_norm", "v_proj", "o_proj"] { ptrs.push(get_ptr(weights, &format!("{lp}.self_attn.{w}.weight"))); @@ -266,10 +346,22 @@ pub fn build_mtp_weight_ptrs( if config.is_moe { // MoE: gate, shared_expert_gate, shared_gate_proj, shared_up_proj, shared_down_proj, experts_gate_up, experts_down ptrs.push(get_ptr(weights, &format!("{lp}.mlp.gate.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert_gate.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert.gate_proj.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert.up_proj.weight"))); - ptrs.push(get_ptr(weights, &format!("{lp}.mlp.shared_expert.down_proj.weight"))); + ptrs.push(get_ptr( + weights, + &format!("{lp}.mlp.shared_expert_gate.weight"), + )); + ptrs.push(get_ptr( + weights, + &format!("{lp}.mlp.shared_expert.gate_proj.weight"), + )); + ptrs.push(get_ptr( + weights, + &format!("{lp}.mlp.shared_expert.up_proj.weight"), + )); + ptrs.push(get_ptr( + weights, + &format!("{lp}.mlp.shared_expert.down_proj.weight"), + )); ptrs.push(get_ptr(weights, &format!("{lp}.mlp.experts.gate_up_proj"))); ptrs.push(get_ptr(weights, &format!("{lp}.mlp.experts.down_proj"))); } else { @@ -288,31 +380,33 @@ pub fn build_mtp_layer_configs( expert_start: usize, expert_count: usize, ) -> Vec { - (0..config.mtp_num_hidden_layers).map(|_| { - CppLayerConfig { - layer_type: 0, // MTP layers are always full attention - num_heads: config.num_attention_heads, - num_kv_heads: config.num_key_value_heads, - head_dim: config.head_dim, - num_k_heads: config.linear_num_key_heads, - key_dim: config.linear_key_head_dim, - num_v_heads: config.linear_num_value_heads, - val_dim: config.linear_value_head_dim, - conv_kernel: config.linear_conv_kernel_dim, - partial_rotary_factor: config.partial_rotary_factor, - rope_theta: config.rope_theta, - rms_eps: config.rms_norm_eps, - num_experts: config.num_experts as i64, - top_k: config.num_experts_per_tok as i64, - moe_intermediate: config.moe_intermediate_size, - norm_topk_prob: if config.norm_topk_prob { 1 } else { 0 }, - expert_start: expert_start as i64, - expert_count: expert_count as i64, - intermediate_size: config.intermediate_size, - nccl_comm: std::ptr::null_mut(), - nccl_stream: std::ptr::null_mut(), - } - }).collect() + (0..config.mtp_num_hidden_layers) + .map(|_| { + CppLayerConfig { + layer_type: 0, // MTP layers are always full attention + num_heads: config.num_attention_heads, + num_kv_heads: config.num_key_value_heads, + head_dim: config.head_dim, + num_k_heads: config.linear_num_key_heads, + key_dim: config.linear_key_head_dim, + num_v_heads: config.linear_num_value_heads, + val_dim: config.linear_value_head_dim, + conv_kernel: config.linear_conv_kernel_dim, + partial_rotary_factor: config.partial_rotary_factor, + rope_theta: config.rope_theta, + rms_eps: config.rms_norm_eps, + num_experts: config.num_experts as i64, + top_k: config.num_experts_per_tok as i64, + moe_intermediate: config.moe_intermediate_size, + norm_topk_prob: if config.norm_topk_prob { 1 } else { 0 }, + expert_start: expert_start as i64, + expert_count: expert_count as i64, + intermediate_size: config.intermediate_size, + nccl_comm: std::ptr::null_mut(), + nccl_stream: std::ptr::null_mut(), + } + }) + .collect() } /// Opaque training context handle. @@ -327,10 +421,14 @@ impl CppTrainingContext { weights: &std::collections::BTreeMap, config: &crate::config::Qwen36RuntimeConfig, compute_kind: Kind, - lr: f64, beta1: f64, beta2: f64, eps: f64, + lr: f64, + beta1: f64, + beta2: f64, + eps: f64, lora_scaling: f64, lora_rank: i64, target_layers: &[usize], + target_modules: &[Qwen36LoraTargetModule], expert_start: usize, expert_count: usize, ) -> Result { @@ -338,7 +436,10 @@ impl CppTrainingContext { let mut weight_ptrs = build_weight_ptrs(weights, config); let layer_configs = build_layer_configs(config, expert_start, expert_count); - let embed_ptr = get_ptr(weights, &format!("{}embed_tokens.weight", config.weight_prefix)); + let embed_ptr = get_ptr( + weights, + &format!("{}embed_tokens.weight", config.weight_prefix), + ); let final_norm_ptr = get_ptr(weights, &format!("{}norm.weight", config.weight_prefix)); let lm_head_ptr = if config.tie_word_embeddings { // Tied embeddings: use embed_tokens as lm_head @@ -367,15 +468,40 @@ impl CppTrainingContext { }; let tl_len = target_layers.len() as i64; + let module_names = target_modules + .iter() + .map(Qwen36LoraTargetModule::cpp_name) + .collect::>() + .join(","); + let module_names_c = std::ffi::CString::new(module_names) + .map_err(|_| anyhow::anyhow!("LoRA target module contains NUL"))?; + let modules_ptr = if target_modules.is_empty() { + std::ptr::null() + } else { + module_names_c.as_ptr() + }; + let ptr = unsafe { (kh.create_ctx)( - wp_ptr, wp_len as i64, - embed_ptr, final_norm_ptr, lm_head_ptr, - lc_ptr, config.num_hidden_layers as i64, - compute_type, lora_scaling, - lr, beta1, beta2, eps, - config.vocab_size, config.rms_norm_eps, - lora_rank, tl_ptr, tl_len, + wp_ptr, + wp_len as i64, + embed_ptr, + final_norm_ptr, + lm_head_ptr, + lc_ptr, + config.num_hidden_layers as i64, + compute_type, + lora_scaling, + lr, + beta1, + beta2, + eps, + config.vocab_size, + config.rms_norm_eps, + lora_rank, + tl_ptr, + tl_len, + modules_ptr, ) }; if ptr.is_null() { @@ -387,7 +513,12 @@ impl CppTrainingContext { /// Run one training step: forward + loss + backward + Adam update. /// Returns loss value. - pub fn train_step(&self, input_ids: &Tensor, target_mask: &Tensor, attention_mask: &Tensor) -> Result { + pub fn train_step( + &self, + input_ids: &Tensor, + target_mask: &Tensor, + attention_mask: &Tensor, + ) -> Result { let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; let loss = unsafe { (kh.train_step)( @@ -408,8 +539,12 @@ impl CppTrainingContext { /// n_total: total number of adapters. lora_rank: LoRA rank for N_max calc. /// Returns average loss across chunks. pub fn train_multi_lora( - &self, input_ids: &Tensor, target_mask: &Tensor, attention_mask: &Tensor, - n_total: i32, lora_rank: i32, + &self, + input_ids: &Tensor, + target_mask: &Tensor, + attention_mask: &Tensor, + n_total: i32, + lora_rank: i32, ) -> Result { let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; let loss = unsafe { @@ -432,7 +567,9 @@ impl CppTrainingContext { pub fn get_lora_a(&self, index: i64) -> Option { let kh = get_kernels()?; let ptr = unsafe { (kh.get_lora_a)(self.ptr, index) }; - if ptr.is_null() { return None; } + if ptr.is_null() { + return None; + } Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) }) } @@ -440,10 +577,28 @@ impl CppTrainingContext { pub fn get_lora_b(&self, index: i64) -> Option { let kh = get_kernels()?; let ptr = unsafe { (kh.get_lora_b)(self.ptr, index) }; - if ptr.is_null() { return None; } + if ptr.is_null() { + return None; + } Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) }) } + pub fn set_lora_tensor(&self, index: i64, is_b: bool, tensor: &Tensor) -> Result<()> { + let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; + let status = unsafe { + (kh.set_lora_tensor)( + self.ptr, + index, + if is_b { 1 } else { 0 }, + tensor.as_ptr() as *mut c_void, + ) + }; + if status != 0 { + bail!("C++ set_lora_tensor failed for slot {index}"); + } + Ok(()) + } + pub fn lora_count(&self) -> i64 { self.lora_count } @@ -502,7 +657,13 @@ impl CppTrainingContext { /// Set NCCL communicator for Expert Parallel all-reduce. /// Must be called after `new()` if EP is enabled. /// comm_ptr / stream_ptr from `NcclPersistentComm::raw_comm_ptr()` / `raw_stream_ptr()`. - pub fn set_nccl_comm(&self, comm_ptr: *mut c_void, stream_ptr: *mut c_void, ep_rank: i32, ep_world_size: i32) { + pub fn set_nccl_comm( + &self, + comm_ptr: *mut c_void, + stream_ptr: *mut c_void, + ep_rank: i32, + ep_world_size: i32, + ) { let kh = get_kernels().expect("kernels not loaded"); unsafe { (kh.set_nccl_comm)(self.ptr, comm_ptr, stream_ptr, ep_rank, ep_world_size); @@ -528,17 +689,26 @@ impl CppTrainingContext { /// target_layers: empty = all layers /// target_modules: comma-separated, e.g. "q_proj,k_proj,v_proj,o_proj". Empty = all. pub fn add_lora( - &self, rank: i64, alpha: f64, - target_layers: &[i64], target_modules: &str, + &self, + rank: i64, + alpha: f64, + target_layers: &[i64], + target_modules: &str, ) -> Result { let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; - let tl_ptr = if target_layers.is_empty() { std::ptr::null() } else { target_layers.as_ptr() }; + let tl_ptr = if target_layers.is_empty() { + std::ptr::null() + } else { + target_layers.as_ptr() + }; let tl_len = target_layers.len() as i64; let modules_c = std::ffi::CString::new(target_modules).unwrap(); - let modules_ptr = if target_modules.is_empty() { std::ptr::null() } else { modules_c.as_ptr() }; - let id = unsafe { - (kh.add_lora)(self.ptr, rank, alpha, tl_ptr, tl_len, modules_ptr) + let modules_ptr = if target_modules.is_empty() { + std::ptr::null() + } else { + modules_c.as_ptr() }; + let id = unsafe { (kh.add_lora)(self.ptr, rank, alpha, tl_ptr, tl_len, modules_ptr) }; if id < 0 { bail!("C++ add_lora failed"); } @@ -547,6 +717,9 @@ impl CppTrainingContext { /// Remove a LoRA adapter by ID. pub fn remove_lora(&self, adapter_id: i64) -> Result { + if adapter_id == 0 { + return Ok(false); + } let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; let found = unsafe { (kh.remove_lora)(self.ptr, adapter_id) }; Ok(found != 0) @@ -554,15 +727,80 @@ impl CppTrainingContext { /// List all active adapter IDs. pub fn list_lora(&self) -> Vec { - let kh = match get_kernels() { Some(k) => k, None => return Vec::new() }; - let mut ids = vec![0i64; 64]; - let count = unsafe { (kh.list_lora)(self.ptr, ids.as_mut_ptr(), 64) }; - ids.truncate(count as usize); + let kh = match get_kernels() { + Some(k) => k, + None => return Vec::new(), + }; + let mut ids = Vec::with_capacity(65); + if self.lora_count > 0 { + // ID 0 is the fixed adapter created with the training context. + ids.push(0); + } + let mut dynamic_ids = vec![0i64; 64]; + let count = unsafe { (kh.list_lora)(self.ptr, dynamic_ids.as_mut_ptr(), 64) }; + ids.extend_from_slice(&dynamic_ids[..count as usize]); ids } + pub fn get_adapter_lora_tensor( + &self, + adapter_id: i64, + layer: i64, + module: &str, + is_b: bool, + ) -> Option { + let kh = get_kernels()?; + let module = std::ffi::CString::new(module).ok()?; + let ptr = unsafe { + (kh.get_adapter_lora_tensor)( + self.ptr, + adapter_id, + layer, + module.as_ptr(), + if is_b { 1 } else { 0 }, + ) + }; + if ptr.is_null() { + return None; + } + Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) }) + } + + pub fn set_adapter_lora_tensor( + &self, + adapter_id: i64, + layer: i64, + module: &str, + is_b: bool, + tensor: &Tensor, + ) -> Result<()> { + let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; + let module = std::ffi::CString::new(module)?; + let status = unsafe { + (kh.set_adapter_lora_tensor)( + self.ptr, + adapter_id, + layer, + module.as_ptr(), + if is_b { 1 } else { 0 }, + tensor.as_ptr() as *mut c_void, + ) + }; + if status != 0 { + bail!( + "C++ set_adapter_lora_tensor failed for adapter {adapter_id}, layer {layer}, module {module:?}" + ); + } + Ok(()) + } + /// Eval step: forward + loss, no backward, no Adam update. - pub fn eval_step(&self, input_ids: &Tensor, target_mask: &Tensor, attention_mask: &Tensor) -> Result { + pub fn eval_step( + &self, + input_ids: &Tensor, + target_mask: &Tensor, + attention_mask: &Tensor, + ) -> Result { let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; let loss = unsafe { (kh.eval_step)( @@ -580,7 +818,10 @@ impl CppTrainingContext { /// Get current training step count. pub fn get_step_count(&self) -> i64 { - let kh = match get_kernels() { Some(k) => k, None => return 0 }; + let kh = match get_kernels() { + Some(k) => k, + None => return 0, + }; unsafe { (kh.get_step_count)(self.ptr) } } @@ -588,16 +829,11 @@ impl CppTrainingContext { /// Returns (m_tensors, v_tensors) — owned copies on CPU. pub fn export_optimizer_state(&self) -> Result<(Vec, Vec)> { let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; - let count = self.lora_count * 2; // m and v per LoRA param (a+b) + let count = self.lora_count * 2; // m and v per LoRA param (a+b) let mut m_ptrs: Vec<*mut c_void> = vec![std::ptr::null_mut(); count as usize]; let mut v_ptrs: Vec<*mut c_void> = vec![std::ptr::null_mut(); count as usize]; let actual = unsafe { - (kh.export_optimizer)( - self.ptr, - m_ptrs.as_mut_ptr(), - v_ptrs.as_mut_ptr(), - count, - ) + (kh.export_optimizer)(self.ptr, m_ptrs.as_mut_ptr(), v_ptrs.as_mut_ptr(), count) }; let mut m_tensors = Vec::new(); let mut v_tensors = Vec::new(); @@ -613,11 +849,21 @@ impl CppTrainingContext { } /// Import Adam optimizer state (m and v vectors). - pub fn import_optimizer_state(&self, m_tensors: &[Tensor], v_tensors: &[Tensor]) -> Result { + pub fn import_optimizer_state( + &self, + m_tensors: &[Tensor], + v_tensors: &[Tensor], + ) -> Result { let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; let count = m_tensors.len().min(v_tensors.len()); - let m_ptrs: Vec<*mut c_void> = m_tensors.iter().map(|t| t.as_ptr() as *mut c_void).collect(); - let v_ptrs: Vec<*mut c_void> = v_tensors.iter().map(|t| t.as_ptr() as *mut c_void).collect(); + let m_ptrs: Vec<*mut c_void> = m_tensors + .iter() + .map(|t| t.as_ptr() as *mut c_void) + .collect(); + let v_ptrs: Vec<*mut c_void> = v_tensors + .iter() + .map(|t| t.as_ptr() as *mut c_void) + .collect(); let imported = unsafe { (kh.import_optimizer)( self.ptr, diff --git a/crates/rustrain-qwen3-6/src/lora.rs b/crates/rustrain-qwen3-6/src/lora.rs index f88c13ee..a3396164 100644 --- a/crates/rustrain-qwen3-6/src/lora.rs +++ b/crates/rustrain-qwen3-6/src/lora.rs @@ -1,14 +1,18 @@ //! Qwen3.6 LoRA adapter registry — VarStore-backed, stores tensors directly. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; +use std::fs; +use std::path::{Path, PathBuf}; -use anyhow::{Context, Result, bail}; +use anyhow::{bail, Context, Result}; +use serde::{Deserialize, Serialize}; use tch::{nn, Kind, Tensor}; use tracing::info; -use crate::config::Qwen36RuntimeConfig; +use crate::config::{LayerType, Qwen36RuntimeConfig}; -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] pub enum Qwen36LoraTargetModule { QProj, KProj, @@ -16,10 +20,95 @@ pub enum Qwen36LoraTargetModule { OProj, InProjQkv, InProjZ, + InProjA, + InProjB, OutProj, + GateProj, + UpProj, + DownProj, SharedGateProj, SharedUpProj, SharedDownProj, + ExpertsGateUpProj, + ExpertsDownProj, +} + +#[derive(Debug, Clone)] +pub struct Qwen36NativeLoraSlot { + pub index: usize, + pub layer: usize, + pub module: Qwen36LoraTargetModule, + pub active: bool, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct Qwen36AdapterConfig { + pub format_version: u32, + pub peft_type: String, + pub task_type: String, + pub base_model_name_or_path: String, + pub rustrain_architecture: String, + pub model_family: String, + pub r: i64, + pub lora_alpha: f64, + pub target_layers: Vec, + pub target_modules: Vec, + pub adapter_dtype: String, + pub bias: String, + pub inference_mode: bool, +} + +pub struct Qwen36AdapterArtifact { + pub config: Qwen36AdapterConfig, + pub tensors: BTreeMap, +} + +#[derive(Debug, Serialize)] +struct PeftAdapterConfig<'a> { + peft_type: &'static str, + task_type: &'static str, + base_model_name_or_path: &'a str, + r: i64, + lora_alpha: f64, + lora_dropout: f64, + target_modules: &'a [String], + layers_to_transform: &'a [usize], + layers_pattern: &'static str, + bias: &'static str, + fan_in_fan_out: bool, + inference_mode: bool, +} + +#[derive(Debug, Deserialize)] +struct PeftAdapterConfigOwned { + #[serde(default = "default_lora_type")] + peft_type: String, + #[serde(default = "default_task_type")] + task_type: String, + #[serde(default)] + base_model_name_or_path: String, + r: i64, + lora_alpha: f64, + #[serde(default)] + target_modules: Vec, + #[serde(default)] + layers_to_transform: Vec, + #[serde(default = "default_bias")] + bias: String, + #[serde(default)] + inference_mode: bool, +} + +fn default_lora_type() -> String { + "LORA".to_string() +} + +fn default_task_type() -> String { + "CAUSAL_LM".to_string() +} + +fn default_bias() -> String { + "none".to_string() } impl Qwen36LoraTargetModule { @@ -31,10 +120,17 @@ impl Qwen36LoraTargetModule { "o_proj" => Ok(Self::OProj), "in_proj_qkv" => Ok(Self::InProjQkv), "in_proj_z" => Ok(Self::InProjZ), + "in_proj_a" => Ok(Self::InProjA), + "in_proj_b" => Ok(Self::InProjB), "out_proj" => Ok(Self::OutProj), + "gate_proj" => Ok(Self::GateProj), + "up_proj" => Ok(Self::UpProj), + "down_proj" => Ok(Self::DownProj), "shared_gate_proj" => Ok(Self::SharedGateProj), "shared_up_proj" => Ok(Self::SharedUpProj), "shared_down_proj" => Ok(Self::SharedDownProj), + "experts_gate_up_proj" => Ok(Self::ExpertsGateUpProj), + "experts_down_proj" => Ok(Self::ExpertsDownProj), other => bail!("unknown LoRA target module: {other}"), } } @@ -47,12 +143,529 @@ impl Qwen36LoraTargetModule { Self::OProj => "self_attn.o_proj", Self::InProjQkv => "linear_attn.in_proj_qkv", Self::InProjZ => "linear_attn.in_proj_z", + Self::InProjA => "linear_attn.in_proj_a", + Self::InProjB => "linear_attn.in_proj_b", Self::OutProj => "linear_attn.out_proj", + Self::GateProj => "mlp.gate_proj", + Self::UpProj => "mlp.up_proj", + Self::DownProj => "mlp.down_proj", Self::SharedGateProj => "mlp.shared_expert.gate_proj", Self::SharedUpProj => "mlp.shared_expert.up_proj", Self::SharedDownProj => "mlp.shared_expert.down_proj", + Self::ExpertsGateUpProj => "mlp.experts.gate_up_proj", + Self::ExpertsDownProj => "mlp.experts.down_proj", + } + } + + /// Stable C++ projection identifier used by the native training context. + pub fn cpp_name(&self) -> &'static str { + match self { + Self::QProj => "q_proj", + Self::KProj => "k_proj", + Self::VProj => "v_proj", + Self::OProj => "o_proj", + Self::InProjQkv => "in_proj_qkv", + Self::InProjZ => "in_proj_z", + Self::InProjA => "in_proj_a", + Self::InProjB => "in_proj_b", + Self::OutProj => "out_proj", + Self::GateProj => "gate_proj", + Self::UpProj => "up_proj", + Self::DownProj => "down_proj", + Self::SharedGateProj => "shared_gate_proj", + Self::SharedUpProj => "shared_up_proj", + Self::SharedDownProj => "shared_down_proj", + Self::ExpertsGateUpProj => "experts_gate_up_proj", + Self::ExpertsDownProj => "experts_down_proj", + } + } +} + +pub fn native_lora_slots( + config: &Qwen36RuntimeConfig, + lora_config: &Qwen36LoraConfig, +) -> Vec { + let target_layers = lora_config + .target_layers + .iter() + .copied() + .collect::>(); + let target_modules = lora_config + .target_modules + .iter() + .copied() + .collect::>(); + let all_layers = target_layers.is_empty(); + let all_modules = target_modules.is_empty(); + let mut slots = Vec::new(); + + for (layer, layer_type) in config.layer_types.iter().enumerate() { + let modules: &[Qwen36LoraTargetModule] = match layer_type { + LayerType::FullAttention => &[ + Qwen36LoraTargetModule::QProj, + Qwen36LoraTargetModule::KProj, + Qwen36LoraTargetModule::VProj, + Qwen36LoraTargetModule::OProj, + ], + LayerType::LinearAttention => &[ + Qwen36LoraTargetModule::InProjQkv, + Qwen36LoraTargetModule::InProjZ, + Qwen36LoraTargetModule::InProjA, + Qwen36LoraTargetModule::InProjB, + Qwen36LoraTargetModule::OutProj, + ], + }; + for &module in modules { + slots.push(Qwen36NativeLoraSlot { + index: slots.len(), + layer, + module, + active: (all_layers || target_layers.contains(&layer)) + && (all_modules || target_modules.contains(&module)), + }); + } + let mlp_modules: &[Qwen36LoraTargetModule] = if config.is_moe { + &[ + Qwen36LoraTargetModule::SharedGateProj, + Qwen36LoraTargetModule::SharedUpProj, + Qwen36LoraTargetModule::SharedDownProj, + Qwen36LoraTargetModule::ExpertsGateUpProj, + Qwen36LoraTargetModule::ExpertsDownProj, + ] + } else { + &[ + Qwen36LoraTargetModule::GateProj, + Qwen36LoraTargetModule::UpProj, + Qwen36LoraTargetModule::DownProj, + ] + }; + for &module in mlp_modules { + slots.push(Qwen36NativeLoraSlot { + index: slots.len(), + layer, + module, + active: (all_layers || target_layers.contains(&layer)) + && (all_modules || target_modules.contains(&module)), + }); + } + } + slots +} + +pub fn validate_lora_targets( + runtime_config: &Qwen36RuntimeConfig, + lora_config: &Qwen36LoraConfig, +) -> Result<()> { + let slots = native_lora_slots(runtime_config, lora_config); + if !slots.iter().any(|slot| slot.active) { + bail!("LoRA targets do not resolve to any projection in this model"); + } + if !runtime_config.is_moe + && lora_config.target_modules.iter().any(|module| { + matches!( + module, + Qwen36LoraTargetModule::SharedGateProj + | Qwen36LoraTargetModule::SharedUpProj + | Qwen36LoraTargetModule::SharedDownProj + | Qwen36LoraTargetModule::ExpertsGateUpProj + | Qwen36LoraTargetModule::ExpertsDownProj + ) + }) + { + bail!("shared expert LoRA targets require a MoE Qwen model"); + } + if runtime_config.is_moe + && lora_config.target_modules.iter().any(|module| { + matches!( + module, + Qwen36LoraTargetModule::GateProj + | Qwen36LoraTargetModule::UpProj + | Qwen36LoraTargetModule::DownProj + ) + }) + { + bail!("dense MLP LoRA targets require a dense Qwen model"); + } + Ok(()) +} + +fn adapter_tensor_prefix( + config: &Qwen36RuntimeConfig, + layer: usize, + module: Qwen36LoraTargetModule, +) -> String { + format!( + "base_model.model.{}layers.{layer}.{}", + config.weight_prefix, + module.suffix() + ) +} + +/// Parse a canonical PEFT tensor name back to its layer, projection and side. +pub fn parse_adapter_tensor_name( + config: &Qwen36RuntimeConfig, + name: &str, +) -> Result<(usize, Qwen36LoraTargetModule, bool)> { + let (base, is_b) = if let Some(base) = name.strip_suffix(".lora_A.weight") { + (base, false) + } else if let Some(base) = name.strip_suffix(".lora_B.weight") { + (base, true) + } else { + bail!("invalid LoRA tensor name: {name}"); + }; + let prefix = format!("base_model.model.{}layers.", config.weight_prefix); + let rest = base + .strip_prefix(&prefix) + .with_context(|| format!("LoRA tensor has unexpected module prefix: {name}"))?; + let (layer, suffix) = rest + .split_once('.') + .with_context(|| format!("LoRA tensor is missing projection path: {name}"))?; + let layer = layer + .parse::() + .with_context(|| format!("invalid LoRA layer in tensor name: {name}"))?; + let module = [ + Qwen36LoraTargetModule::QProj, + Qwen36LoraTargetModule::KProj, + Qwen36LoraTargetModule::VProj, + Qwen36LoraTargetModule::OProj, + Qwen36LoraTargetModule::InProjQkv, + Qwen36LoraTargetModule::InProjZ, + Qwen36LoraTargetModule::InProjA, + Qwen36LoraTargetModule::InProjB, + Qwen36LoraTargetModule::OutProj, + Qwen36LoraTargetModule::GateProj, + Qwen36LoraTargetModule::UpProj, + Qwen36LoraTargetModule::DownProj, + Qwen36LoraTargetModule::SharedGateProj, + Qwen36LoraTargetModule::SharedUpProj, + Qwen36LoraTargetModule::SharedDownProj, + Qwen36LoraTargetModule::ExpertsGateUpProj, + Qwen36LoraTargetModule::ExpertsDownProj, + ] + .into_iter() + .find(|module| module.suffix() == suffix) + .with_context(|| format!("unknown Qwen LoRA projection suffix: {suffix}"))?; + Ok((layer, module, is_b)) +} + +fn resolve_artifact_paths(path: &Path) -> (PathBuf, PathBuf) { + if path.extension().and_then(|extension| extension.to_str()) == Some("safetensors") { + let config_path = path + .parent() + .unwrap_or_else(|| Path::new(".")) + .join("adapter_config.json"); + (path.to_path_buf(), config_path) + } else { + ( + path.join("adapter_model.safetensors"), + path.join("adapter_config.json"), + ) + } +} + +fn rustrain_config_path(config_path: &Path) -> PathBuf { + config_path + .parent() + .unwrap_or_else(|| Path::new(".")) + .join("rustrain_adapter.json") +} + +impl Qwen36AdapterArtifact { + pub fn from_native_exports( + model_name: &str, + architecture: &str, + base_model_path: Option<&Path>, + runtime_config: &Qwen36RuntimeConfig, + lora_config: &Qwen36LoraConfig, + exported: Vec<(Tensor, Tensor)>, + ) -> Result { + let slots = native_lora_slots(runtime_config, lora_config); + if exported.len() != slots.len() { + bail!( + "native LoRA export returned {} slots, expected {} for this model", + exported.len(), + slots.len() + ); + } + + let target_layers = slots + .iter() + .filter(|slot| slot.active) + .map(|slot| slot.layer) + .collect::>() + .into_iter() + .collect::>(); + let target_modules = slots + .iter() + .filter(|slot| slot.active) + .map(|slot| slot.module.cpp_name().to_string()) + .collect::>() + .into_iter() + .collect::>(); + + let mut tensors = BTreeMap::new(); + for (slot, (a, b)) in slots.into_iter().zip(exported) { + if !slot.active { + continue; + } + let prefix = adapter_tensor_prefix(runtime_config, slot.layer, slot.module); + tensors.insert( + format!("{prefix}.lora_A.weight"), + a.to_device(tch::Device::Cpu).to_kind(Kind::Float), + ); + tensors.insert( + format!("{prefix}.lora_B.weight"), + b.to_device(tch::Device::Cpu).to_kind(Kind::Float), + ); + } + if tensors.is_empty() { + bail!("LoRA export contains no active target modules"); + } + + Ok(Self { + config: Qwen36AdapterConfig { + format_version: 1, + peft_type: "LORA".to_string(), + task_type: "CAUSAL_LM".to_string(), + base_model_name_or_path: base_model_path + .map(|path| path.to_string_lossy().into_owned()) + .unwrap_or_else(|| model_name.to_string()), + rustrain_architecture: architecture.to_string(), + model_family: "qwen3_hybrid_text".to_string(), + r: lora_config.rank, + lora_alpha: lora_config.alpha, + target_layers, + target_modules, + adapter_dtype: "float32".to_string(), + bias: "none".to_string(), + inference_mode: true, + }, + tensors, + }) + } + + pub fn save(&self, path: &Path) -> Result { + let (tensor_path, config_path) = resolve_artifact_paths(path); + if let Some(parent) = tensor_path.parent() { + fs::create_dir_all(parent) + .with_context(|| format!("failed to create {}", parent.display()))?; + } + let refs = self + .tensors + .iter() + .map(|(name, tensor)| (name.as_str(), tensor)) + .collect::>(); + Tensor::write_safetensors(&refs, &tensor_path) + .with_context(|| format!("failed to write {}", tensor_path.display()))?; + let peft_config = PeftAdapterConfig { + peft_type: "LORA", + task_type: "CAUSAL_LM", + base_model_name_or_path: &self.config.base_model_name_or_path, + r: self.config.r, + lora_alpha: self.config.lora_alpha, + lora_dropout: 0.0, + target_modules: &self.config.target_modules, + layers_to_transform: &self.config.target_layers, + layers_pattern: "layers", + bias: "none", + fan_in_fan_out: false, + inference_mode: self.config.inference_mode, + }; + fs::write(&config_path, serde_json::to_vec_pretty(&peft_config)?) + .with_context(|| format!("failed to write {}", config_path.display()))?; + let rustrain_path = rustrain_config_path(&config_path); + fs::write(&rustrain_path, serde_json::to_vec_pretty(&self.config)?) + .with_context(|| format!("failed to write {}", rustrain_path.display()))?; + Ok(tensor_path) + } + + pub fn load(path: &Path) -> Result { + let (tensor_path, config_path) = resolve_artifact_paths(path); + let rustrain_path = rustrain_config_path(&config_path); + let metadata_path = if rustrain_path.exists() { + &rustrain_path + } else { + // Backward compatibility with the first named artifact draft, + // which stored rustrain metadata directly in adapter_config.json. + &config_path + }; + let config: Qwen36AdapterConfig = if rustrain_path.exists() { + serde_json::from_slice( + &fs::read(metadata_path) + .with_context(|| format!("failed to read {}", metadata_path.display()))?, + ) + .with_context(|| format!("failed to parse {}", metadata_path.display()))? + } else { + // Standard PEFT exports do not carry rustrain-specific fields. + // Preserve their canonical metadata and fill only the native + // runtime fields that are not part of the PEFT schema. + let peft: PeftAdapterConfigOwned = serde_json::from_slice( + &fs::read(&config_path) + .with_context(|| format!("failed to read {}", config_path.display()))?, + ) + .with_context(|| format!("failed to parse {}", config_path.display()))?; + if peft.target_modules.is_empty() { + bail!("PEFT adapter_config.json has no target_modules"); + } + Qwen36AdapterConfig { + format_version: 1, + peft_type: peft.peft_type, + task_type: peft.task_type, + base_model_name_or_path: peft.base_model_name_or_path, + rustrain_architecture: "qwen3_hybrid_lora_sft".to_string(), + model_family: "qwen3_hybrid_text".to_string(), + r: peft.r, + lora_alpha: peft.lora_alpha, + target_layers: peft.layers_to_transform, + target_modules: peft.target_modules, + adapter_dtype: "float32".to_string(), + bias: peft.bias, + inference_mode: peft.inference_mode, + } + }; + if config.format_version != 1 || config.peft_type != "LORA" { + bail!( + "unsupported adapter format version/type: {}/{}", + config.format_version, + config.peft_type + ); + } + if config.r <= 0 || config.lora_alpha <= 0.0 { + bail!("adapter rank and alpha must be positive"); + } + + let tensors = Tensor::read_safetensors(&tensor_path) + .with_context(|| format!("failed to read {}", tensor_path.display()))? + .into_iter() + .collect::>(); + validate_adapter_tensors(&config, &tensors)?; + Ok(Self { config, tensors }) + } + + /// Convert the pre-registry positional export into the named v1 format. + /// No metadata is inferred: callers must provide the runtime target + /// contract that defined the positional slots. + pub fn load_legacy( + path: &Path, + model_name: &str, + architecture: &str, + runtime_config: &Qwen36RuntimeConfig, + lora_config: &Qwen36LoraConfig, + ) -> Result { + let positional = Tensor::read_safetensors(path) + .with_context(|| format!("failed to read legacy adapter {}", path.display()))? + .into_iter() + .collect::>(); + let slots = native_lora_slots(runtime_config, lora_config); + let active_slots = slots.iter().filter(|slot| slot.active).collect::>(); + let a_count = positional + .keys() + .filter(|key| key.starts_with("lora_a_")) + .count(); + let b_count = positional + .keys() + .filter(|key| key.starts_with("lora_b_")) + .count(); + if a_count != b_count || (a_count != slots.len() && a_count != active_slots.len()) { + bail!( + "legacy adapter slot count {a_count}/{b_count} does not match fixed layout {} or compact target layout {}", + slots.len(), + active_slots.len() + ); + } + let compact = a_count == active_slots.len() && a_count != slots.len(); + let mut exported = Vec::with_capacity(slots.len()); + let mut compact_index = 0usize; + for slot in &slots { + let source_index = if compact { + if !slot.active { + let placeholder = Tensor::zeros([], (Kind::Float, tch::Device::Cpu)); + exported.push((placeholder.shallow_clone(), placeholder)); + continue; + } + let index = compact_index; + compact_index += 1; + index + } else { + slot.index + }; + let a_name = format!("lora_a_{source_index}"); + let b_name = format!("lora_b_{source_index}"); + let a = positional + .get(&a_name) + .with_context(|| format!("legacy adapter missing {a_name}"))? + .shallow_clone(); + let b = positional + .get(&b_name) + .with_context(|| format!("legacy adapter missing {b_name}"))? + .shallow_clone(); + exported.push((a, b)); + } + Self::from_native_exports( + model_name, + architecture, + None, + runtime_config, + lora_config, + exported, + ) + } +} + +fn validate_adapter_tensors( + config: &Qwen36AdapterConfig, + tensors: &BTreeMap, +) -> Result<()> { + if tensors.is_empty() || tensors.len() % 2 != 0 { + bail!("adapter must contain paired LoRA A/B tensors"); + } + let mut a_count = 0usize; + for (name, a) in tensors { + let Some(prefix) = name.strip_suffix(".lora_A.weight") else { + if name.ends_with(".lora_B.weight") { + let prefix = name.trim_end_matches(".lora_B.weight"); + if !tensors.contains_key(&format!("{prefix}.lora_A.weight")) { + bail!("missing paired tensor {prefix}.lora_A.weight"); + } + continue; + } + bail!("unexpected adapter tensor name: {name}"); + }; + a_count += 1; + let b_name = format!("{prefix}.lora_B.weight"); + let b = tensors + .get(&b_name) + .with_context(|| format!("missing paired tensor {b_name}"))?; + let grouped_expert = prefix.contains(".mlp.experts."); + if grouped_expert { + if a.dim() != 3 || b.dim() != 3 { + bail!("routed expert adapter tensors must be rank-3: {name}, {b_name}"); + } + if a.size()[0] != b.size()[0] || a.size()[1] != config.r || b.size()[2] != config.r { + bail!( + "routed expert adapter shape/rank mismatch for {prefix}: expected rank {}", + config.r + ); + } + if a.size()[2] <= 0 || b.size()[1] <= 0 { + bail!("routed expert adapter has an empty feature dimension: {prefix}"); + } + } else { + if a.dim() != 2 || b.dim() != 2 { + bail!("adapter tensors must be matrices: {name}, {b_name}"); + } + if a.size()[0] != config.r || b.size()[1] != config.r { + bail!("adapter rank mismatch for {prefix}: expected {}", config.r); + } + if a.size()[1] <= 0 || b.size()[0] <= 0 { + bail!("adapter tensor has an empty feature dimension: {prefix}"); + } } } + if a_count == 0 { + bail!("adapter must contain at least one LoRA A tensor"); + } + Ok(()) } #[derive(Debug, Clone)] @@ -84,16 +697,57 @@ impl Qwen36LoraRegistry { let layer_prefix = format!("{}layers.{}", config.weight_prefix, layer_idx); for &module in &lora_config.target_modules { let weight_name = match module { - Qwen36LoraTargetModule::QProj => format!("{layer_prefix}.self_attn.q_proj.weight"), - Qwen36LoraTargetModule::KProj => format!("{layer_prefix}.self_attn.k_proj.weight"), - Qwen36LoraTargetModule::VProj => format!("{layer_prefix}.self_attn.v_proj.weight"), - Qwen36LoraTargetModule::OProj => format!("{layer_prefix}.self_attn.o_proj.weight"), - Qwen36LoraTargetModule::InProjQkv => format!("{layer_prefix}.linear_attn.in_proj_qkv.weight"), - Qwen36LoraTargetModule::InProjZ => format!("{layer_prefix}.linear_attn.in_proj_z.weight"), - Qwen36LoraTargetModule::OutProj => format!("{layer_prefix}.linear_attn.out_proj.weight"), - Qwen36LoraTargetModule::SharedGateProj => format!("{layer_prefix}.mlp.shared_expert.gate_proj.weight"), - Qwen36LoraTargetModule::SharedUpProj => format!("{layer_prefix}.mlp.shared_expert.up_proj.weight"), - Qwen36LoraTargetModule::SharedDownProj => format!("{layer_prefix}.mlp.shared_expert.down_proj.weight"), + Qwen36LoraTargetModule::QProj => { + format!("{layer_prefix}.self_attn.q_proj.weight") + } + Qwen36LoraTargetModule::KProj => { + format!("{layer_prefix}.self_attn.k_proj.weight") + } + Qwen36LoraTargetModule::VProj => { + format!("{layer_prefix}.self_attn.v_proj.weight") + } + Qwen36LoraTargetModule::OProj => { + format!("{layer_prefix}.self_attn.o_proj.weight") + } + Qwen36LoraTargetModule::InProjQkv => { + format!("{layer_prefix}.linear_attn.in_proj_qkv.weight") + } + Qwen36LoraTargetModule::InProjZ => { + format!("{layer_prefix}.linear_attn.in_proj_z.weight") + } + Qwen36LoraTargetModule::InProjA => { + format!("{layer_prefix}.linear_attn.in_proj_a.weight") + } + Qwen36LoraTargetModule::InProjB => { + format!("{layer_prefix}.linear_attn.in_proj_b.weight") + } + Qwen36LoraTargetModule::OutProj => { + format!("{layer_prefix}.linear_attn.out_proj.weight") + } + Qwen36LoraTargetModule::GateProj => { + format!("{layer_prefix}.mlp.gate_proj.weight") + } + Qwen36LoraTargetModule::UpProj => { + format!("{layer_prefix}.mlp.up_proj.weight") + } + Qwen36LoraTargetModule::DownProj => { + format!("{layer_prefix}.mlp.down_proj.weight") + } + Qwen36LoraTargetModule::SharedGateProj => { + format!("{layer_prefix}.mlp.shared_expert.gate_proj.weight") + } + Qwen36LoraTargetModule::SharedUpProj => { + format!("{layer_prefix}.mlp.shared_expert.up_proj.weight") + } + Qwen36LoraTargetModule::SharedDownProj => { + format!("{layer_prefix}.mlp.shared_expert.down_proj.weight") + } + Qwen36LoraTargetModule::ExpertsGateUpProj => { + format!("{layer_prefix}.mlp.experts.gate_up_proj") + } + Qwen36LoraTargetModule::ExpertsDownProj => { + format!("{layer_prefix}.mlp.experts.down_proj") + } }; let base_weight = match weights.get(&weight_name) { @@ -101,34 +755,77 @@ impl Qwen36LoraRegistry { None => { // Skip modules that don't exist for this layer type // (e.g., q_proj doesn't exist for linear attention layers) - tracing::debug!("skipping LoRA target {weight_name} — not found (layer {layer_idx} may be different attention type)"); + tracing::debug!( + "skipping LoRA target {weight_name} — not found (layer {layer_idx} may be different attention type)" + ); continue; } }; - let (out_features, in_features) = (base_weight.size()[0], base_weight.size()[1]); let name = module.suffix().replace('.', "_"); let scale = 1.0 / (lora_config.rank as f64).sqrt(); - let lora_a = p.randn(&format!("lora_a_{layer_idx}_{name}"), &[lora_config.rank, in_features], 0.0, scale); - let lora_b = p.zeros(&format!("lora_b_{layer_idx}_{name}"), &[out_features, lora_config.rank]); + let (lora_a, lora_b) = if base_weight.dim() == 3 { + let experts = base_weight.size()[0]; + let out_features = base_weight.size()[1]; + let in_features = base_weight.size()[2]; + ( + p.randn( + &format!("lora_a_{layer_idx}_{name}"), + &[experts, lora_config.rank, in_features], + 0.0, + scale, + ), + p.zeros( + &format!("lora_b_{layer_idx}_{name}"), + &[experts, out_features, lora_config.rank], + ), + ) + } else { + let out_features = base_weight.size()[0]; + let in_features = base_weight.size()[1]; + ( + p.randn( + &format!("lora_a_{layer_idx}_{name}"), + &[lora_config.rank, in_features], + 0.0, + scale, + ), + p.zeros( + &format!("lora_b_{layer_idx}_{name}"), + &[out_features, lora_config.rank], + ), + ) + }; adapters.insert((layer_idx, module), (lora_a, lora_b)); } } - Ok(Self { config: lora_config, var_store, adapters }) + Ok(Self { + config: lora_config, + var_store, + adapters, + }) } pub fn trainable_variables(&self) -> Vec { self.var_store.trainable_variables() } - pub fn adapter_tensors(&self, layer: usize, module: Qwen36LoraTargetModule) -> Option<(Tensor, Tensor)> { + pub fn adapter_tensors( + &self, + layer: usize, + module: Qwen36LoraTargetModule, + ) -> Option<(Tensor, Tensor)> { let (a, b) = self.adapters.get(&(layer, module))?; Some((a.shallow_clone(), b.shallow_clone())) } /// Get references to the actual VarStore adapter tensors (for training backward). - pub fn adapter_ref(&self, layer: usize, module: Qwen36LoraTargetModule) -> Option<(&Tensor, &Tensor)> { + pub fn adapter_ref( + &self, + layer: usize, + module: Qwen36LoraTargetModule, + ) -> Option<(&Tensor, &Tensor)> { let (a, b) = self.adapters.get(&(layer, module))?; Some((a, b)) } @@ -138,7 +835,10 @@ impl Qwen36LoraRegistry { } pub fn trainable_param_count(&self) -> usize { - self.trainable_variables().iter().map(|v| v.numel() as usize).sum() + self.trainable_variables() + .iter() + .map(|v| v.numel() as usize) + .sum() } pub fn save(&self, path: &std::path::Path) -> Result<()> { @@ -156,16 +856,22 @@ impl Qwen36LoraRegistry { let mut header = serde_json::Map::new(); let mut offset = 0u64; for (name, tensor) in &named { - let t = tensor.to_device(tch::Device::Cpu).contiguous().to_kind(Kind::Float); + let t = tensor + .to_device(tch::Device::Cpu) + .contiguous() + .to_kind(Kind::Float); let shape: Vec = t.size().iter().copied().map(|d| d).collect(); let t_flat = t.reshape([-1]); let data: Vec = Vec::::try_from(&t_flat)?; let bytes: Vec = data.iter().flat_map(|f| f.to_le_bytes()).collect(); - header.insert(name.clone(), serde_json::json!({ - "dtype": "F32", - "shape": shape, - "data_offsets": [offset, offset + bytes.len() as u64], - })); + header.insert( + name.clone(), + serde_json::json!({ + "dtype": "F32", + "shape": shape, + "data_offsets": [offset, offset + bytes.len() as u64], + }), + ); offset += bytes.len() as u64; tensors_data.push(bytes); } @@ -185,3 +891,133 @@ impl Qwen36LoraRegistry { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn runtime(prefix: &str, layer_types: Vec) -> Qwen36RuntimeConfig { + Qwen36RuntimeConfig { + num_hidden_layers: layer_types.len(), + hidden_size: 16, + vocab_size: 32, + rms_norm_eps: 1e-6, + tie_word_embeddings: true, + hidden_act: "silu".into(), + layer_types, + full_attention_interval: 4, + num_attention_heads: 2, + num_key_value_heads: 2, + head_dim: 8, + attention_bias: false, + attn_output_gate: false, + rope_theta: 1e6, + partial_rotary_factor: 1.0, + mrope_interleaved: false, + mrope_section: vec![], + linear_num_key_heads: 2, + linear_key_head_dim: 8, + linear_num_value_heads: 2, + linear_value_head_dim: 8, + linear_conv_kernel_dim: 4, + mamba_ssm_dtype: "float32".into(), + is_moe: false, + num_experts: 0, + num_experts_per_tok: 0, + moe_intermediate_size: 0, + shared_expert_intermediate_size: 0, + norm_topk_prob: true, + router_aux_loss_coef: 0.0, + intermediate_size: 32, + mtp_num_hidden_layers: 0, + mtp_use_dedicated_embeddings: false, + has_vision: prefix.contains("language_model"), + vision_depth: 0, + vision_hidden_size: 0, + vision_num_heads: 0, + vision_patch_size: 0, + vision_spatial_merge_size: 0, + vision_temporal_patch_size: 0, + vision_out_hidden_size: 0, + weight_prefix: prefix.into(), + } + } + + #[test] + fn native_slots_match_cpp_projection_order() { + let config = runtime( + "model.", + vec![LayerType::FullAttention, LayerType::LinearAttention], + ); + let lora = Qwen36LoraConfig { + rank: 2, + alpha: 4.0, + target_layers: vec![], + target_modules: vec![], + }; + let slots = native_lora_slots(&config, &lora); + let modules = slots + .iter() + .map(|slot| slot.module.cpp_name()) + .collect::>(); + assert_eq!( + modules, + [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "in_proj_qkv", + "in_proj_z", + "out_proj" + ] + ); + } + + #[test] + fn parses_text_and_multimodal_peft_names() { + let text = runtime("model.", vec![LayerType::FullAttention]); + assert_eq!( + parse_adapter_tensor_name( + &text, + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" + ) + .unwrap(), + (0, Qwen36LoraTargetModule::QProj, false) + ); + let multimodal = runtime("model.language_model.", vec![LayerType::LinearAttention]); + assert_eq!( + parse_adapter_tensor_name( + &multimodal, + "base_model.model.model.language_model.layers.0.linear_attn.in_proj_qkv.lora_B.weight" + ) + .unwrap(), + (0, Qwen36LoraTargetModule::InProjQkv, true) + ); + } + + #[test] + fn rejects_unpaired_adapter_tensor() { + let config = Qwen36AdapterConfig { + format_version: 1, + peft_type: "LORA".into(), + task_type: "CAUSAL_LM".into(), + base_model_name_or_path: "Qwen/test".into(), + rustrain_architecture: "qwen3.6".into(), + model_family: "qwen3_hybrid_text".into(), + r: 2, + lora_alpha: 4.0, + target_layers: vec![0], + target_modules: vec!["q_proj".into()], + adapter_dtype: "float32".into(), + bias: "none".into(), + inference_mode: true, + }; + let mut tensors = BTreeMap::new(); + tensors.insert( + "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight".into(), + Tensor::zeros([4, 2], (Kind::Float, tch::Device::Cpu)), + ); + assert!(validate_adapter_tensors(&config, &tensors).is_err()); + } +} diff --git a/crates/rustrain-qwen3-6/src/session.rs b/crates/rustrain-qwen3-6/src/session.rs index ee2a90ab..3e4be84d 100644 --- a/crates/rustrain-qwen3-6/src/session.rs +++ b/crates/rustrain-qwen3-6/src/session.rs @@ -7,11 +7,15 @@ use anyhow::{Context, Result, anyhow, bail}; use tch::{Kind, Tensor}; use tracing::info; -use crate::config::{read_qwen36_runtime_config, resolve_qwen36_model_path, Qwen36RuntimeConfig, LayerType}; -use crate::lora::{Qwen36LoraConfig, Qwen36LoraTargetModule}; +use crate::config::{ + LayerType, Qwen36RuntimeConfig, read_qwen36_runtime_config, resolve_qwen36_model_path, +}; +use crate::lora::{ + Qwen36AdapterArtifact, Qwen36LoraConfig, Qwen36LoraTargetModule, validate_lora_targets, +}; use crate::sft::SftDataset; +use rustrain_checkpoint::safetensors::read_safetensors_dir_filtered; use rustrain_core::runtime::{Config, RunPaths}; -use rustrain_checkpoint::safetensors::{read_safetensors_dir_filtered}; // ────────────────────────────────────────────────────────────────────── // EP Shard @@ -27,7 +31,10 @@ pub struct EpShard { impl EpShard { pub fn new(rank: usize, world_size: usize, num_experts: usize) -> Self { - assert!(num_experts % world_size == 0, "num_experts {num_experts} not divisible by world_size {world_size}"); + assert!( + num_experts % world_size == 0, + "num_experts {num_experts} not divisible by world_size {world_size}" + ); let epr = num_experts / world_size; let start = rank * epr; Self { @@ -62,11 +69,17 @@ pub struct Qwen36LoraSftSummary { fn parse_env_usize(key: &str) -> Result { env::var(key) .with_context(|| format!("{key} not set")) - .and_then(|v| v.parse::().with_context(|| format!("invalid {key}: {v}"))) + .and_then(|v| { + v.parse::() + .with_context(|| format!("invalid {key}: {v}")) + }) } fn lora_config_from_config(config: &Config) -> Result { - let lora = config.lora.as_ref().ok_or_else(|| anyhow!("[lora] section required"))?; + let lora = config + .lora + .as_ref() + .ok_or_else(|| anyhow!("[lora] section required"))?; let target_layers: Vec = lora.target_layers.clone(); let target_modules: Vec = lora .target_modules @@ -116,7 +129,13 @@ fn build_needed_weights( needed.insert(format!("{lp}.linear_attn.conv1d.weight")); needed.insert(format!("{lp}.linear_attn.dt_bias")); needed.insert(format!("{lp}.linear_attn.norm.weight")); - for w in &["in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b", "out_proj"] { + for w in &[ + "in_proj_qkv", + "in_proj_z", + "in_proj_a", + "in_proj_b", + "out_proj", + ] { needed.insert(format!("{lp}.linear_attn.{w}.weight")); } } @@ -172,7 +191,10 @@ pub fn train_qwen3_6_lora_sft_ep( let rank = parse_env_usize("RANK")?; let world_size = parse_env_usize("WORLD_SIZE")?; // For MoE models, shard experts. For dense models, EP is a no-op (no experts to shard). - let model_path = config.model.model_path.as_ref() + let model_path = config + .model + .model_path + .as_ref() .ok_or_else(|| anyhow!("model.model_path required"))?; let model_path = resolve_qwen36_model_path(model_path)?; let runtime_config = read_qwen36_runtime_config(&model_path)?; @@ -193,11 +215,15 @@ fn train_impl( run_paths: &RunPaths, ep_shard: Option, ) -> Result { - let model_path = config.model.model_path.as_ref() + let model_path = config + .model + .model_path + .as_ref() .ok_or_else(|| anyhow!("model.model_path required"))?; let model_path = resolve_qwen36_model_path(model_path)?; let runtime_config = read_qwen36_runtime_config(&model_path)?; let lora_config = lora_config_from_config(config)?; + validate_lora_targets(&runtime_config, &lora_config)?; let device = match config.train.device { rustrain_core::runtime::Device::Cuda => { // EP mode: use LOCAL_RANK to select the correct GPU @@ -214,6 +240,12 @@ fn train_impl( rustrain_core::runtime::DType::Bf16 => Kind::BFloat16, rustrain_core::runtime::DType::Fp32 => Kind::Float, }; + if compute_kind != Kind::BFloat16 { + bail!( + "native Qwen3.5/3.6 LoRA currently supports bf16 only; {:?} would not be updated by the fused Adam kernel", + config.train.dtype + ); + } let shard_ref = ep_shard.as_ref(); let is_ep = shard_ref.is_some(); @@ -228,7 +260,11 @@ fn train_impl( std::thread::sleep(std::time::Duration::from_secs(rank as u64 * 5)); } - info!("loading {} weight tensors from {}", needed.len(), model_path.display()); + info!( + "loading {} weight tensors from {}", + needed.len(), + model_path.display() + ); let weights = read_safetensors_dir_filtered(&model_path, &needed)?; // Move to device — for EP, narrow expert tensors on CPU first to save GPU memory @@ -256,14 +292,20 @@ fn train_impl( weights_gpu.insert(name.clone(), tensor.to_device(device).to_kind(compute_kind)); } } - info!("EP{}: narrowed expert tensors to {} experts per rank", world_size, shard.experts_per_rank); + info!( + "EP{}: narrowed expert tensors to {} experts per rank", + world_size, shard.experts_per_rank + ); } else { for (name, tensor) in &weights { weights_gpu.insert(name.clone(), tensor.to_device(device).to_kind(compute_kind)); } } - info!("LoRA config: rank={}, alpha={}", lora_config.rank, lora_config.alpha); + info!( + "LoRA config: rank={}, alpha={}", + lora_config.rank, lora_config.alpha + ); // Load SFT data let tokenizer_path = model_path.join("tokenizer.json"); @@ -294,18 +336,23 @@ fn train_impl( // ── C++ all-in-C++ training path (required) ── // LoRA A/B, Adam optimizer, forward, loss, backward all in C++. if !crate::kernel::kernels_available() { - bail!("C++ kernels (libqwen36_kernels.so) not found — required for training. Ensure the .so is in LD_LIBRARY_PATH."); + bail!( + "C++ kernels (libqwen36_kernels.so) not found — required for training. Ensure the .so is in LD_LIBRARY_PATH." + ); } let ctx = crate::kernel::CppTrainingContext::new( - &weights_gpu, &runtime_config, compute_kind, + &weights_gpu, + &runtime_config, + compute_kind, config.train.learning_rate as f64, config.train.adam_beta1 as f64, config.train.adam_beta2 as f64, config.train.adam_eps as f64, - lora_config.alpha as f64 / lora_config.rank as f64, // lora scaling = alpha / rank + lora_config.alpha as f64 / lora_config.rank as f64, // lora scaling = alpha / rank lora_config.rank as i64, &lora_config.target_layers, + &lora_config.target_modules, shard_ref.map(|s| s.expert_start).unwrap_or(0), shard_ref.map(|s| s.experts_per_rank).unwrap_or(0), )?; @@ -314,11 +361,15 @@ fn train_impl( // Set MTP weights if available if runtime_config.mtp_num_hidden_layers > 0 { ctx.set_mtp_weights( - &weights_gpu, &runtime_config, + &weights_gpu, + &runtime_config, shard_ref.map(|s| s.expert_start).unwrap_or(0), shard_ref.map(|s| s.experts_per_rank).unwrap_or(0), )?; - info!("C++ TrainingContext: MTP weights set ({} layers)", runtime_config.mtp_num_hidden_layers); + info!( + "C++ TrainingContext: MTP weights set ({} layers)", + runtime_config.mtp_num_hidden_layers + ); } // Enable gradient checkpointing if env var set @@ -349,53 +400,47 @@ fn train_impl( // but padding also has mask=0, we can't distinguish prompt from padding using mask alone. // Solution: use the pad_token_id to build attention mask from input_ids. let pad_id = data.pad_token_id(); - let attention_mask = input_ids.ne(pad_id).to_kind(Kind::Float).unsqueeze(0); // [1, seq] + let attention_mask = input_ids.ne(pad_id).to_kind(Kind::Float).unsqueeze(0); // [1, seq] // C++ all-in-C++ path: single call does forward + loss + backward + Adam let loss_value = ctx.train_step(&input_ids, &target_mask, &attention_mask)?; - if step == 0 { initial_loss = loss_value; } + if step == 0 { + initial_loss = loss_value; + } final_loss = loss_value; if step % 10 == 0 || step == max_steps - 1 { info!("step {step}/{max_steps} loss={loss_value:.6}"); } } - // Save adapter — export LoRA A/B from C++ to safetensors - let adapter_path = run_paths.root.join("adapter.safetensors"); - { - let mut named_tensors: BTreeMap = BTreeMap::new(); - for i in 0..ctx.lora_count() { - if let (Some(a), Some(b)) = (ctx.get_lora_a(i), ctx.get_lora_b(i)) { - named_tensors.insert(format!("lora_a_{i}"), a.to_kind(Kind::Float).to_device(tch::Device::Cpu)); - named_tensors.insert(format!("lora_b_{i}"), b.to_kind(Kind::Float).to_device(tch::Device::Cpu)); - } - } - use std::io::Write; - let mut tensors_data = Vec::new(); - let mut header = serde_json::Map::new(); - let mut offset = 0u64; - for (name, tensor) in &named_tensors { - let t = tensor.contiguous().to_kind(Kind::Float); - let shape: Vec = t.size().iter().copied().collect(); - let data: Vec = Vec::::try_from(&t.reshape([-1]))?; - let bytes: Vec = data.iter().flat_map(|f| f.to_le_bytes()).collect(); - header.insert(name.clone(), serde_json::json!({"dtype":"F32","shape":shape,"data_offsets":[offset,offset+bytes.len() as u64]})); - offset += bytes.len() as u64; - tensors_data.push(bytes); - } - let header_str = serde_json::to_string(&serde_json::Value::Object(header))?; - let file = std::fs::File::create(&adapter_path).with_context(|| format!("create {}", adapter_path.display()))?; - let mut writer = std::io::BufWriter::new(file); - writer.write_all(&(header_str.len() as u64).to_le_bytes())?; - writer.write_all(header_str.as_bytes())?; - for data in &tensors_data { writer.write_all(data)?; } - info!("saved adapter to {}", adapter_path.display()); + // Export every positional native slot, then let the artifact mapper omit + // inactive slots and assign stable projection-aware tensor names. + let mut exported = Vec::with_capacity(ctx.lora_count() as usize); + for index in 0..ctx.lora_count() { + let a = ctx + .get_lora_a(index) + .with_context(|| format!("native LoRA slot {index} is missing A"))?; + let b = ctx + .get_lora_b(index) + .with_context(|| format!("native LoRA slot {index} is missing B"))?; + exported.push((a, b)); } + let artifact = Qwen36AdapterArtifact::from_native_exports( + &config.model.name, + &config.model.architecture, + Some(&model_path), + &runtime_config, + &lora_config, + exported, + )?; + let trainable_params = artifact.tensors.len(); + let adapter_path = artifact.save(&run_paths.root)?; + info!("saved adapter to {}", adapter_path.display()); Ok(Qwen36LoraSftSummary { adapter_output: adapter_path.to_string_lossy().to_string(), initial_loss, final_loss, - trainable_params: ctx.lora_count() as usize * 2, + trainable_params, }) } diff --git a/crates/rustrain-qwen3-6/tests/integration.rs b/crates/rustrain-qwen3-6/tests/integration.rs index cb291a6d..865c48fe 100644 --- a/crates/rustrain-qwen3-6/tests/integration.rs +++ b/crates/rustrain-qwen3-6/tests/integration.rs @@ -1,11 +1,205 @@ //! Integration tests for Qwen3.6 config parsing and forward pass -use std::collections::BTreeMap; - -use rustrain_qwen3_6::config::{read_qwen36_runtime_config, resolve_qwen36_model_path, LayerType}; -use rustrain_qwen3_6::model::{qwen36_forward_from_ids, Qwen36LayerWeights}; +use rustrain_qwen3_6::config::{LayerType, read_qwen36_runtime_config}; +use rustrain_qwen3_6::lora::{ + Qwen36AdapterArtifact, Qwen36LoraConfig, Qwen36LoraTargetModule, native_lora_slots, + validate_lora_targets, +}; +use rustrain_qwen3_6::model::Qwen36LayerWeights; const MODEL_PATH: &str = "/vePFS-Mindverse/share/huggingface/hub/models--Qwen--Qwen3.6-35B-A3B/snapshots/995ad96eacd98c81ed38be0c5b274b04031597b0"; +const QWEN35_MODEL_PATH: &str = "/vePFS-Mindverse/share/huggingface/hub/models--Qwen--Qwen3.5-0.8B/snapshots/2fc06364715b967f1860aea9cf38778875588b17"; + +#[test] +fn test_native_lora_target_module_contract() { + let names = [ + ("q_proj", "self_attn.q_proj"), + ("k_proj", "self_attn.k_proj"), + ("v_proj", "self_attn.v_proj"), + ("o_proj", "self_attn.o_proj"), + ("in_proj_qkv", "linear_attn.in_proj_qkv"), + ("in_proj_z", "linear_attn.in_proj_z"), + ("out_proj", "linear_attn.out_proj"), + ]; + for (cpp_name, suffix) in names { + let module = Qwen36LoraTargetModule::parse(cpp_name).expect("supported target"); + assert_eq!(module.cpp_name(), cpp_name); + assert_eq!(module.suffix(), suffix); + } + for name in [ + "in_proj_a", + "in_proj_b", + "gate_proj", + "up_proj", + "down_proj", + "shared_gate_proj", + "shared_up_proj", + "shared_down_proj", + "experts_gate_up_proj", + "experts_down_proj", + ] { + assert!(Qwen36LoraTargetModule::parse(name).is_ok(), "{name}"); + } +} + +#[test] +fn test_routed_expert_adapter_roundtrip_uses_rank3_local_shards() { + let runtime = + read_qwen36_runtime_config(std::path::Path::new(MODEL_PATH)).expect("Qwen3.6 config parse"); + let lora = Qwen36LoraConfig { + rank: 2, + alpha: 8.0, + target_layers: vec![0], + target_modules: vec![ + Qwen36LoraTargetModule::ExpertsGateUpProj, + Qwen36LoraTargetModule::ExpertsDownProj, + ], + }; + validate_lora_targets(&runtime, &lora).expect("MoE routed expert targets"); + let slots = native_lora_slots(&runtime, &lora); + let exported = slots + .iter() + .map(|slot| match slot.module { + Qwen36LoraTargetModule::ExpertsGateUpProj if slot.active => ( + tch::Tensor::ones([2, 2, 3], (tch::Kind::Float, tch::Device::Cpu)), + tch::Tensor::zeros([2, 4, 2], (tch::Kind::Float, tch::Device::Cpu)), + ), + Qwen36LoraTargetModule::ExpertsDownProj if slot.active => ( + tch::Tensor::ones([2, 2, 4], (tch::Kind::Float, tch::Device::Cpu)), + tch::Tensor::zeros([2, 3, 2], (tch::Kind::Float, tch::Device::Cpu)), + ), + _ => ( + tch::Tensor::zeros([], (tch::Kind::Float, tch::Device::Cpu)), + tch::Tensor::zeros([], (tch::Kind::Float, tch::Device::Cpu)), + ), + }) + .collect(); + let artifact = Qwen36AdapterArtifact::from_native_exports( + "Qwen3.6-35B-A3B", + "qwen3_6_lora_sft", + Some(std::path::Path::new(MODEL_PATH)), + &runtime, + &lora, + exported, + ) + .expect("build expert adapter artifact"); + + assert_eq!(artifact.tensors.len(), 4); + let key = format!( + "base_model.model.{}layers.0.mlp.experts.gate_up_proj.lora_A.weight", + runtime.weight_prefix + ); + assert_eq!(artifact.tensors[&key].size(), [2, 2, 3]); + + let temp = tempfile::tempdir().expect("temporary expert adapter directory"); + artifact.save(temp.path()).expect("save expert adapter"); + let loaded = Qwen36AdapterArtifact::load(temp.path()).expect("reload expert adapter"); + assert_eq!(loaded.tensors[&key].size(), [2, 2, 3]); + + let dense_runtime = read_qwen36_runtime_config(std::path::Path::new(QWEN35_MODEL_PATH)) + .expect("Qwen3.5 dense config parse"); + assert!(validate_lora_targets(&dense_runtime, &lora).is_err()); +} + +#[test] +fn test_adapter_artifact_roundtrip_uses_projection_names() { + let runtime = read_qwen36_runtime_config(std::path::Path::new(QWEN35_MODEL_PATH)) + .expect("Qwen3.5 config parse"); + let lora = Qwen36LoraConfig { + rank: 2, + alpha: 4.0, + target_layers: vec![0, 3], + target_modules: vec![ + Qwen36LoraTargetModule::InProjQkv, + Qwen36LoraTargetModule::QProj, + ], + }; + let slots = native_lora_slots(&runtime, &lora); + let exported = slots + .iter() + .map(|slot| { + let value = slot.index as f64 + 1.0; + ( + tch::Tensor::full([2, 3], value, (tch::Kind::Float, tch::Device::Cpu)), + tch::Tensor::full([5, 2], -value, (tch::Kind::Float, tch::Device::Cpu)), + ) + }) + .collect(); + let artifact = Qwen36AdapterArtifact::from_native_exports( + "Qwen3.5-0.8B", + "qwen3_5_lora_sft", + Some(std::path::Path::new(QWEN35_MODEL_PATH)), + &runtime, + &lora, + exported, + ) + .expect("build adapter artifact"); + + assert_eq!(artifact.tensors.len(), 4); + let linear_key = format!( + "base_model.model.{}layers.0.linear_attn.in_proj_qkv.lora_A.weight", + runtime.weight_prefix + ); + let full_key = format!( + "base_model.model.{}layers.3.self_attn.q_proj.lora_B.weight", + runtime.weight_prefix + ); + assert!(artifact.tensors.contains_key(&linear_key)); + assert!(artifact.tensors.contains_key(&full_key)); + + let temp = tempfile::tempdir().expect("temporary adapter directory"); + let tensor_path = artifact.save(temp.path()).expect("save adapter artifact"); + assert_eq!( + tensor_path.file_name().and_then(|name| name.to_str()), + Some("adapter_model.safetensors") + ); + assert!(temp.path().join("rustrain_adapter.json").is_file()); + let peft_config: serde_json::Value = serde_json::from_slice( + &std::fs::read(temp.path().join("adapter_config.json")).expect("read PEFT config"), + ) + .expect("parse PEFT config"); + assert_eq!(peft_config["peft_type"], "LORA"); + assert_eq!(peft_config["r"], 2); + assert_eq!( + peft_config["layers_to_transform"], + serde_json::json!([0, 3]) + ); + assert_eq!( + peft_config["target_modules"], + serde_json::json!(["in_proj_qkv", "q_proj"]) + ); + let loaded = Qwen36AdapterArtifact::load(temp.path()).expect("reload adapter artifact"); + assert_eq!(loaded.config, artifact.config); + assert_eq!(loaded.tensors.len(), artifact.tensors.len()); + for (name, expected) in &artifact.tensors { + let actual = loaded.tensors.get(name).expect("roundtrip tensor name"); + assert_eq!(actual.size(), expected.size()); + assert_eq!( + Vec::::try_from(&actual.reshape([-1])).expect("actual values"), + Vec::::try_from(&expected.reshape([-1])).expect("expected values") + ); + } + // A PEFT-only export is accepted when rustrain_adapter.json is absent. + std::fs::remove_file(temp.path().join("rustrain_adapter.json")) + .expect("remove private rustrain metadata"); + let peft_only = Qwen36AdapterArtifact::load(temp.path()).expect("load PEFT-only artifact"); + assert_eq!(peft_only.config.r, 2); + assert_eq!(peft_only.config.target_layers, vec![0, 3]); +} + +#[test] +fn test_qwen35_text_config_parsing() { + let config = read_qwen36_runtime_config(std::path::Path::new(QWEN35_MODEL_PATH)) + .expect("Qwen3.5 config parse"); + assert_eq!(config.num_hidden_layers, 24); + assert_eq!(config.hidden_size, 1024); + assert_eq!(config.layer_types[3], LayerType::FullAttention); + assert_eq!(config.full_attention_interval, 4); + assert_eq!(config.linear_num_key_heads, 16); + assert_eq!(config.linear_num_value_heads, 16); + assert!(config.attn_output_gate); + assert!(config.tie_word_embeddings); +} #[test] fn test_config_parsing() { @@ -99,21 +293,18 @@ fn test_forward_single_layer() { .map(|suffix| format!("{layer_prefix}.{suffix}")) .collect(); - let weights = read_safetensors_dir_filtered(model_path, &needed) - .expect("load layer 0 weights"); + let weights = read_safetensors_dir_filtered(model_path, &needed).expect("load layer 0 weights"); // Load layer 0 (linear attention) - let layer = Qwen36LayerWeights::load(&weights, &config, 0, tch::Kind::BFloat16) - .expect("load layer 0"); + let layer = + Qwen36LayerWeights::load(&weights, &config, 0, tch::Kind::BFloat16).expect("load layer 0"); // Run forward on CPU (no CUDA in this env) let device = tch::Device::Cpu; - let input = tch::Tensor::randn( - [1, 5, config.hidden_size], - (tch::Kind::BFloat16, device), - ); + let input = tch::Tensor::randn([1, 5, config.hidden_size], (tch::Kind::BFloat16, device)); - let output = rustrain_qwen3_6::model::qwen36_layer(&input, &layer, &config, tch::Kind::BFloat16); + let output = + rustrain_qwen3_6::model::qwen36_layer(&input, &layer, &config, tch::Kind::BFloat16); assert_eq!(output.size()[0], 1); assert_eq!(output.size()[1], 5); diff --git a/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp new file mode 100644 index 00000000..29799d57 --- /dev/null +++ b/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp @@ -0,0 +1,123 @@ +#include +#include + +#include +#include +#include +#include + +struct LayerConfig { + int64_t layer_type, num_heads, num_kv_heads, head_dim; + int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel; + double partial_rotary_factor, rope_theta, rms_eps; + int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count; + int64_t intermediate_size; + int32_t norm_topk_prob; + void* nccl_comm; + void* nccl_stream; +}; + +extern "C" void* qwen36_create_training_context( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*); +extern "C" int32_t qwen36_init_nccl(void*); +extern "C" void* qwen36_get_lora_b(void*, int64_t); +extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*); +extern "C" double qwen36_train_step(void*, void*, void*, void*); +extern "C" void qwen36_free_training_context(void*); + +static at::Tensor cuda_rand(std::initializer_list shape) { + return at::randn(shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)); +} + +int main() { + const int rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0"); + const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); + const int local_rank = std::atoi(std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); + assert(world == 2 && rank >= 0 && rank < world); + c10::cuda::CUDAGuard guard(local_rank); + // Replicated weights must be identical across EP ranks. + at::manual_seed(100); + + constexpr int64_t hidden = 16; + constexpr int64_t vocab = 8; + constexpr int64_t experts = 2; + constexpr int64_t head_dim = 8; + constexpr int64_t intermediate = 8; + constexpr int64_t rank_lora = 4; + std::vector weights; + weights.push_back(cuda_rand({hidden})); + weights.push_back(cuda_rand({hidden})); + weights.push_back(cuda_rand({2 * head_dim, hidden})); + weights.push_back(cuda_rand({head_dim})); + weights.push_back(cuda_rand({head_dim, hidden})); + weights.push_back(cuda_rand({head_dim})); + weights.push_back(cuda_rand({head_dim, hidden})); + weights.push_back(cuda_rand({hidden, head_dim})); + weights.push_back(cuda_rand({experts, hidden})); + weights.push_back(cuda_rand({1, hidden})); + weights.push_back(cuda_rand({intermediate, hidden})); + weights.push_back(cuda_rand({intermediate, hidden})); + weights.push_back(cuda_rand({hidden, intermediate})); + // Each process owns exactly one expert row. + weights.push_back(cuda_rand({1, 2 * intermediate, hidden})); + weights.push_back(cuda_rand({1, hidden, intermediate})); + for (auto& weight : weights) weight.set_requires_grad(false); + auto embed = cuda_rand({vocab, hidden}); + auto final_norm = cuda_rand({hidden}); + auto lm_head = cuda_rand({vocab, hidden}); + embed.set_requires_grad(false); + final_norm.set_requires_grad(false); + lm_head.set_requires_grad(false); + + std::vector weight_ptrs; + for (auto& weight : weights) weight_ptrs.push_back(&weight); + LayerConfig config{}; + config.layer_type = 0; + config.num_heads = 1; + config.num_kv_heads = 1; + config.head_dim = head_dim; + config.partial_rotary_factor = 1.0; + config.rope_theta = 10000.0; + config.rms_eps = 1e-5; + config.num_experts = experts; + // Route every token to both experts so both local shards exercise their + // expert LoRA optimizer path in this two-rank smoke. + config.top_k = 2; + config.moe_intermediate = intermediate; + config.expert_start = rank; + config.expert_count = 1; + config.norm_topk_prob = 1; + + const int64_t target_layer = 0; + void* ctx = qwen36_create_training_context( + weight_ptrs.data(), static_cast(weight_ptrs.size()), + &embed, &final_norm, &lm_head, &config, 1, + static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank_lora, + &target_layer, 1, "experts_gate_up_proj,experts_down_proj"); + assert(ctx); + assert(qwen36_init_nccl(ctx) == 0); + auto* lora_b = reinterpret_cast(qwen36_get_lora_b(ctx, 7)); + assert(lora_b); + auto lora_b_value = at::ones(lora_b->sizes(), lora_b->options()); + assert(qwen36_set_lora_tensor(ctx, 7, 1, &lora_b_value) == 0); + auto before = lora_b->clone(); + + auto input_ids = at::tensor({1, 2}, + at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 2}); + auto target_mask = at::ones({1, 2}, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); + auto attention_mask = at::ones({1, 2}, + at::TensorOptions().device(at::kCUDA).dtype(at::kBool)); + const double loss = qwen36_train_step(ctx, &input_ids, &target_mask, &attention_mask); + c10::cuda::device_synchronize(); + const double update = (*lora_b - before).abs().sum().item(); + std::printf("native_qwen36_ep_smoke rank=%d world=%d loss=%0.8f lora_b_update=%0.8e\n", + rank, world, loss, update); + assert(loss == loss && loss > 0.0); + assert(update > 0.0); + qwen36_free_training_context(ctx); + return 0; +} diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp new file mode 100644 index 00000000..b9364758 --- /dev/null +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -0,0 +1,390 @@ +#include +#include + +#include +#include +#include +#include +#include +#include + +struct LayerConfig { + int64_t layer_type, num_heads, num_kv_heads, head_dim; + int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel; + double partial_rotary_factor, rope_theta, rms_eps; + int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count; + int64_t intermediate_size; + int32_t norm_topk_prob; + void* nccl_comm; + void* nccl_stream; +}; + +extern "C" void* qwen36_create_training_context( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*); +extern "C" int64_t qwen36_get_lora_count(void*); +extern "C" void* qwen36_get_lora_a(void*, int64_t); +extern "C" void* qwen36_get_lora_b(void*, int64_t); +extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*); +extern "C" double qwen36_train_step(void*, void*, void*, void*); +extern "C" double qwen36_eval_step(void*, void*, void*, void*); +extern "C" double qwen36_train_multi_lora( + void*, void*, void*, void*, int32_t, int32_t); +extern "C" int64_t qwen36_add_lora( + void*, int64_t, double, const int64_t*, int64_t, const char*); +extern "C" void* qwen36_get_adapter_lora_tensor( + void*, int64_t, int64_t, const char*, int32_t); +extern "C" int32_t qwen36_set_adapter_lora_tensor( + void*, int64_t, int64_t, const char*, int32_t, void*); +extern "C" void qwen36_free_training_context(void*); + +static at::Tensor cuda_rand(std::initializer_list shape) { + return at::randn(shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)); +} + +int main() { + c10::cuda::CUDAGuard guard(0); + at::manual_seed(7); + + constexpr int64_t hidden = 16; + constexpr int64_t vocab = 8; + constexpr int64_t experts = 2; + constexpr int64_t head_dim = 8; + constexpr int64_t intermediate = 8; + constexpr int64_t rank = 8; + + // One full-attention MoE layer. Shapes intentionally match the native + // weight order used by build_weight_ptrs/kernel.cpp. + std::vector weights; + weights.push_back(cuda_rand({hidden})); // input RMSNorm + weights.push_back(cuda_rand({hidden})); // post-attention RMSNorm + weights.push_back(cuda_rand({2 * head_dim, hidden})); // q_proj + weights.push_back(cuda_rand({head_dim})); // q_norm + weights.push_back(cuda_rand({head_dim, hidden})); // k_proj + weights.push_back(cuda_rand({head_dim})); // k_norm + weights.push_back(cuda_rand({head_dim, hidden})); // v_proj + weights.push_back(cuda_rand({hidden, head_dim})); // o_proj + weights.push_back(cuda_rand({experts, hidden})); // router + weights.push_back(cuda_rand({1, hidden})); // shared expert gate + weights.push_back(cuda_rand({intermediate, hidden})); // shared gate + weights.push_back(cuda_rand({intermediate, hidden})); // shared up + weights.push_back(cuda_rand({hidden, intermediate})); // shared down + weights.push_back(cuda_rand({experts, 2 * intermediate, hidden})); + weights.push_back(cuda_rand({experts, hidden, intermediate})); + for (auto& weight : weights) weight.set_requires_grad(false); + + auto embed = cuda_rand({vocab, hidden}); + auto final_norm = cuda_rand({hidden}); + auto lm_head = cuda_rand({vocab, hidden}); + embed.set_requires_grad(false); + final_norm.set_requires_grad(false); + lm_head.set_requires_grad(false); + + std::vector weight_ptrs; + weight_ptrs.reserve(weights.size()); + for (auto& weight : weights) weight_ptrs.push_back(&weight); + LayerConfig config{}; + config.layer_type = 0; + config.num_heads = 1; + config.num_kv_heads = 1; + config.head_dim = head_dim; + config.num_k_heads = 0; + config.key_dim = 0; + config.num_v_heads = 0; + config.val_dim = 0; + config.conv_kernel = 0; + config.partial_rotary_factor = 1.0; + config.rope_theta = 10000.0; + config.rms_eps = 1e-5; + config.num_experts = experts; + config.top_k = 1; + config.moe_intermediate = intermediate; + config.expert_start = 0; + config.expert_count = experts; + config.intermediate_size = 0; + config.norm_topk_prob = 1; + config.nccl_comm = nullptr; + config.nccl_stream = nullptr; + + const int64_t target_layer = 0; + void* ctx = qwen36_create_training_context( + weight_ptrs.data(), static_cast(weight_ptrs.size()), + &embed, &final_norm, &lm_head, &config, 1, + static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank, + &target_layer, 1, "experts_gate_up_proj,experts_down_proj"); + if (!ctx) return 2; + + const int64_t count = qwen36_get_lora_count(ctx); + assert(count == 9); + auto* expert_a = reinterpret_cast(qwen36_get_lora_a(ctx, 7)); + auto* expert_b = reinterpret_cast(qwen36_get_lora_b(ctx, 7)); + auto* down_a = reinterpret_cast(qwen36_get_lora_a(ctx, 8)); + auto* down_b = reinterpret_cast(qwen36_get_lora_b(ctx, 8)); + assert(expert_a && expert_b && down_a && down_b); + assert(expert_a->sizes() == at::IntArrayRef({experts, rank, hidden})); + assert(expert_b->sizes() == at::IntArrayRef({experts, 2 * intermediate, rank})); + assert(down_a->sizes() == at::IntArrayRef({experts, rank, intermediate})); + assert(down_b->sizes() == at::IntArrayRef({experts, hidden, rank})); + + // Make both B tensors nonzero so the step exercises the LoRA branches. + auto expert_b_value = at::ones(expert_b->sizes(), expert_b->options()); + auto down_b_value = at::ones(down_b->sizes(), down_b->options()); + assert(qwen36_set_lora_tensor(ctx, 7, 1, &expert_b_value) == 0); + assert(qwen36_set_lora_tensor(ctx, 8, 1, &down_b_value) == 0); + auto expert_a_before = expert_a->clone(); + auto expert_b_before = expert_b->clone(); + + auto input_ids = at::arange(1, 3, at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 2}); + auto target_mask = at::ones({1, 2}, at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); + auto attention_mask = at::ones({1, 2}, at::TensorOptions().device(at::kCUDA).dtype(at::kBool)); + // Distinct tenant rows exercise the production [n_total, seq] path. The + // rows are intentionally different so a repeated batch-1 implementation + // cannot satisfy the assertions below. + auto multi_input_ids = at::tensor({1, 2, 3, 4}, + at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({2, 2}); + auto multi_target_mask = at::ones({2, 2}, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); + auto multi_attention_mask = at::ones({2, 2}, + at::TensorOptions().device(at::kCUDA).dtype(at::kBool)); + + // The optimized and fallback routed-expert paths must agree before the + // optimizer changes any parameters. On older libtorch both calls use the + // fallback, preserving the same compatibility check. + setenv("QWEN36_DISABLE_GROUPED_MM", "1", 1); + const double fallback_loss = + qwen36_eval_step(ctx, &input_ids, &target_mask, &attention_mask); + unsetenv("QWEN36_DISABLE_GROUPED_MM"); + setenv("QWEN36_REPORT_GROUPED_MM", "1", 1); + const double grouped_loss = + qwen36_eval_step(ctx, &input_ids, &target_mask, &attention_mask); + std::printf( + "native_qwen36_moe_lora_parity fallback=%0.8f grouped=%0.8f diff=%0.8e\n", + fallback_loss, grouped_loss, std::abs(fallback_loss - grouped_loss)); + assert(fallback_loss > 0.0 && grouped_loss > 0.0); + assert(std::abs(fallback_loss - grouped_loss) <= 2e-2); + + const double loss = qwen36_train_step(ctx, &input_ids, &target_mask, &attention_mask); + c10::cuda::device_synchronize(); + std::printf("native_qwen36_moe_lora_smoke loss=%0.8f\n", loss); + assert(loss == loss); + assert(loss > 0.0); + const double update_norm = (*expert_a - expert_a_before).abs().sum().item(); + const double b_update_norm = (*expert_b - expert_b_before).abs().sum().item(); + std::printf("native_qwen36_moe_lora_smoke expert_a_update=%0.8e expert_b_update=%0.8e\n", + update_norm, b_update_norm); + assert(update_norm > 0.0 || b_update_norm > 0.0); + + // Dynamic multi-adapter batches must apply shared-expert MLP LoRA per + // sample and preserve the parameter updates after chunk registry restore. + const char* shared_targets = + "shared_gate_proj,shared_up_proj,shared_down_proj," + "experts_gate_up_proj,experts_down_proj"; + const int64_t adapter_one = qwen36_add_lora( + ctx, rank, 1.0, &target_layer, 1, shared_targets); + const int64_t adapter_two = qwen36_add_lora( + ctx, rank, 1.0, &target_layer, 1, shared_targets); + assert(adapter_one > 0 && adapter_two > adapter_one); + auto* dynamic_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_one, 0, "shared_gate_proj", 1)); + assert(dynamic_b && dynamic_b->sizes() == at::IntArrayRef({intermediate, rank})); + auto dynamic_b_value = at::ones(dynamic_b->sizes(), dynamic_b->options()); + assert(qwen36_set_adapter_lora_tensor( + ctx, adapter_one, 0, "shared_gate_proj", 1, &dynamic_b_value) == 0); + auto dynamic_b_before = dynamic_b->clone(); + auto* dynamic_b_two = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_two, 0, "shared_gate_proj", 1)); + assert(dynamic_b_two && dynamic_b_two->sizes() == dynamic_b->sizes()); + auto dynamic_b_two_value = at::full(dynamic_b_two->sizes(), -1.0, dynamic_b_two->options()); + assert(qwen36_set_adapter_lora_tensor( + ctx, adapter_two, 0, "shared_gate_proj", 1, &dynamic_b_two_value) == 0); + auto dynamic_b_two_before = dynamic_b_two->clone(); + auto* dynamic_expert_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_one, 0, "experts_gate_up_proj", 1)); + assert(dynamic_expert_b && dynamic_expert_b->sizes() == + at::IntArrayRef({experts, 2 * intermediate, rank})); + auto dynamic_expert_b_value = at::ones( + dynamic_expert_b->sizes(), dynamic_expert_b->options()); + assert(qwen36_set_adapter_lora_tensor( + ctx, adapter_one, 0, "experts_gate_up_proj", 1, + &dynamic_expert_b_value) == 0); + auto dynamic_expert_b_before = dynamic_expert_b->clone(); + const double multi_loss = qwen36_train_multi_lora( + ctx, &multi_input_ids, &multi_target_mask, &multi_attention_mask, 2, rank); + c10::cuda::device_synchronize(); + dynamic_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_one, 0, "shared_gate_proj", 1)); + assert(dynamic_b); + dynamic_expert_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_one, 0, "experts_gate_up_proj", 1)); + assert(dynamic_expert_b); + const double dynamic_update = + (*dynamic_b - dynamic_b_before).abs().sum().item(); + const double dynamic_expert_update = + (*dynamic_expert_b - dynamic_expert_b_before).abs().sum().item(); + dynamic_b_two = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_two, 0, "shared_gate_proj", 1)); + assert(dynamic_b_two); + const double dynamic_two_update = + (*dynamic_b_two - dynamic_b_two_before).abs().sum().item(); + std::printf( + "native_qwen36_multi_lora_smoke loss=%0.8f shared_gate_b_update=%0.8e " + "expert_gate_up_b_update=%0.8e adapter_two_gate_b_update=%0.8e\n", + multi_loss, dynamic_update, dynamic_expert_update, dynamic_two_update); + assert(multi_loss == multi_loss && multi_loss > 0.0); + assert(dynamic_update > 0.0); + assert(dynamic_expert_update > 0.0); + assert(dynamic_two_update > 0.0); + qwen36_free_training_context(ctx); + + // Dense Qwen3.5 variants use the same per-sample activation path for + // gate/up/down projections. + std::vector dense_weights(weights.begin(), weights.begin() + 8); + dense_weights.push_back(cuda_rand({intermediate, hidden})); + dense_weights.push_back(cuda_rand({intermediate, hidden})); + dense_weights.push_back(cuda_rand({hidden, intermediate})); + weight_ptrs.clear(); + for (auto& weight : dense_weights) weight_ptrs.push_back(&weight); + LayerConfig dense_config = config; + dense_config.num_experts = 0; + dense_config.top_k = 0; + dense_config.moe_intermediate = 0; + dense_config.expert_count = 0; + dense_config.intermediate_size = intermediate; + ctx = qwen36_create_training_context( + weight_ptrs.data(), static_cast(weight_ptrs.size()), + &embed, &final_norm, &lm_head, &dense_config, 1, + static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank, + &target_layer, 1, "q_proj"); + assert(ctx); + const char* dense_targets = "gate_proj,up_proj,down_proj"; + const int64_t dense_one = qwen36_add_lora( + ctx, rank, 1.0, &target_layer, 1, dense_targets); + const int64_t dense_two = qwen36_add_lora( + ctx, rank, 1.0, &target_layer, 1, dense_targets); + assert(dense_one > 0 && dense_two > dense_one); + auto* dense_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor(ctx, dense_one, 0, "gate_proj", 1)); + assert(dense_b && dense_b->sizes() == at::IntArrayRef({intermediate, rank})); + auto dense_b_value = at::ones(dense_b->sizes(), dense_b->options()); + assert(qwen36_set_adapter_lora_tensor( + ctx, dense_one, 0, "gate_proj", 1, &dense_b_value) == 0); + auto dense_b_before = dense_b->clone(); + const double dense_loss = qwen36_train_multi_lora( + ctx, &multi_input_ids, &multi_target_mask, &multi_attention_mask, 2, rank); + c10::cuda::device_synchronize(); + dense_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor(ctx, dense_one, 0, "gate_proj", 1)); + assert(dense_b); + const double dense_update = + (*dense_b - dense_b_before).abs().sum().item(); + std::printf( + "native_qwen35_dense_multi_lora_smoke loss=%0.8f gate_b_update=%0.8e\n", + dense_loss, dense_update); + assert(dense_loss == dense_loss && dense_loss > 0.0); + assert(dense_update > 0.0); + qwen36_free_training_context(ctx); + + // Qwen3.5/3.6 linear-attention layers use the model's actual 128-wide + // delta-rule state. Exercise both fixed and per-sample dynamic GDN LoRA, + // including the custom CUDA backward for q/k/v/g/beta. + constexpr int64_t linear_heads = 1; + constexpr int64_t linear_dim = 128; + constexpr int64_t linear_qkv = 3 * linear_dim; + std::vector linear_weights; + linear_weights.push_back(cuda_rand({hidden})); + linear_weights.push_back(cuda_rand({hidden})); + linear_weights.push_back(cuda_rand({linear_qkv, hidden})); + linear_weights.push_back(cuda_rand({linear_dim, hidden})); + linear_weights.push_back(cuda_rand({linear_heads, hidden})); + linear_weights.push_back(cuda_rand({linear_heads, hidden})); + linear_weights.push_back(cuda_rand({linear_heads})); + linear_weights.push_back(cuda_rand({linear_heads})); + linear_weights.push_back(cuda_rand({linear_qkv, 1, 4})); + linear_weights.push_back(cuda_rand({linear_dim})); + linear_weights.push_back(cuda_rand({hidden, linear_dim})); + linear_weights.push_back(cuda_rand({intermediate, hidden})); + linear_weights.push_back(cuda_rand({intermediate, hidden})); + linear_weights.push_back(cuda_rand({hidden, intermediate})); + for (auto& weight : linear_weights) weight.set_requires_grad(false); + weight_ptrs.clear(); + for (auto& weight : linear_weights) weight_ptrs.push_back(&weight); + LayerConfig linear_config = dense_config; + linear_config.layer_type = 1; + linear_config.num_heads = 0; + linear_config.num_kv_heads = 0; + linear_config.head_dim = 0; + linear_config.num_k_heads = linear_heads; + linear_config.key_dim = linear_dim; + linear_config.num_v_heads = linear_heads; + linear_config.val_dim = linear_dim; + linear_config.conv_kernel = 4; + ctx = qwen36_create_training_context( + weight_ptrs.data(), static_cast(weight_ptrs.size()), + &embed, &final_norm, &lm_head, &linear_config, 1, + static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank, + &target_layer, 1, + "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj"); + assert(ctx); + assert(qwen36_get_lora_count(ctx) == 8); + auto* linear_a = reinterpret_cast(qwen36_get_lora_a(ctx, 0)); + auto* linear_b = reinterpret_cast(qwen36_get_lora_b(ctx, 0)); + assert(linear_a && linear_b); + assert(linear_a->sizes() == at::IntArrayRef({rank, hidden})); + assert(linear_b->sizes() == at::IntArrayRef({linear_qkv, rank})); + auto linear_b_value = at::ones(linear_b->sizes(), linear_b->options()); + assert(qwen36_set_lora_tensor(ctx, 0, 1, &linear_b_value) == 0); + auto linear_a_before = linear_a->clone(); + const double linear_loss = qwen36_train_step( + ctx, &input_ids, &target_mask, &attention_mask); + c10::cuda::device_synchronize(); + const double linear_update = + (*linear_a - linear_a_before).abs().sum().item(); + std::printf( + "native_qwen35_linear_lora_smoke loss=%0.8f qkv_a_update=%0.8e\n", + linear_loss, linear_update); + assert(linear_loss == linear_loss && linear_loss > 0.0); + assert(linear_update > 0.0); + + const int64_t linear_adapter_one = qwen36_add_lora( + ctx, rank, 1.0, &target_layer, 1, "in_proj_qkv"); + const int64_t linear_adapter_two = qwen36_add_lora( + ctx, rank, 1.0, &target_layer, 1, "in_proj_qkv"); + assert(linear_adapter_one > 0 && linear_adapter_two > linear_adapter_one); + auto* dynamic_linear_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, linear_adapter_one, 0, "in_proj_qkv", 1)); + assert(dynamic_linear_b && dynamic_linear_b->sizes() == + at::IntArrayRef({linear_qkv, rank})); + auto dynamic_linear_b_value = at::ones( + dynamic_linear_b->sizes(), dynamic_linear_b->options()); + assert(qwen36_set_adapter_lora_tensor( + ctx, linear_adapter_one, 0, "in_proj_qkv", 1, + &dynamic_linear_b_value) == 0); + auto dynamic_linear_b_before = dynamic_linear_b->clone(); + const double dynamic_linear_loss = qwen36_train_multi_lora( + ctx, &multi_input_ids, &multi_target_mask, &multi_attention_mask, 2, rank); + c10::cuda::device_synchronize(); + dynamic_linear_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, linear_adapter_one, 0, "in_proj_qkv", 1)); + assert(dynamic_linear_b); + const double dynamic_linear_update = + (*dynamic_linear_b - dynamic_linear_b_before).abs().sum().item(); + std::printf( + "native_qwen35_linear_multi_lora_smoke loss=%0.8f qkv_b_update=%0.8e\n", + dynamic_linear_loss, dynamic_linear_update); + assert(dynamic_linear_loss == dynamic_linear_loss && dynamic_linear_loss > 0.0); + assert(dynamic_linear_update > 0.0); + qwen36_free_training_context(ctx); + return 0; +} diff --git a/crates/rustrain-server/Cargo.toml b/crates/rustrain-server/Cargo.toml index cf36ea5a..3b5f52d1 100644 --- a/crates/rustrain-server/Cargo.toml +++ b/crates/rustrain-server/Cargo.toml @@ -30,3 +30,6 @@ tracing.workspace = true [build-dependencies] tonic-build.workspace = true + +[dev-dependencies] +tempfile.workspace = true diff --git a/crates/rustrain-server/proto/train.proto b/crates/rustrain-server/proto/train.proto index 644f72c1..d7f413b4 100644 --- a/crates/rustrain-server/proto/train.proto +++ b/crates/rustrain-server/proto/train.proto @@ -10,6 +10,7 @@ service TrainService { rpc SaveCheckpoint(SaveCheckpointRequest) returns (CheckpointInfo); rpc LoadCheckpoint(LoadCheckpointRequest) returns (CheckpointInfo); rpc ExportAdapter(ExportAdapterRequest) returns (AdapterInfo); + rpc ImportAdapter(ImportAdapterRequest) returns (ImportAdapterResponse); rpc StreamMetrics(StreamMetricsRequest) returns (stream StepMetric); rpc GetStatus(GetStatusRequest) returns (SessionStatus); rpc AddLoRA(AddLoRARequest) returns (AddLoRAResponse); @@ -94,12 +95,21 @@ message CheckpointInfo { message ExportAdapterRequest { string session_id = 1; string path = 2; + int64 adapter_id = 3; // 0 exports the fixed adapter } message AdapterInfo { string path = 1; int64 param_count = 2; } +message ImportAdapterRequest { + string session_id = 1; + string path = 2; +} +message ImportAdapterResponse { + int64 adapter_id = 1; +} + message StreamMetricsRequest { string session_id = 1; } diff --git a/crates/rustrain-server/src/api.rs b/crates/rustrain-server/src/api.rs index 3ed5b008..4896f450 100644 --- a/crates/rustrain-server/src/api.rs +++ b/crates/rustrain-server/src/api.rs @@ -1,11 +1,11 @@ //! HTTP API (axum) — RESTful endpoints for training session management. use axum::{ + Json, Router, extract::{Path, State}, http::StatusCode, response::sse::{Event, KeepAlive, Sse}, routing::{get, post}, - Json, Router, }; use futures::stream::{self, Stream}; use serde::{Deserialize, Serialize}; @@ -39,6 +39,7 @@ pub fn router(state: Arc) -> Router { .route("/v1/sessions/{id}/save_checkpoint", post(save_checkpoint)) .route("/v1/sessions/{id}/load_checkpoint", post(load_checkpoint)) .route("/v1/sessions/{id}/export_adapter", post(export_adapter)) + .route("/v1/sessions/{id}/import_adapter", post(import_adapter)) .route("/v1/sessions/{id}/add_lora", post(add_lora)) .route("/v1/sessions/{id}/remove_lora", post(remove_lora)) .route("/v1/sessions/{id}/list_lora", get(list_lora)) @@ -55,7 +56,9 @@ struct ErrorResponse { fn err_resp(msg: &str) -> (StatusCode, Json) { ( StatusCode::BAD_REQUEST, - Json(ErrorResponse { error: msg.to_string() }), + Json(ErrorResponse { + error: msg.to_string(), + }), ) } @@ -76,7 +79,11 @@ async fn delete_session( State(state): State>, axum::extract::Path(id): axum::extract::Path, ) -> Result, (StatusCode, Json)> { - state.manager.delete_session(&id).await.map_err(|e| err_resp(&e))?; + state + .manager + .delete_session(&id) + .await + .map_err(|e| err_resp(&e))?; Ok(Json(serde_json::json!({"deleted": id}))) } @@ -89,9 +96,7 @@ struct CreateSessionResponse { session_id: String, } -async fn list_sessions( - State(state): State>, -) -> Json> { +async fn list_sessions(State(state): State>) -> Json> { Json(state.manager.list_sessions().await) } @@ -181,7 +186,7 @@ async fn init_lora( #[derive(Deserialize)] struct InitLoRAHttp { rank: i64, - alpha: i64, + alpha: f64, target_layers: Vec, target_modules: Vec, lr: f64, @@ -204,12 +209,9 @@ async fn train_step( .get_session(&id) .await .ok_or_else(|| err_resp("session not found"))?; - let input_ids = decode_tensor(&req.input_ids) - .map_err(|e| err_resp(&e))?; - let target_mask = decode_tensor(&req.target_mask) - .map_err(|e| err_resp(&e))?; - let attention_mask = decode_tensor(&req.attention_mask) - .map_err(|e| err_resp(&e))?; + let input_ids = decode_tensor(&req.input_ids).map_err(|e| err_resp(&e))?; + let target_mask = decode_tensor(&req.target_mask).map_err(|e| err_resp(&e))?; + let attention_mask = decode_tensor(&req.attention_mask).map_err(|e| err_resp(&e))?; let mut s = session.lock().await; let result = s @@ -254,18 +256,73 @@ struct TrainStepResponse { /// Decode a base64-encoded int64 tensor to Vec (for EP IPC). fn decode_int64_vec(t: &TensorHttp) -> Result, String> { - use base64::{engine::general_purpose, Engine}; + use base64::{Engine, engine::general_purpose}; let bytes = general_purpose::STANDARD .decode(&t.data) .map_err(|e| format!("base64 decode: {e}"))?; - Ok(bytes + let values = bytes .chunks_exact(8) .map(|c| i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]])) - .collect()) + .collect::>(); + let expected = t + .shape + .iter() + .try_fold(1usize, |acc, dim| { + usize::try_from(*dim) + .ok() + .and_then(|dim| acc.checked_mul(dim)) + }) + .ok_or_else(|| format!("invalid tensor shape {:?}", t.shape))?; + if expected != values.len() { + return Err(format!( + "tensor shape {:?} expects {} int64 values, got {}", + t.shape, + expected, + values.len() + )); + } + Ok(values) +} + +fn validate_multi_lora_http_shapes( + input_ids: &TensorHttp, + target_mask: &TensorHttp, + attention_mask: &TensorHttp, + n_total: i32, +) -> Result { + if n_total <= 0 { + return Err(format!("n_total must be positive, got {n_total}")); + } + if input_ids.shape.len() != 1 && input_ids.shape.len() != 2 { + return Err(format!("input_ids must have shape [seq] or [batch, seq], got {:?}", input_ids.shape)); + } + if target_mask.shape != input_ids.shape || attention_mask.shape != input_ids.shape { + return Err(format!( + "multi-LoRA masks must have the same shape as input_ids: input={:?} target={:?} attention={:?}", + input_ids.shape, target_mask.shape, attention_mask.shape + )); + } + let seq_len = *input_ids + .shape + .last() + .ok_or_else(|| "input_ids shape is empty".to_string())?; + if seq_len <= 0 { + return Err(format!("sequence length must be positive, got {seq_len}")); + } + if input_ids.shape.len() == 2 { + let batch = input_ids.shape[0]; + if batch != 1 && batch != i64::from(n_total) { + return Err(format!( + "multi-LoRA batch must be 1 or n_total={}, got {}", + n_total, batch + )); + } + } + usize::try_from(seq_len).map_err(|_| format!("invalid sequence length {seq_len}")) } fn decode_tensor(t: &TensorHttp) -> Result { - use base64::{engine::general_purpose, Engine}; + use base64::{Engine, engine::general_purpose}; let bytes = general_purpose::STANDARD .decode(&t.data) .map_err(|e| format!("base64 decode: {e}"))?; @@ -293,8 +350,13 @@ fn decode_tensor(t: &TensorHttp) -> Result { } _ => return Err("only int64 and float32 supported via HTTP".into()), }; - let local_rank = std::env::var("LOCAL_RANK").ok().and_then(|s| s.parse::().ok()).unwrap_or(0); - Ok(tensor.reshape(&t.shape).to_device(tch::Device::Cuda(local_rank))) + let local_rank = std::env::var("LOCAL_RANK") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + Ok(tensor + .reshape(&t.shape) + .to_device(tch::Device::Cuda(local_rank))) } async fn eval_step( @@ -381,6 +443,8 @@ async fn load_checkpoint( #[derive(Deserialize)] struct ExportHttp { path: String, + #[serde(default)] + adapter_id: Option, } #[derive(Serialize)] struct ExportResponse { @@ -400,7 +464,7 @@ async fn export_adapter( .ok_or_else(|| err_resp("session not found"))?; let s = session.lock().await; let count = s - .export_adapter(&req.path) + .export_adapter(&req.path, req.adapter_id) .map_err(|e| err_resp(&e.to_string()))?; Ok(Json(ExportResponse { path: req.path, @@ -408,6 +472,33 @@ async fn export_adapter( })) } +#[derive(Deserialize)] +struct ImportHttp { + path: String, +} + +#[derive(Serialize)] +struct ImportResponse { + adapter_id: i64, +} + +async fn import_adapter( + State(state): State>, + Path(id): Path, + Json(req): Json, +) -> Result, (StatusCode, Json)> { + let session = state + .manager + .get_session(&id) + .await + .ok_or_else(|| err_resp("session not found"))?; + let mut s = session.lock().await; + let adapter_id = s + .import_adapter(&req.path) + .map_err(|e| err_resp(&e.to_string()))?; + Ok(Json(ImportResponse { adapter_id })) +} + async fn get_status( State(state): State>, Path(id): Path, @@ -522,13 +613,15 @@ async fn stream_metrics( drop(s); let stream = stream::iter(metrics.into_iter().map(|m| { - Ok(Event::default().json_data(StepMetricJson { - step: m.step, - loss: m.loss, - lr: m.lr, - mem_gb: m.mem_gb, - timestamp_unix: m.timestamp_unix, - }).unwrap_or_default()) + Ok(Event::default() + .json_data(StepMetricJson { + step: m.step, + loss: m.loss, + lr: m.lr, + mem_gb: m.mem_gb, + timestamp_unix: m.timestamp_unix, + }) + .unwrap_or_default()) })); Ok(Sse::new(stream).keep_alive(KeepAlive::default())) @@ -550,7 +643,10 @@ struct StepMetricJson { pub fn ep_router(state: Arc) -> Router { Router::new() .route("/v1/sessions", post(ep_create_session)) - .route("/v1/sessions/{id}", axum::routing::delete(ep_delete_session)) + .route( + "/v1/sessions/{id}", + axum::routing::delete(ep_delete_session), + ) .route("/v1/sessions/{id}/load_model", post(ep_load_model)) .route("/v1/sessions/{id}/load_dataset", post(ep_load_dataset)) .route("/v1/sessions/{id}/init_lora", post(ep_init_lora)) @@ -575,9 +671,13 @@ async fn ep_create_session( State(state): State>, Json(req): Json, ) -> Result, (StatusCode, Json)> { - let cmd = rustrain_ipc::EpCommand::CreateSession { session_id: req.session_id.clone() }; + let cmd = rustrain_ipc::EpCommand::CreateSession { + session_id: req.session_id.clone(), + }; match state.coordinator.dispatch(&cmd) { - rustrain_ipc::EpResult::Ok => Ok(Json(CreateSessionResponse { session_id: req.session_id })), + rustrain_ipc::EpResult::Ok => Ok(Json(CreateSessionResponse { + session_id: req.session_id, + })), rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)), _ => Err(err_resp("unexpected result")), } @@ -587,7 +687,9 @@ async fn ep_delete_session( State(state): State>, Path(id): Path, ) -> Result, (StatusCode, Json)> { - let cmd = rustrain_ipc::EpCommand::DeleteSession { session_id: id.clone() }; + let cmd = rustrain_ipc::EpCommand::DeleteSession { + session_id: id.clone(), + }; match state.coordinator.dispatch(&cmd) { rustrain_ipc::EpResult::Ok => Ok(Json(serde_json::json!({"deleted": id}))), rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)), @@ -708,7 +810,13 @@ async fn ep_train_multi_lora( let input_ids = decode_int64_vec(&req.input_ids).map_err(|e| err_resp(&e))?; let target_mask = decode_int64_vec(&req.target_mask).map_err(|e| err_resp(&e))?; let attention_mask = decode_int64_vec(&req.attention_mask).map_err(|e| err_resp(&e))?; - let seq_len = input_ids.len(); + let seq_len = validate_multi_lora_http_shapes( + &req.input_ids, + &req.target_mask, + &req.attention_mask, + req.n_total, + ) + .map_err(|e| err_resp(&e))?; let cmd = rustrain_ipc::EpCommand::TrainMultiLora { session_id: id, @@ -779,7 +887,10 @@ async fn ep_remove_lora( Path(id): Path, Json(req): Json, ) -> Result, (StatusCode, Json)> { - let cmd = rustrain_ipc::EpCommand::RemoveLora { session_id: id, adapter_id: req.adapter_id }; + let cmd = rustrain_ipc::EpCommand::RemoveLora { + session_id: id, + adapter_id: req.adapter_id, + }; match state.coordinator.dispatch(&cmd) { rustrain_ipc::EpResult::Ok => Ok(Json(serde_json::json!({"removed": true}))), rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)), @@ -804,7 +915,11 @@ async fn ep_export_adapter( Path(id): Path, Json(req): Json, ) -> Result, (StatusCode, Json)> { - let cmd = rustrain_ipc::EpCommand::ExportAdapter { session_id: id, path: req.path }; + let cmd = rustrain_ipc::EpCommand::ExportAdapter { + session_id: id, + path: req.path, + adapter_id: req.adapter_id, + }; match state.coordinator.dispatch(&cmd) { rustrain_ipc::EpResult::Count(n) => Ok(Json(serde_json::json!({"exported": n}))), rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)), @@ -818,9 +933,17 @@ async fn ep_get_status( ) -> Result, (StatusCode, Json)> { let cmd = rustrain_ipc::EpCommand::Status { session_id: id }; match state.coordinator.dispatch(&cmd) { - rustrain_ipc::EpResult::Status { state, step, last_loss, model_path } => { - Ok(Json(StatusResponse { state, step, last_loss, model_path })) - } + rustrain_ipc::EpResult::Status { + state, + step, + last_loss, + model_path, + } => Ok(Json(StatusResponse { + state, + step, + last_loss, + model_path, + })), rustrain_ipc::EpResult::Error(e) => Err(err_resp(&e)), _ => Err(err_resp("unexpected result")), } diff --git a/crates/rustrain-server/src/checkpoint.rs b/crates/rustrain-server/src/checkpoint.rs index c3164fa1..18dd40d1 100644 --- a/crates/rustrain-server/src/checkpoint.rs +++ b/crates/rustrain-server/src/checkpoint.rs @@ -12,7 +12,7 @@ pub struct CheckpointManifest { pub loss: f64, pub model_path: String, pub lora_rank: i64, - pub lora_alpha: i64, + pub lora_alpha: f64, pub files: Vec, } @@ -32,7 +32,7 @@ pub fn save_checkpoint( loss: f64, model_path: &str, lora_rank: i64, - lora_alpha: i64, + lora_alpha: f64, lora_a: &[Tensor], lora_b: &[Tensor], adam_m: &[Tensor], @@ -64,7 +64,8 @@ pub fn save_checkpoint( .with_context(|| "write manifest.json")?; tracing::info!( - step, loss, + step, + loss, path = dir.display().to_string(), "checkpoint saved" ); @@ -102,102 +103,89 @@ pub fn load_checkpoint(dir: &Path) -> Result { } fn save_tensors(path: &Path, a: &[Tensor], b: &[Tensor]) -> Result<()> { - use std::io::Write; let mut named: Vec<(String, Tensor)> = Vec::new(); for (i, t) in a.iter().enumerate() { - named.push((format!("a_{i}"), t.to_kind(tch::Kind::Float).to_device(tch::Device::Cpu))); + named.push(( + format!("a_{i}"), + t.to_kind(tch::Kind::Float).to_device(tch::Device::Cpu), + )); } for (i, t) in b.iter().enumerate() { - named.push((format!("b_{i}"), t.to_kind(tch::Kind::Float).to_device(tch::Device::Cpu))); + named.push(( + format!("b_{i}"), + t.to_kind(tch::Kind::Float).to_device(tch::Device::Cpu), + )); } - // Build safetensors manually (header + data) - let mut header = serde_json::Map::new(); - let mut offset = 0u64; - let mut all_bytes: Vec = Vec::new(); - for (name, t) in &named { - let t = t.contiguous().to_kind(tch::Kind::Float); - let shape: Vec = t.size().iter().copied().collect(); - let data: Vec = Vec::::try_from(&t.reshape([-1]))?; - let bytes: Vec = data.iter().flat_map(|f| f.to_le_bytes()).collect(); - header.insert( - name.clone(), - serde_json::json!({"dtype":"F32","shape":shape,"data_offsets":[offset, offset + bytes.len() as u64]}), - ); - offset += bytes.len() as u64; - all_bytes.extend_from_slice(&bytes); - } - let header_str = serde_json::to_string(&serde_json::Value::Object(header))?; - let file = std::fs::File::create(path).with_context(|| format!("create {}", path.display()))?; - let mut writer = std::io::BufWriter::new(file); - writer.write_all(&(header_str.len() as u64).to_le_bytes())?; - writer.write_all(header_str.as_bytes())?; - writer.write_all(&all_bytes)?; + Tensor::write_safetensors(&named, path).with_context(|| format!("write {}", path.display()))?; Ok(()) } fn load_tensors(path: &Path) -> Result<(Vec, Vec)> { - let data = std::fs::read(path).with_context(|| format!("read {}", path.display()))?; - if data.len() < 8 { - anyhow::bail!("safetensors file too small"); + let named = + Tensor::read_safetensors(path).with_context(|| format!("read {}", path.display()))?; + let mut by_name = std::collections::BTreeMap::new(); + for (name, tensor) in named { + by_name.insert(name, tensor); } - let header_len = u64::from_le_bytes(data[..8].try_into().unwrap()) as usize; - let header_str = std::str::from_utf8(&data[8..8 + header_len])?; - let header: serde_json::Value = serde_json::from_str(header_str)?; - - // Count a_N and b_N entries - let mut a_count = 0; - let mut b_count = 0; - if let serde_json::Value::Object(map) = &header { - for key in map.keys() { - if key.starts_with("a_") { - a_count += 1; - } else if key.starts_with("b_") { - b_count += 1; + fn collect_side( + tensors: &std::collections::BTreeMap, + prefix: &str, + ) -> Result> { + let mut indices = tensors + .keys() + .filter_map(|name| name.strip_prefix(prefix)?.parse::().ok()) + .collect::>(); + indices.sort_unstable(); + let mut result = Vec::with_capacity(indices.len()); + for (expected, index) in indices.into_iter().enumerate() { + if index != expected { + anyhow::bail!("checkpoint tensor indices for {prefix} are not contiguous"); } + result.push( + tensors + .get(&format!("{prefix}{index}")) + .expect("index collected from map") + .shallow_clone(), + ); } + Ok(result) } - - let mut a_tensors = Vec::new(); - let mut b_tensors = Vec::new(); - - for i in 0..a_count { - let key = format!("a_{i}"); - let (tensor, shape) = load_one_tensor(&header, &key, &data, 8 + header_len)?; - a_tensors.push(tensor.reshape(&shape)); - } - for i in 0..b_count { - let key = format!("b_{i}"); - let (tensor, shape) = load_one_tensor(&header, &key, &data, 8 + header_len)?; - b_tensors.push(tensor.reshape(&shape)); - } - - Ok((a_tensors, b_tensors)) + Ok((collect_side(&by_name, "a_")?, collect_side(&by_name, "b_")?)) } -fn load_one_tensor( - header: &serde_json::Value, - key: &str, - data: &[u8], - data_offset: usize, -) -> Result<(Tensor, Vec)> { - let entry = header - .get(key) - .ok_or_else(|| anyhow::anyhow!("key {key} not found in safetensors"))?; - let shape: Vec = entry["shape"] - .as_array() - .unwrap() - .iter() - .map(|v| v.as_i64().unwrap()) - .collect(); - let offsets = entry["data_offsets"].as_array().unwrap(); - let start = data_offset + offsets[0].as_u64().unwrap() as usize; - let end = data_offset + offsets[1].as_u64().unwrap() as usize; - let bytes = &data[start..end]; - let floats: Vec = bytes - .chunks_exact(4) - .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) - .collect(); - let tensor = Tensor::from_slice(&floats); - Ok((tensor, shape)) +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn checkpoint_adapter_and_optimizer_roundtrip() { + let dir = tempfile::tempdir().unwrap(); + let a = Tensor::arange(6, (tch::Kind::Float, tch::Device::Cpu)).reshape([2, 3]); + let b = Tensor::arange(8, (tch::Kind::Float, tch::Device::Cpu)).reshape([4, 2]); + let m = Tensor::ones([2, 3], (tch::Kind::Float, tch::Device::Cpu)); + let v = Tensor::full([4, 2], 2.0, (tch::Kind::Float, tch::Device::Cpu)); + save_checkpoint( + dir.path(), + 7, + 1.25, + "Qwen/test", + 2, + 4.0, + &[a.shallow_clone()], + &[b.shallow_clone()], + &[m.shallow_clone()], + &[v.shallow_clone()], + ) + .unwrap(); + let loaded = load_checkpoint(dir.path()).unwrap(); + assert_eq!(loaded.manifest.lora_alpha, 4.0); + assert_eq!(loaded.manifest.step, 7); + assert_eq!(loaded.lora_a[0].size(), [2, 3]); + assert_eq!(loaded.lora_b[0].size(), [4, 2]); + assert!(loaded.lora_a[0].allclose(&a, 1e-6, 1e-6, false)); + assert!(loaded.lora_b[0].allclose(&b, 1e-6, 1e-6, false)); + assert!(loaded.adam_m[0].allclose(&m, 1e-6, 1e-6, false)); + assert!(loaded.adam_v[0].allclose(&v, 1e-6, 1e-6, false)); + } } diff --git a/crates/rustrain-server/src/ep.rs b/crates/rustrain-server/src/ep.rs index e84a5f35..3b5bf7a7 100644 --- a/crates/rustrain-server/src/ep.rs +++ b/crates/rustrain-server/src/ep.rs @@ -55,6 +55,7 @@ impl EpCoordinator { .env("RANK", rank.to_string()) .env("WORLD_SIZE", world_size.to_string()) .env("LOCAL_RANK", rank.to_string()) + .env("RUSTRAIN_NCCL_RUN_ID", &shm_name) .env("QWEN36_LOSS_DIAG", std::env::var("QWEN36_LOSS_DIAG").unwrap_or_default()) .env("QWEN36_GROUP_SIZE", std::env::var("QWEN36_GROUP_SIZE").unwrap_or_default()) .env("QWEN36_FUSED_CE", std::env::var("QWEN36_FUSED_CE").unwrap_or_default()) @@ -263,9 +264,30 @@ fn execute_command(session: &mut Qwen36Session, cmd: &EpCommand) -> EpResult { } EpCommand::TrainMultiLora { input_ids, target_mask, attention_mask, seq_len, n_total, lora_rank, .. } => { let sl = *seq_len as i64; - let input_ids_tensor = tch::Tensor::from_slice(input_ids).reshape(&[1, sl]).to_device(session.device()); - let target_mask_tensor = tch::Tensor::from_slice(target_mask).reshape(&[1, sl]).to_device(session.device()); - let attention_mask_tensor = tch::Tensor::from_slice(attention_mask).reshape(&[1, sl]).to_device(session.device()); + let batch = if *n_total > 0 + && input_ids.len() == (*n_total as usize).saturating_mul(*seq_len) + { + *n_total as i64 + } else if input_ids.len() == *seq_len { + 1 + } else { + return EpResult::Error(format!( + "multi-LoRA input length {} is incompatible with n_total={} seq_len={}", + input_ids.len(), n_total, seq_len + )); + }; + let expected = (batch as usize).saturating_mul(*seq_len); + if target_mask.len() != expected || attention_mask.len() != expected { + return EpResult::Error(format!( + "multi-LoRA mask lengths must equal {}, got target={} attention={}", + expected, + target_mask.len(), + attention_mask.len() + )); + } + let input_ids_tensor = tch::Tensor::from_slice(input_ids).reshape(&[batch, sl]).to_device(session.device()); + let target_mask_tensor = tch::Tensor::from_slice(target_mask).reshape(&[batch, sl]).to_device(session.device()); + let attention_mask_tensor = tch::Tensor::from_slice(attention_mask).reshape(&[batch, sl]).to_device(session.device()); match session.train_multi_lora(TrainInput { input_ids: input_ids_tensor, @@ -291,8 +313,10 @@ fn execute_command(session: &mut Qwen36Session, cmd: &EpCommand) -> EpResult { Err(e) => EpResult::Error(e.to_string()), } } - EpCommand::ExportAdapter { path, .. } => { - match session.export_adapter(path) { + EpCommand::ExportAdapter { + path, adapter_id, .. + } => { + match session.export_adapter(path, *adapter_id) { Ok(n) => EpResult::Count(n), Err(e) => EpResult::Error(e.to_string()), } diff --git a/crates/rustrain-server/src/grpc.rs b/crates/rustrain-server/src/grpc.rs index 1d3f20bc..3d788c72 100644 --- a/crates/rustrain-server/src/grpc.rs +++ b/crates/rustrain-server/src/grpc.rs @@ -75,7 +75,7 @@ impl TrainService for TrainServiceImpl { let count = s .init_lora(InitLoRARequest { rank: req.rank, - alpha: req.alpha as i64, + alpha: req.alpha, target_layers: req.target_layers.iter().map(|&l| l as usize).collect(), target_modules: req.target_modules, lr: req.lr, @@ -141,9 +141,7 @@ impl TrainService for TrainServiceImpl { attention_mask, }) .map_err(|e| Status::internal(e.to_string()))?; - Ok(Response::new(EvalStepResponse { - loss: result.loss, - })) + Ok(Response::new(EvalStepResponse { loss: result.loss })) } async fn save_checkpoint( @@ -200,7 +198,7 @@ impl TrainService for TrainServiceImpl { .ok_or_else(|| Status::not_found("session not found"))?; let s = session.lock().await; let count = s - .export_adapter(&req.path) + .export_adapter(&req.path, Some(req.adapter_id)) .map_err(|e| Status::internal(e.to_string()))?; Ok(Response::new(AdapterInfo { path: req.path, @@ -208,7 +206,25 @@ impl TrainService for TrainServiceImpl { })) } - type StreamMetricsStream = std::pin::Pin> + Send>>; + async fn import_adapter( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let session = self + .manager + .get_session(&req.session_id) + .await + .ok_or_else(|| Status::not_found("session not found"))?; + let mut s = session.lock().await; + let adapter_id = s + .import_adapter(&req.path) + .map_err(|e| Status::internal(e.to_string()))?; + Ok(Response::new(ImportAdapterResponse { adapter_id })) + } + + type StreamMetricsStream = + std::pin::Pin> + Send>>; async fn stream_metrics( &self, @@ -321,7 +337,12 @@ fn decode_tensor_data(td: &Option) -> Result { "int64" => tch::Kind::Int64, "float32" => tch::Kind::Float, "bfloat16" => tch::Kind::BFloat16, - _ => return Err(Status::invalid_argument(format!("unsupported dtype: {}", td.dtype))), + _ => { + return Err(Status::invalid_argument(format!( + "unsupported dtype: {}", + td.dtype + ))); + } }; let tensor = match kind { tch::Kind::Int64 => { @@ -342,6 +363,11 @@ fn decode_tensor_data(td: &Option) -> Result { } _ => return Err(Status::invalid_argument("only int64 and float32 supported")), }; - let local_rank = std::env::var("LOCAL_RANK").ok().and_then(|s| s.parse::().ok()).unwrap_or(0); - Ok(tensor.reshape(&td.shape).to_device(tch::Device::Cuda(local_rank))) + let local_rank = std::env::var("LOCAL_RANK") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + Ok(tensor + .reshape(&td.shape) + .to_device(tch::Device::Cuda(local_rank))) } diff --git a/crates/rustrain-server/src/session.rs b/crates/rustrain-server/src/session.rs index e115c794..d057293b 100644 --- a/crates/rustrain-server/src/session.rs +++ b/crates/rustrain-server/src/session.rs @@ -1,6 +1,6 @@ //! Training session trait + Qwen3.6 implementation. -use anyhow::{anyhow, Context, Result}; +use anyhow::{anyhow, bail, Context, Result}; use std::path::PathBuf; use std::sync::Arc; use tch::{Device, Kind, Tensor}; @@ -8,6 +8,7 @@ use tokio::sync::Mutex; use crate::checkpoint; use crate::metrics::{FileMetricsSink, MetricsSink, StepMetric}; +use rustrain_qwen3_6::lora::{Qwen36AdapterArtifact, Qwen36LoraConfig, Qwen36LoraTargetModule}; /// Session states. #[derive(Debug, Clone)] @@ -36,7 +37,7 @@ pub struct SessLoadDatasetRequest { #[derive(Debug)] pub struct InitLoRARequest { pub rank: i64, - pub alpha: i64, + pub alpha: f64, pub target_layers: Vec, pub target_modules: Vec, pub lr: f64, @@ -77,11 +78,19 @@ pub trait TrainingSession: Send { fn load_dataset(&mut self, req: SessLoadDatasetRequest) -> Result; fn init_lora(&mut self, req: InitLoRARequest) -> Result; fn train_step(&mut self, input: TrainInput) -> Result; - fn train_multi_lora(&mut self, input: TrainInput, n_total: i32, rank: i32) -> Result; + fn train_multi_lora( + &mut self, + input: TrainInput, + n_total: i32, + rank: i32, + ) -> Result; fn eval_step(&self, input: TrainInput) -> Result; fn save_checkpoint(&self, path: &str) -> Result<(u64, f64)>; fn load_checkpoint(&mut self, path: &str) -> Result<(u64, f64)>; - fn export_adapter(&self, path: &str) -> Result; + fn export_adapter(&self, path: &str, adapter_id: Option) -> Result; + /// Import one PEFT-style adapter as a new dynamic adapter. The fixed + /// adapter remains reserved as ID 0 and is never overwritten by import. + fn import_adapter(&mut self, path: &str) -> Result; fn status(&self) -> SessionStatus; fn get_metrics(&self) -> Vec; @@ -101,7 +110,7 @@ pub struct AddLoRARequest { pub rank: i64, pub alpha: f64, pub target_layers: Vec, - pub target_modules: String, // comma-separated, empty = all + pub target_modules: String, // comma-separated, empty = all } /// Qwen3.6 training session — wraps CppTrainingContext. @@ -116,7 +125,10 @@ pub struct Qwen36Session { weights: Option>, dataset: Option, lora_rank: i64, - lora_alpha: i64, + lora_alpha: f64, + lora_target_layers: Vec, + lora_target_modules: Vec, + dynamic_lora_configs: std::collections::BTreeMap, lr: f64, metrics: Option>, last_loss: f64, @@ -142,7 +154,10 @@ impl Qwen36Session { weights: None, dataset: None, lora_rank: 0, - lora_alpha: 0, + lora_alpha: 0.0, + lora_target_layers: Vec::new(), + lora_target_modules: Vec::new(), + dynamic_lora_configs: std::collections::BTreeMap::new(), lr: 1e-4, metrics: Some(Arc::new(FileMetricsSink::new(metrics_path))), last_loss: 0.0, @@ -200,10 +215,7 @@ impl TrainingSession for Qwen36Session { // Load runtime config from model's config.json (no need to parse full TOML) let model_path_obj = std::path::Path::new(model_path); - let runtime_config = - rustrain_qwen3_6::config::read_qwen36_runtime_config( - model_path_obj, - )?; + let runtime_config = rustrain_qwen3_6::config::read_qwen36_runtime_config(model_path_obj)?; // Build needed weight set — keys must match safetensors (with model prefix) let n_layers = runtime_config.num_hidden_layers; @@ -289,15 +301,28 @@ impl TrainingSession for Qwen36Session { // ── Expert Parallel support ── // Read EP params from env vars (set by launcher script) - let ep_rank = std::env::var("RANK").ok().and_then(|s| s.parse::().ok()).unwrap_or(0); - let ep_world_size = std::env::var("WORLD_SIZE").ok().and_then(|s| s.parse::().ok()).unwrap_or(1); - let local_rank = std::env::var("LOCAL_RANK").ok().and_then(|s| s.parse::().ok()).unwrap_or(0); + let ep_rank = std::env::var("RANK") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + let ep_world_size = std::env::var("WORLD_SIZE") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(1); + let local_rank = std::env::var("LOCAL_RANK") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); let is_ep = ep_world_size > 1 && runtime_config.is_moe; // Compute expert shard let (expert_start, expert_count) = if is_ep { - assert!(runtime_config.num_experts % ep_world_size == 0, - "num_experts {} not divisible by ep_world_size {}", runtime_config.num_experts, ep_world_size); + assert!( + runtime_config.num_experts % ep_world_size == 0, + "num_experts {} not divisible by ep_world_size {}", + runtime_config.num_experts, + ep_world_size + ); let epr = runtime_config.num_experts / ep_world_size; (ep_rank * epr, epr) } else { @@ -311,10 +336,12 @@ impl TrainingSession for Qwen36Session { // Move to device — for EP, narrow expert tensors before GPU transfer let num_experts = runtime_config.num_experts as i64; - let mut weights: std::collections::BTreeMap = std::collections::BTreeMap::new(); + let mut weights: std::collections::BTreeMap = + std::collections::BTreeMap::new(); for (name, tensor) in raw_weights { let needs_narrow = is_ep - && (name.contains(".mlp.experts.gate_up_proj") || name.contains(".mlp.experts.down_proj")); + && (name.contains(".mlp.experts.gate_up_proj") + || name.contains(".mlp.experts.down_proj")); if needs_narrow && tensor.size()[0] == num_experts { let narrowed = tensor .narrow(0, expert_start as i64, expert_count as i64) @@ -336,6 +363,20 @@ impl TrainingSession for Qwen36Session { } else { req.target_layers.clone() }; + let target_modules = req + .target_modules + .iter() + .map(|name| rustrain_qwen3_6::lora::Qwen36LoraTargetModule::parse(name)) + .collect::>>()?; + rustrain_qwen3_6::lora::validate_lora_targets( + &runtime_config, + &Qwen36LoraConfig { + rank: req.rank, + alpha: req.alpha, + target_layers: all_layers.clone(), + target_modules: target_modules.clone(), + }, + )?; let lora_scaling = req.alpha as f64 / req.rank as f64; let ctx = rustrain_qwen3_6::kernel::CppTrainingContext::new( &weights, @@ -348,6 +389,7 @@ impl TrainingSession for Qwen36Session { lora_scaling, req.rank, &all_layers, + &target_modules, expert_start, expert_count, )?; @@ -358,7 +400,11 @@ impl TrainingSession for Qwen36Session { if ret != 0 { return Err(anyhow!("C++ NCCL init failed (code {})", ret)); } - tracing::info!(ep_rank, ep_world_size, "NCCL communicator created in C++ for EP"); + tracing::info!( + ep_rank, + ep_world_size, + "NCCL communicator created in C++ for EP" + ); true } else { false @@ -366,10 +412,12 @@ impl TrainingSession for Qwen36Session { let count = ctx.lora_count() as usize; self.ctx = Some(ctx); - self.weights = Some(weights); // Keep alive — C++ holds raw pointers + self.weights = Some(weights); // Keep alive — C++ holds raw pointers self._nccl_ep = nccl_ep; self.lora_rank = req.rank; self.lora_alpha = req.alpha; + self.lora_target_layers = all_layers; + self.lora_target_modules = target_modules; self.lr = req.lr; self.state = SessionState::Ready { model_path: model_path.clone(), @@ -391,7 +439,9 @@ impl TrainingSession for Qwen36Session { // Record metric if let Some(ref metrics) = self.metrics { - let mem_gb = rustrain_train::metrics::gpu_memory_allocated_mb().map(|m| m / 1024.0).unwrap_or(0.0); + let mem_gb = rustrain_train::metrics::gpu_memory_allocated_mb() + .map(|m| m / 1024.0) + .unwrap_or(0.0); metrics.record_step(StepMetric { step: self.step, loss, @@ -401,21 +451,38 @@ impl TrainingSession for Qwen36Session { }); } - Ok(TrainOutput { loss, step: self.step }) + Ok(TrainOutput { + loss, + step: self.step, + }) } - fn train_multi_lora(&mut self, input: TrainInput, n_total: i32, rank: i32) -> Result { + fn train_multi_lora( + &mut self, + input: TrainInput, + n_total: i32, + rank: i32, + ) -> Result { let ctx = self .ctx .as_ref() .ok_or_else(|| anyhow!("LoRA not initialized"))?; - let loss = ctx.train_multi_lora(&input.input_ids, &input.target_mask, &input.attention_mask, n_total, rank)?; + let loss = ctx.train_multi_lora( + &input.input_ids, + &input.target_mask, + &input.attention_mask, + n_total, + rank, + )?; self.step += 1; self.last_loss = loss; self.state = SessionState::Training { step: self.step }; - Ok(TrainOutput { loss, step: self.step }) + Ok(TrainOutput { + loss, + step: self.step, + }) } fn eval_step(&self, input: TrainInput) -> Result { @@ -432,6 +499,11 @@ impl TrainingSession for Qwen36Session { .ctx .as_ref() .ok_or_else(|| anyhow!("LoRA not initialized"))?; + if !self.dynamic_lora_configs.is_empty() { + bail!( + "checkpoint-v1 stores only the fixed adapter; export dynamic adapters individually before checkpointing" + ); + } let lora_count = ctx.lora_count(); let mut lora_a = Vec::new(); @@ -466,12 +538,23 @@ impl TrainingSession for Qwen36Session { let data = checkpoint::load_checkpoint(std::path::Path::new(path))?; // Import Adam optimizer state into C++ context if let Some(ctx) = &self.ctx { + if data.lora_a.len() != data.lora_b.len() + || data.lora_a.len() != ctx.lora_count() as usize + { + return Err(anyhow!( + "checkpoint LoRA slot count mismatch: checkpoint A/B={}/{}, context={}", + data.lora_a.len(), + data.lora_b.len(), + ctx.lora_count() + )); + } + for (index, (a, b)) in data.lora_a.iter().zip(&data.lora_b).enumerate() { + ctx.set_lora_tensor(index as i64, false, a)?; + ctx.set_lora_tensor(index as i64, true, b)?; + } if !data.adam_m.is_empty() && !data.adam_v.is_empty() { ctx.import_optimizer_state(&data.adam_m, &data.adam_v)?; - tracing::info!( - imported = data.adam_m.len(), - "optimizer state imported" - ); + tracing::info!(imported = data.adam_m.len(), "optimizer state imported"); } } self.step = data.manifest.step; @@ -479,56 +562,148 @@ impl TrainingSession for Qwen36Session { Ok((data.manifest.step, data.manifest.loss)) } - fn export_adapter(&self, path: &str) -> Result { + fn export_adapter(&self, path: &str, adapter_id: Option) -> Result { let ctx = self .ctx .as_ref() .ok_or_else(|| anyhow!("LoRA not initialized"))?; - let count = ctx.lora_count() as usize; - let mut named_tensors: std::collections::BTreeMap = - std::collections::BTreeMap::new(); - for i in 0..count { - if let (Some(a), Some(b)) = (ctx.get_lora_a(i as i64), ctx.get_lora_b(i as i64)) { - named_tensors.insert( - format!("lora_a_{i}"), - a.to_kind(Kind::Float).to_device(tch::Device::Cpu), - ); - named_tensors.insert( - format!("lora_b_{i}"), - b.to_kind(Kind::Float).to_device(tch::Device::Cpu), - ); + let model_path = self + .model_path + .as_ref() + .ok_or_else(|| anyhow!("model not loaded"))?; + let runtime_config = + rustrain_qwen3_6::config::read_qwen36_runtime_config(std::path::Path::new(model_path))?; + let adapter_id = adapter_id.unwrap_or(0); + let lora_config = if adapter_id == 0 { + Qwen36LoraConfig { + rank: self.lora_rank, + alpha: self.lora_alpha, + target_layers: self.lora_target_layers.clone(), + target_modules: self.lora_target_modules.clone(), + } + } else { + self.dynamic_lora_configs + .get(&adapter_id) + .cloned() + .ok_or_else(|| anyhow!("unknown dynamic LoRA adapter: {adapter_id}"))? + }; + let slots = rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, &lora_config); + let mut exported = Vec::with_capacity(slots.len()); + for slot in slots { + if adapter_id == 0 { + let a = ctx + .get_lora_a(slot.index as i64) + .with_context(|| format!("native LoRA slot {} is missing A", slot.index))?; + let b = ctx + .get_lora_b(slot.index as i64) + .with_context(|| format!("native LoRA slot {} is missing B", slot.index))?; + exported.push((a, b)); + } else if slot.active { + let module = slot.module.cpp_name(); + let a = ctx + .get_adapter_lora_tensor(adapter_id, slot.layer as i64, module, false) + .with_context(|| { + format!( + "dynamic LoRA A is missing: adapter={adapter_id} layer={} module={module}", + slot.layer + ) + })?; + let b = ctx + .get_adapter_lora_tensor(adapter_id, slot.layer as i64, module, true) + .with_context(|| { + format!( + "dynamic LoRA B is missing: adapter={adapter_id} layer={} module={module}", + slot.layer + ) + })?; + exported.push((a, b)); + } else { + let placeholder = Tensor::zeros([], (Kind::Float, Device::Cpu)); + exported.push((placeholder.shallow_clone(), placeholder)); } } - - // Write safetensors - use std::io::Write; - let mut header = serde_json::Map::new(); - let mut offset = 0u64; - let mut all_bytes: Vec = Vec::new(); - for (name, t) in &named_tensors { - let t = t.contiguous().to_kind(Kind::Float); - let shape: Vec = t.size().iter().copied().collect(); - let data: Vec = Vec::::try_from(&t.reshape([-1]))?; - let bytes: Vec = data.iter().flat_map(|f| f.to_le_bytes()).collect(); - header.insert( - name.clone(), - serde_json::json!({"dtype":"F32","shape":shape,"data_offsets":[offset, offset + bytes.len() as u64]}), - ); - offset += bytes.len() as u64; - all_bytes.extend_from_slice(&bytes); - } - let header_str = serde_json::to_string(&serde_json::Value::Object(header))?; - let file = std::fs::File::create(path)?; - let mut writer = std::io::BufWriter::new(file); - writer.write_all(&(header_str.len() as u64).to_le_bytes())?; - writer.write_all(header_str.as_bytes())?; - writer.write_all(&all_bytes)?; - + let artifact = Qwen36AdapterArtifact::from_native_exports( + model_path, + "qwen3_hybrid_lora_sft", + Some(std::path::Path::new(model_path)), + &runtime_config, + &lora_config, + exported, + )?; + let count = artifact.tensors.len(); + artifact.save(std::path::Path::new(path))?; tracing::info!(params = count, path, "adapter exported"); Ok(count) } + fn import_adapter(&mut self, path: &str) -> Result { + let model_path = self + .model_path + .as_ref() + .ok_or_else(|| anyhow!("model not loaded"))? + .clone(); + let runtime_config = rustrain_qwen3_6::config::read_qwen36_runtime_config( + std::path::Path::new(&model_path), + )?; + let artifact = Qwen36AdapterArtifact::load(std::path::Path::new(path))?; + let target_modules = artifact + .config + .target_modules + .iter() + .map(|name| Qwen36LoraTargetModule::parse(name)) + .collect::>>()?; + let target_layers = artifact.config.target_layers.clone(); + let lora_config = Qwen36LoraConfig { + rank: artifact.config.r, + alpha: artifact.config.lora_alpha, + target_layers: target_layers.clone(), + target_modules: target_modules.clone(), + }; + rustrain_qwen3_6::lora::validate_lora_targets(&runtime_config, &lora_config)?; + + let ctx = self + .ctx + .as_ref() + .ok_or_else(|| anyhow!("LoRA not initialized"))?; + let layer_ids = target_layers + .iter() + .map(|&layer| layer as i64) + .collect::>(); + let module_csv = target_modules + .iter() + .map(Qwen36LoraTargetModule::cpp_name) + .collect::>() + .join(","); + let adapter_id = + ctx.add_lora(lora_config.rank, lora_config.alpha, &layer_ids, &module_csv)?; + + let load_result = (|| -> Result<()> { + for (name, tensor) in &artifact.tensors { + let (layer, module, is_b) = + rustrain_qwen3_6::lora::parse_adapter_tensor_name(&runtime_config, name)?; + if layer >= runtime_config.num_hidden_layers { + bail!("adapter tensor layer {layer} is outside the model"); + } + ctx.set_adapter_lora_tensor( + adapter_id, + layer as i64, + module.cpp_name(), + is_b, + tensor, + )?; + } + Ok(()) + })(); + if let Err(error) = load_result { + let _ = ctx.remove_lora(adapter_id); + return Err(error); + } + self.dynamic_lora_configs.insert(adapter_id, lora_config); + tracing::info!(adapter_id, path, "LoRA adapter imported"); + Ok(adapter_id) + } + fn status(&self) -> SessionStatus { let state = match &self.state { SessionState::Unloaded => "unloaded", @@ -558,7 +733,76 @@ impl TrainingSession for Qwen36Session { .ctx .as_ref() .ok_or_else(|| anyhow!("model not loaded — call load_model + init_lora first"))?; - let id = ctx.add_lora(req.rank, req.alpha, &req.target_layers, &req.target_modules)?; + let model_path = self + .model_path + .as_ref() + .ok_or_else(|| anyhow!("model not loaded"))?; + let runtime_config = + rustrain_qwen3_6::config::read_qwen36_runtime_config(std::path::Path::new(model_path))?; + let target_modules = if req.target_modules.trim().is_empty() { + let mut modules = std::collections::BTreeSet::new(); + for layer_type in &runtime_config.layer_types { + match layer_type { + rustrain_qwen3_6::config::LayerType::FullAttention => { + modules.extend([ + Qwen36LoraTargetModule::QProj, + Qwen36LoraTargetModule::KProj, + Qwen36LoraTargetModule::VProj, + Qwen36LoraTargetModule::OProj, + ]); + } + rustrain_qwen3_6::config::LayerType::LinearAttention => { + modules.extend([ + Qwen36LoraTargetModule::InProjQkv, + Qwen36LoraTargetModule::InProjZ, + Qwen36LoraTargetModule::InProjA, + Qwen36LoraTargetModule::InProjB, + Qwen36LoraTargetModule::OutProj, + ]); + } + } + } + if runtime_config.is_moe { + modules.extend([ + Qwen36LoraTargetModule::SharedGateProj, + Qwen36LoraTargetModule::SharedUpProj, + Qwen36LoraTargetModule::SharedDownProj, + Qwen36LoraTargetModule::ExpertsGateUpProj, + Qwen36LoraTargetModule::ExpertsDownProj, + ]); + } else { + modules.extend([ + Qwen36LoraTargetModule::GateProj, + Qwen36LoraTargetModule::UpProj, + Qwen36LoraTargetModule::DownProj, + ]); + } + modules.into_iter().collect() + } else { + req.target_modules + .split(',') + .filter(|name| !name.is_empty()) + .map(Qwen36LoraTargetModule::parse) + .collect::>>()? + }; + let config = Qwen36LoraConfig { + rank: req.rank, + alpha: req.alpha, + target_layers: req + .target_layers + .iter() + .map(|&layer| layer as usize) + .collect(), + target_modules: target_modules.clone(), + }; + rustrain_qwen3_6::lora::validate_lora_targets(&runtime_config, &config)?; + let native_module_names = target_modules + .iter() + .map(Qwen36LoraTargetModule::cpp_name) + .collect::>(); + let module_csv = native_module_names.join(","); + let id = ctx.add_lora(req.rank, req.alpha, &req.target_layers, &module_csv)?; + self.dynamic_lora_configs.insert(id, config); tracing::info!(adapter_id = id, rank = req.rank, "LoRA adapter added"); Ok(id) } @@ -570,6 +814,7 @@ impl TrainingSession for Qwen36Session { .ok_or_else(|| anyhow!("model not loaded"))?; let removed = ctx.remove_lora(adapter_id)?; if removed { + self.dynamic_lora_configs.remove(&adapter_id); tracing::info!(adapter_id, "LoRA adapter removed"); } Ok(removed) diff --git a/src/main.rs b/src/main.rs index 322a9333..cb5a3c85 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,7 +12,7 @@ use rustrain_tch_tiny::tch_train; use std::path::{Path, PathBuf}; -use anyhow::{anyhow, bail, Context, Result}; +use anyhow::{Context, Result, anyhow, bail}; use clap::{Parser, Subcommand}; use tracing::info; @@ -98,7 +98,7 @@ enum Command { #[arg(long, default_value_t = 1)] lora_rank: i32, #[arg(long, default_value_t = 60)] - duration: u64, // seconds, 0 = single step + duration: u64, // seconds, 0 = single step #[arg(long, default_value = "512")] seq_len: usize, }, @@ -199,7 +199,14 @@ fn main() -> Result<()> { lora_rank, duration, seq_len, - } => run_ep_bench(metrics_dir, world_size, n_adapters, lora_rank, duration, seq_len), + } => run_ep_bench( + metrics_dir, + world_size, + n_adapters, + lora_rank, + duration, + seq_len, + ), Command::EpWorker { shm_name, rank, @@ -595,7 +602,8 @@ fn dispatch_train(config_path: &Path, resume_from: Option) -> Result<() let tp_size = config.parallel.tensor_model_parallel_size; let cp_size = config.parallel.context_parallel_size; if tp_size > 1 || cp_size > 1 { - let summary = rustrain_glm5::session_tp_cp::train_glm5_lora_sft_tp_cp_ep(&config, &run_paths)?; + let summary = + rustrain_glm5::session_tp_cp::train_glm5_lora_sft_tp_cp_ep(&config, &run_paths)?; println!("rustrain GLM-5.2 LoRA SFT TP+CP+EP complete"); println!("run_dir: {}", run_paths.root.display()); println!("adapter_checkpoint: {}", summary.adapter_output); @@ -614,9 +622,9 @@ fn dispatch_train(config_path: &Path, resume_from: Option) -> Result<() return Ok(()); } - if is_tch && arch == "qwen3_6_lora_sft" { + if is_tch && matches!(arch, "qwen3_5_lora_sft" | "qwen3_6_lora_sft") { let summary = rustrain_qwen3_6::session::train_qwen3_6_lora_sft(&config, &run_paths)?; - println!("rustrain Qwen3.6 LoRA SFT complete"); + println!("rustrain Qwen3.5/3.6 LoRA SFT complete"); println!("run_dir: {}", run_paths.root.display()); println!("adapter_checkpoint: {}", summary.adapter_output); println!("initial_loss: {:.9}", summary.initial_loss); @@ -625,9 +633,9 @@ fn dispatch_train(config_path: &Path, resume_from: Option) -> Result<() return Ok(()); } - if is_tch && arch == "qwen3_6_lora_sft_ep" { + if is_tch && matches!(arch, "qwen3_5_lora_sft_ep" | "qwen3_6_lora_sft_ep") { let summary = rustrain_qwen3_6::session::train_qwen3_6_lora_sft_ep(&config, &run_paths)?; - println!("rustrain Qwen3.6 LoRA SFT EP complete"); + println!("rustrain Qwen3.5/3.6 LoRA SFT EP complete"); println!("run_dir: {}", run_paths.root.display()); println!("adapter_checkpoint: {}", summary.adapter_output); println!("initial_loss: {:.9}", summary.initial_loss); @@ -641,8 +649,8 @@ fn dispatch_train(config_path: &Path, resume_from: Option) -> Result<() } fn run_server(http_port: u16, grpc_port: u16, metrics_dir: PathBuf) -> Result<()> { - use rustrain_server::{api, grpc, state::SessionManager}; use rustrain_server::grpc::train::train_service_server::TrainServiceServer; + use rustrain_server::{api, grpc, state::SessionManager}; std::fs::create_dir_all(&metrics_dir)?; let manager = std::sync::Arc::new(SessionManager::new(metrics_dir.clone())); @@ -740,8 +748,10 @@ fn run_ep_bench( std::fs::create_dir_all(&metrics_dir)?; - info!("EP bench: world_size={}, n_adapters={}, rank={}, duration={}s, seq={}", - world_size, n_adapters, lora_rank, duration, seq_len); + info!( + "EP bench: world_size={}, n_adapters={}, rank={}, duration={}s, seq={}", + world_size, n_adapters, lora_rank, duration, seq_len + ); let coordinator = EpCoordinator::launch(world_size, metrics_dir.clone()) .map_err(|e| anyhow!("Failed to launch EP workers: {}", e))?; @@ -752,7 +762,8 @@ fn run_ep_bench( .unwrap_or_else(|_| "/mnt/workspace/huggingface/hub/models--Qwen--Qwen3.6-35B-A3B/snapshots/995ad96eacd98c81ed38be0c5b274b04031597b0".to_string()); let lt = "[\"linear_attention\",\"linear_attention\",\"linear_attention\",\"full_attention\"]"; let lt_full = std::iter::repeat(lt).take(8).collect::>().join(","); - let config_toml = format!(r#" + let config_toml = format!( + r#" [run] name="bench" seed=42 @@ -806,10 +817,13 @@ target_modules=["q_proj","k_proj","v_proj","o_proj","in_proj_qkv","in_proj_z","o [data] kind="instruction_jsonl" paths=["/tmp/qwen3_6_test.jsonl"] -"#); +"# + ); eprintln!("[bench] create session..."); - let _ = coordinator.dispatch(&EpCommand::CreateSession { session_id: sid.to_string() }); + let _ = coordinator.dispatch(&EpCommand::CreateSession { + session_id: sid.to_string(), + }); eprintln!("[bench] load_model..."); match coordinator.dispatch(&EpCommand::LoadModel { @@ -817,7 +831,7 @@ paths=["/tmp/qwen3_6_test.jsonl"] model_path: model_path.to_string(), config_toml: config_toml.clone(), }) { - EpResult::Ok => {}, + EpResult::Ok => {} EpResult::Error(e) => bail!("load_model failed: {}", e), _ => bail!("load_model unexpected result"), } @@ -833,22 +847,31 @@ paths=["/tmp/qwen3_6_test.jsonl"] match coordinator.dispatch(&EpCommand::InitLora { session_id: sid.to_string(), rank: 8, - alpha: 16, + alpha: 16.0, target_layers: vec![], - target_modules: vec!["q_proj".to_string(),"k_proj".to_string(),"v_proj".to_string(), - "o_proj".to_string(),"in_proj_qkv".to_string(), - "in_proj_z".to_string(),"out_proj".to_string()], + target_modules: vec![ + "q_proj".to_string(), + "k_proj".to_string(), + "v_proj".to_string(), + "o_proj".to_string(), + "in_proj_qkv".to_string(), + "in_proj_z".to_string(), + "out_proj".to_string(), + ], lr: 0.0001, beta1: 0.9, beta2: 0.999, eps: 0.00000001, }) { - EpResult::Count(_) => {}, + EpResult::Count(_) => {} EpResult::Error(e) => bail!("init_lora failed: {}", e), _ => bail!("init_lora unexpected result"), } - eprintln!("[bench] batch_add_lora ({} adapters, rank={})...", n_adapters, lora_rank); + eprintln!( + "[bench] batch_add_lora ({} adapters, rank={})...", + n_adapters, lora_rank + ); match coordinator.dispatch(&EpCommand::BatchAddLora { session_id: sid.to_string(), count: n_adapters, @@ -864,7 +887,11 @@ paths=["/tmp/qwen3_6_test.jsonl"] // Build input tensors let ids: Vec = vec![1; seq_len]; - let mask: Vec = { let mut m = vec![0i64; 20]; m.extend(vec![1i64; seq_len - 20]); m }; + let mask: Vec = { + let mut m = vec![0i64; 20]; + m.extend(vec![1i64; seq_len - 20]); + m + }; let attn: Vec = vec![1; seq_len]; // Warmup @@ -880,16 +907,23 @@ paths=["/tmp/qwen3_6_test.jsonl"] lora_rank, }) { EpResult::Loss(l) => l, - EpResult::Error(e) => { bail!("warmup failed: {}", e); } + EpResult::Error(e) => { + bail!("warmup failed: {}", e); + } _ => bail!("warmup unexpected result"), }; let warmup_ms = t0.elapsed().as_millis(); - eprintln!("[bench] warmup: loss={:.6} time={}ms", warmup_loss, warmup_ms); + eprintln!( + "[bench] warmup: loss={:.6} time={}ms", + warmup_loss, warmup_ms + ); if duration == 0 { // Single step only eprintln!("[bench] single step done"); - let _ = coordinator.dispatch(&EpCommand::DeleteSession { session_id: sid.to_string() }); + let _ = coordinator.dispatch(&EpCommand::DeleteSession { + session_id: sid.to_string(), + }); return Ok(()); } @@ -918,8 +952,15 @@ paths=["/tmp/qwen3_6_test.jsonl"] let elapsed = start.elapsed().as_secs_f64(); if total_steps % 5 == 0 || total_steps == 1 { let rate = total_adapters as f64 / elapsed; - eprintln!(" step {}: loss={:.6} time={}ms total={} adp in {:.0}s = {:.1} adp/s", - total_steps, l, t0.elapsed().as_millis(), total_adapters, elapsed, rate); + eprintln!( + " step {}: loss={:.6} time={}ms total={} adp in {:.0}s = {:.1} adp/s", + total_steps, + l, + t0.elapsed().as_millis(), + total_adapters, + elapsed, + rate + ); } } EpResult::Error(e) => { @@ -935,17 +976,30 @@ paths=["/tmp/qwen3_6_test.jsonl"] eprintln!(); eprintln!("============================================================"); - eprintln!(" RESULTS: {} adapters, rank={}, EP={}", n_adapters, lora_rank, world_size); + eprintln!( + " RESULTS: {} adapters, rank={}, EP={}", + n_adapters, lora_rank, world_size + ); eprintln!("============================================================"); eprintln!(" Duration: {:.0}s ({:.1} min)", elapsed, elapsed / 60.0); eprintln!(" Steps: {}", total_steps); eprintln!(" Total adapters processed: {}", total_adapters); - eprintln!(" Throughput: {:.2} adapters/s ({:.0} adapters/min)", rate, rate * 60.0); + eprintln!( + " Throughput: {:.2} adapters/s ({:.0} adapters/min)", + rate, + rate * 60.0 + ); if !losses.is_empty() { - eprintln!(" Loss: {:.6} -> {:.6}", losses[0], losses[losses.len()-1]); + eprintln!( + " Loss: {:.6} -> {:.6}", + losses[0], + losses[losses.len() - 1] + ); } eprintln!(" Failures: {}", total_steps - losses.len() as i64); - let _ = coordinator.dispatch(&EpCommand::DeleteSession { session_id: sid.to_string() }); + let _ = coordinator.dispatch(&EpCommand::DeleteSession { + session_id: sid.to_string(), + }); Ok(()) } From e2a03a933562e275c7c2190274f1f92e1f252f6b Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 04:09:23 +0800 Subject: [PATCH 002/156] feat: persist dynamic lora and add replicated dp --- crates/rustrain-core/src/runtime.rs | 46 ++- .../kernels/qwen3_6_kernels.cpp | 80 ++++- crates/rustrain-qwen3-6/src/kernel.rs | 81 ++++- crates/rustrain-qwen3-6/src/session.rs | 41 ++- .../rustrain-qwen3-6/tests/native_smoke.cpp | 10 +- crates/rustrain-server/src/checkpoint.rs | 243 +++++++++++++-- crates/rustrain-server/src/session.rs | 278 +++++++++++++++++- 7 files changed, 729 insertions(+), 50 deletions(-) diff --git a/crates/rustrain-core/src/runtime.rs b/crates/rustrain-core/src/runtime.rs index 77aa09d1..2c6d8ebd 100644 --- a/crates/rustrain-core/src/runtime.rs +++ b/crates/rustrain-core/src/runtime.rs @@ -494,6 +494,11 @@ pub fn validate_config(config: &Config) -> Result<()> { config.model.architecture.as_str(), "qwen3_5_lora_sft" | "qwen3_5_lora_sft_ep" | "qwen3_6_lora_sft" | "qwen3_6_lora_sft_ep" ); + let is_qwen3_hybrid_lora_sft_dp = matches!(config.train.backend, BackendKind::Tch) + && matches!( + config.model.architecture.as_str(), + "qwen3_5_lora_sft" | "qwen3_6_lora_sft" + ); let is_tch_moe_ep_session = matches!(config.train.backend, BackendKind::Tch) && config.model.architecture == "tch_moe_ep_session"; let is_v4_tp_rank = matches!(config.train.backend, BackendKind::Tch) @@ -527,6 +532,13 @@ pub fn validate_config(config: &Config) -> Result<()> { && !((is_tch_tiny_lm || is_qwen_trainable_session) && name == "data_parallel_size" && value == 2) + && !(is_qwen3_hybrid_lora_sft_dp + && name == "data_parallel_size" + && value >= 2 + && parallel.tensor_model_parallel_size == 1 + && parallel.pipeline_model_parallel_size == 1 + && parallel.expert_model_parallel_size == 1 + && parallel.context_parallel_size == 1) && !(is_qwen_trainable_session && name == "tensor_model_parallel_size" && value == 2 @@ -677,9 +689,21 @@ pub fn validate_config(config: &Config) -> Result<()> { config.model.architecture )); } - let expected_global_batch_size = - config.train.micro_batch_size * config.train.gradient_accumulation_steps; + let data_parallel_factor = if is_qwen3_hybrid_lora_sft_dp { + config.parallel.data_parallel_size + } else { + 1 + }; + let expected_global_batch_size = config.train.micro_batch_size + * config.train.gradient_accumulation_steps + * data_parallel_factor; if config.train.global_batch_size != expected_global_batch_size { + if data_parallel_factor > 1 { + return Err(anyhow!( + "{} requires global_batch_size = micro_batch_size * gradient_accumulation_steps * data_parallel_size", + config.model.architecture + )); + } return Err(anyhow!( "{} requires global_batch_size = micro_batch_size * gradient_accumulation_steps", config.model.architecture @@ -1331,6 +1355,24 @@ mod tests { validate_config(&config).expect("native Qwen3.5 LoRA targets should validate"); } + #[test] + fn qwen_hybrid_lora_sft_accepts_replicated_data_parallelism() { + let mut config = qwen_lora_sft_config(); + config.model.architecture = "qwen3_6_lora_sft".to_string(); + config.parallel.data_parallel_size = 2; + config.train.global_batch_size = config.train.micro_batch_size + * config.train.gradient_accumulation_steps + * config.parallel.data_parallel_size; + validate_config(&config).expect("Qwen3.6 LoRA DP should validate"); + + config.model.architecture = "qwen3_5_lora_sft".to_string(); + validate_config(&config).expect("Qwen3.5 LoRA DP should validate"); + + config.train.global_batch_size /= 2; + let error = validate_config(&config).expect_err("DP global batch mismatch should fail"); + assert!(error.to_string().contains("data_parallel_size")); + } + #[test] fn data_max_samples_must_be_positive_when_set() { let mut config = qwen_lora_sft_config(); diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 3e30f4a6..38d2d247 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -3412,7 +3412,7 @@ static at::Tensor mtp_compute_loss( extern "C" { __attribute__((visibility("default"))) int64_t qwen36_kernel_abi_version() { - return 4; + return 6; } // Create training context — called once at startup @@ -4558,6 +4558,35 @@ int64_t qwen36_add_lora( } } +// Restore a dynamic adapter's externally visible ID during checkpoint load. +// IDs are positive and unique; the monotonic allocator is advanced so future +// additions cannot collide with a restored tenant. +__attribute__((visibility("default"))) +int32_t qwen36_set_adapter_id(void* ctx_ptr, int64_t current_id, int64_t requested_id) { + try { + auto* ctx = reinterpret_cast(ctx_ptr); + TORCH_CHECK(ctx && current_id > 0 && requested_id > 0, + "dynamic adapter IDs must be positive"); + for (const auto& adapter : ctx->adapters) { + TORCH_CHECK(adapter.id != requested_id || adapter.id == current_id, + "dynamic adapter ID already exists: ", requested_id); + } + for (auto& adapter : ctx->adapters) { + if (adapter.id == current_id) { + adapter.id = requested_id; + ctx->next_adapter_id = std::max(ctx->next_adapter_id, requested_id); + ctx->lora_cache_valid = false; + ctx->lora_batch_valid = false; + return 0; + } + } + TORCH_CHECK(false, "dynamic adapter not found: ", current_id); + } catch (const std::exception& e) { + fprintf(stderr, "[q36] set_adapter_id FAILED: %s\n", e.what()); + return -1; + } +} + __attribute__((visibility("default"))) int32_t qwen36_remove_lora(void* ctx_ptr, int64_t adapter_id) { auto* ctx = reinterpret_cast(ctx_ptr); @@ -4635,6 +4664,55 @@ int32_t qwen36_set_adapter_lora_tensor( } } +// Access one dynamic adapter's Adam state. The per-layer state is stored as +// {m_a, v_a, m_b, v_b}; `is_b` selects A/B and `is_v` selects m/v. +__attribute__((visibility("default"))) +void* qwen36_get_adapter_optimizer_tensor( + void* ctx_ptr, int64_t adapter_id, int64_t layer_idx, + const char* module_name, int32_t is_b, int32_t is_v +) { + auto* ctx = reinterpret_cast(ctx_ptr); + if (!ctx || !module_name || layer_idx < 0 || layer_idx >= ctx->num_layers) + return nullptr; + const int64_t pair_idx = lora_pair_index( + ctx->layer_configs[layer_idx], module_name); + if (pair_idx < 0) return nullptr; + for (auto& adapter : ctx->adapters) { + if (adapter.id != adapter_id) continue; + auto state_it = adapter.adam_state.find(layer_idx); + if (state_it == adapter.adam_state.end() || + pair_idx >= static_cast(state_it->second.size())) + return nullptr; + // array order: m_a, v_a, m_b, v_b + const int index = (is_b ? 2 : 0) + (is_v ? 1 : 0); + return &state_it->second[pair_idx][index]; + } + return nullptr; +} + +__attribute__((visibility("default"))) +int32_t qwen36_set_adapter_optimizer_tensor( + void* ctx_ptr, int64_t adapter_id, int64_t layer_idx, + const char* module_name, int32_t is_b, int32_t is_v, void* tensor_ptr +) { + try { + auto* target = reinterpret_cast( + qwen36_get_adapter_optimizer_tensor( + ctx_ptr, adapter_id, layer_idx, module_name, is_b, is_v)); + TORCH_CHECK(target && tensor_ptr, "dynamic optimizer tensor not found"); + auto& source = *reinterpret_cast(tensor_ptr); + TORCH_CHECK(source.sizes() == target->sizes(), + "dynamic optimizer tensor shape mismatch: expected ", target->sizes(), + " got ", source.sizes()); + at::NoGradGuard guard; + target->copy_(source.to(target->device()).to(target->scalar_type())); + return 0; + } catch (const std::exception& e) { + fprintf(stderr, "[q36] set_adapter_optimizer_tensor FAILED: %s\n", e.what()); + return -1; + } +} + __attribute__((visibility("default"))) int64_t qwen36_get_lora_count(void* ctx_ptr) { auto* ctx = reinterpret_cast(ctx_ptr); diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs index 529399fe..1554c23a 100644 --- a/crates/rustrain-qwen3-6/src/kernel.rs +++ b/crates/rustrain-qwen3-6/src/kernel.rs @@ -4,7 +4,7 @@ //! Rust only handles: weight loading, data loading, training loop orchestration. use crate::lora::Qwen36LoraTargetModule; -use anyhow::{Result, bail}; +use anyhow::{bail, Result}; use std::ffi::c_void; use std::sync::OnceLock; use tch::{Kind, Tensor}; @@ -71,6 +71,11 @@ type FnGetAdapterLoraTensor = unsafe extern "C" fn(*mut c_void, i64, i64, *const i8, i32) -> *mut c_void; type FnSetAdapterLoraTensor = unsafe extern "C" fn(*mut c_void, i64, i64, *const i8, i32, *mut c_void) -> i32; +type FnSetAdapterId = unsafe extern "C" fn(*mut c_void, i64, i64) -> i32; +type FnGetAdapterOptimizerTensor = + unsafe extern "C" fn(*mut c_void, i64, i64, *const i8, i32, i32) -> *mut c_void; +type FnSetAdapterOptimizerTensor = + unsafe extern "C" fn(*mut c_void, i64, i64, *const i8, i32, i32, *mut c_void) -> i32; #[repr(C)] pub struct CppLayerConfig { @@ -123,6 +128,9 @@ struct KernelHandles { list_lora: FnListLora, get_adapter_lora_tensor: FnGetAdapterLoraTensor, set_adapter_lora_tensor: FnSetAdapterLoraTensor, + set_adapter_id: FnSetAdapterId, + get_adapter_optimizer_tensor: FnGetAdapterOptimizerTensor, + set_adapter_optimizer_tensor: FnSetAdapterOptimizerTensor, } static KERNELS: OnceLock> = OnceLock::new(); @@ -163,7 +171,7 @@ unsafe fn load_kernels() -> Option { }}; } let abi_version: FnKernelAbiVersion = sym!("qwen36_kernel_abi_version"); - if abi_version() != 4 { + if abi_version() != 6 { return None; } Some(KernelHandles { @@ -191,6 +199,9 @@ unsafe fn load_kernels() -> Option { list_lora: sym!("qwen36_list_lora"), get_adapter_lora_tensor: sym!("qwen36_get_adapter_lora_tensor"), set_adapter_lora_tensor: sym!("qwen36_set_adapter_lora_tensor"), + set_adapter_id: sym!("qwen36_set_adapter_id"), + get_adapter_optimizer_tensor: sym!("qwen36_get_adapter_optimizer_tensor"), + set_adapter_optimizer_tensor: sym!("qwen36_set_adapter_optimizer_tensor"), }) } @@ -725,6 +736,16 @@ impl CppTrainingContext { Ok(found != 0) } + /// Restore a dynamic adapter's stable external ID from a checkpoint. + pub fn set_adapter_id(&self, current_id: i64, requested_id: i64) -> Result<()> { + let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; + let status = unsafe { (kh.set_adapter_id)(self.ptr, current_id, requested_id) }; + if status != 0 { + bail!("C++ set_adapter_id failed: {current_id} -> {requested_id}"); + } + Ok(()) + } + /// List all active adapter IDs. pub fn list_lora(&self) -> Vec { let kh = match get_kernels() { @@ -794,6 +815,62 @@ impl CppTrainingContext { Ok(()) } + pub fn get_adapter_optimizer_tensor( + &self, + adapter_id: i64, + layer: i64, + module: &str, + is_b: bool, + is_v: bool, + ) -> Option { + let kh = get_kernels()?; + let module = std::ffi::CString::new(module).ok()?; + let ptr = unsafe { + (kh.get_adapter_optimizer_tensor)( + self.ptr, + adapter_id, + layer, + module.as_ptr(), + if is_b { 1 } else { 0 }, + if is_v { 1 } else { 0 }, + ) + }; + if ptr.is_null() { + return None; + } + Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) }) + } + + pub fn set_adapter_optimizer_tensor( + &self, + adapter_id: i64, + layer: i64, + module: &str, + is_b: bool, + is_v: bool, + tensor: &Tensor, + ) -> Result<()> { + let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; + let module = std::ffi::CString::new(module)?; + let status = unsafe { + (kh.set_adapter_optimizer_tensor)( + self.ptr, + adapter_id, + layer, + module.as_ptr(), + if is_b { 1 } else { 0 }, + if is_v { 1 } else { 0 }, + tensor.as_ptr() as *mut c_void, + ) + }; + if status != 0 { + bail!( + "C++ set_adapter_optimizer_tensor failed for adapter {adapter_id}, layer {layer}, module {module:?}" + ); + } + Ok(()) + } + /// Eval step: forward + loss, no backward, no Adam update. pub fn eval_step( &self, diff --git a/crates/rustrain-qwen3-6/src/session.rs b/crates/rustrain-qwen3-6/src/session.rs index 3e4be84d..3fb07321 100644 --- a/crates/rustrain-qwen3-6/src/session.rs +++ b/crates/rustrain-qwen3-6/src/session.rs @@ -249,8 +249,31 @@ fn train_impl( let shard_ref = ep_shard.as_ref(); let is_ep = shard_ref.is_some(); - let world_size = shard_ref.map(|s| s.world_size).unwrap_or(1); - let rank = shard_ref.map(|s| s.rank).unwrap_or(0); + // Non-EP Qwen sessions may run replicated-weight LoRA data parallelism. + // The launcher supplies the standard torchrun environment; EP keeps its + // explicit shard metadata as the source of truth. + let env_world_size = std::env::var("WORLD_SIZE") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(1); + let env_rank = std::env::var("RANK") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + let world_size = shard_ref.map(|s| s.world_size).unwrap_or(env_world_size); + let rank = shard_ref.map(|s| s.rank).unwrap_or(env_rank); + let is_data_parallel = !is_ep && world_size > 1; + if is_data_parallel && runtime_config.is_moe { + bail!("replicated Qwen data parallelism is only supported for dense/linear-attention models; use *_ep for MoE"); + } + if is_data_parallel { + crate::kernel::CppTrainingContext::set_cuda_device( + std::env::var("LOCAL_RANK") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0), + ); + } // Build needed weight set let needed = build_needed_weights(&runtime_config, &lora_config, shard_ref); @@ -356,6 +379,14 @@ fn train_impl( shard_ref.map(|s| s.expert_start).unwrap_or(0), shard_ref.map(|s| s.experts_per_rank).unwrap_or(0), )?; + + if world_size > 1 { + let ret = ctx.init_nccl(); + if ret != 0 { + bail!("C++ NCCL init failed (code {})", ret); + } + info!(rank, world_size, is_ep, "NCCL communicator created for Qwen LoRA parallel training"); + } info!("C++ TrainingContext: {} LoRA params", ctx.lora_count()); // Set MTP weights if available @@ -384,7 +415,11 @@ fn train_impl( let mut final_loss = 0.0_f64; for step in 0..max_steps { - let data_start = (step * batch_size) % data.len(); + let data_start = if is_data_parallel { + (step * batch_size * world_size + rank * batch_size) % data.len() + } else { + (step * batch_size) % data.len() + }; let sft_batch = data.batch(data_start, batch_size); let (input_ids, target_mask) = sft_batch.to_tensors(device, compute_kind); diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index b9364758..ddebdfbb 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -23,6 +23,7 @@ extern "C" void* qwen36_create_training_context( void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, double, double, double, double, double, int64_t, double, int64_t, const int64_t*, int64_t, const char*); +extern "C" int32_t qwen36_init_nccl(void*); extern "C" int64_t qwen36_get_lora_count(void*); extern "C" void* qwen36_get_lora_a(void*, int64_t); extern "C" void* qwen36_get_lora_b(void*, int64_t); @@ -44,7 +45,11 @@ static at::Tensor cuda_rand(std::initializer_list shape) { } int main() { - c10::cuda::CUDAGuard guard(0); + const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); + const int process_rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0"); + const int local_rank = std::atoi(std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); + assert(world == 1 || (world == 2 && process_rank >= 0 && process_rank < world)); + c10::cuda::CUDAGuard guard(local_rank); at::manual_seed(7); constexpr int64_t hidden = 16; @@ -265,6 +270,9 @@ int main() { 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank, &target_layer, 1, "q_proj"); assert(ctx); + if (world > 1) { + assert(qwen36_init_nccl(ctx) == 0); + } const char* dense_targets = "gate_proj,up_proj,down_proj"; const int64_t dense_one = qwen36_add_lora( ctx, rank, 1.0, &target_layer, 1, dense_targets); diff --git a/crates/rustrain-server/src/checkpoint.rs b/crates/rustrain-server/src/checkpoint.rs index 18dd40d1..2042255e 100644 --- a/crates/rustrain-server/src/checkpoint.rs +++ b/crates/rustrain-server/src/checkpoint.rs @@ -14,6 +14,27 @@ pub struct CheckpointManifest { pub lora_rank: i64, pub lora_alpha: f64, pub files: Vec, + #[serde(default)] + pub dynamic_adapters: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DynamicAdapterManifest { + pub id: i64, + pub rank: i64, + pub alpha: f64, + pub target_layers: Vec, + pub target_modules: Vec, + pub parameter_count: usize, + pub optimizer_count: usize, +} + +pub struct DynamicAdapterCheckpoint { + pub manifest: DynamicAdapterManifest, + pub lora_a: Vec, + pub lora_b: Vec, + pub adam_m: Vec, + pub adam_v: Vec, } pub struct CheckpointData { @@ -22,6 +43,7 @@ pub struct CheckpointData { pub lora_b: Vec, pub adam_m: Vec, pub adam_v: Vec, + pub dynamic_adapters: Vec, } /// Save checkpoint to a directory. @@ -37,27 +59,92 @@ pub fn save_checkpoint( lora_b: &[Tensor], adam_m: &[Tensor], adam_v: &[Tensor], +) -> Result<()> { + save_checkpoint_with_dynamic( + dir, + step, + loss, + model_path, + lora_rank, + lora_alpha, + lora_a, + lora_b, + adam_m, + adam_v, + &[], + ) +} + +pub fn save_checkpoint_with_dynamic( + dir: &Path, + step: u64, + loss: f64, + model_path: &str, + lora_rank: i64, + lora_alpha: f64, + lora_a: &[Tensor], + lora_b: &[Tensor], + adam_m: &[Tensor], + adam_v: &[Tensor], + dynamic_adapters: &[DynamicAdapterCheckpoint], ) -> Result<()> { std::fs::create_dir_all(dir) .with_context(|| format!("create checkpoint dir {}", dir.display()))?; // Save adapter (LoRA A/B) as safetensors let adapter_path = dir.join("adapter.safetensors"); - save_tensors(&adapter_path, &lora_a, &lora_b)?; + let mut adapter_tensors = named_tensors(lora_a, lora_b, "a_", "b_"); + let mut dynamic_manifests = Vec::with_capacity(dynamic_adapters.len()); + for adapter in dynamic_adapters { + if adapter.lora_a.len() != adapter.lora_b.len() + || adapter.adam_m.len() != adapter.adam_v.len() + || adapter.manifest.parameter_count != adapter.lora_a.len() + || adapter.manifest.optimizer_count != adapter.adam_m.len() + { + anyhow::bail!( + "dynamic adapter {} checkpoint count mismatch", + adapter.manifest.id + ); + } + let id = adapter.manifest.id; + adapter_tensors.extend(named_tensors( + &adapter.lora_a, + &adapter.lora_b, + &format!("dynamic_{id}_a_"), + &format!("dynamic_{id}_b_"), + )); + dynamic_manifests.push(adapter.manifest.clone()); + } + save_named_tensors(&adapter_path, adapter_tensors)?; // Save optimizer state (Adam m/v) as safetensors let optimizer_path = dir.join("optimizer.safetensors"); - save_tensors(&optimizer_path, &adam_m, &adam_v)?; + let mut optimizer_tensors = named_tensors(adam_m, adam_v, "a_", "b_"); + for adapter in dynamic_adapters { + let id = adapter.manifest.id; + optimizer_tensors.extend(named_tensors( + &adapter.adam_m, + &adapter.adam_v, + &format!("dynamic_{id}_a_"), + &format!("dynamic_{id}_b_"), + )); + } + save_named_tensors(&optimizer_path, optimizer_tensors)?; // Write manifest let manifest = CheckpointManifest { - format: "rustrain-checkpoint-v1".to_string(), + format: if dynamic_manifests.is_empty() { + "rustrain-checkpoint-v1".to_string() + } else { + "rustrain-checkpoint-v2".to_string() + }, step, loss, model_path: model_path.to_string(), lora_rank, lora_alpha, files: vec!["adapter.safetensors".into(), "optimizer.safetensors".into()], + dynamic_adapters: dynamic_manifests, }; let manifest_path = dir.join("manifest.json"); std::fs::write(&manifest_path, serde_json::to_string_pretty(&manifest)?) @@ -82,10 +169,36 @@ pub fn load_checkpoint(dir: &Path) -> Result { .with_context(|| "parse manifest.json")?; let adapter_path = dir.join("adapter.safetensors"); - let (lora_a, lora_b) = load_tensors(&adapter_path)?; + let adapter_named = read_named_tensors(&adapter_path)?; + let lora_a = collect_side(&adapter_named, "a_")?; + let lora_b = collect_side(&adapter_named, "b_")?; let optimizer_path = dir.join("optimizer.safetensors"); - let (adam_m, adam_v) = load_tensors(&optimizer_path)?; + let optimizer_named = read_named_tensors(&optimizer_path)?; + let adam_m = collect_side(&optimizer_named, "a_")?; + let adam_v = collect_side(&optimizer_named, "b_")?; + let mut dynamic_adapters = Vec::with_capacity(manifest.dynamic_adapters.len()); + for dynamic_manifest in &manifest.dynamic_adapters { + let id = dynamic_manifest.id; + let dynamic_lora_a = collect_side(&adapter_named, &format!("dynamic_{id}_a_"))?; + let dynamic_lora_b = collect_side(&adapter_named, &format!("dynamic_{id}_b_"))?; + let dynamic_adam_m = collect_side(&optimizer_named, &format!("dynamic_{id}_a_"))?; + let dynamic_adam_v = collect_side(&optimizer_named, &format!("dynamic_{id}_b_"))?; + if dynamic_lora_a.len() != dynamic_manifest.parameter_count + || dynamic_lora_b.len() != dynamic_manifest.parameter_count + || dynamic_adam_m.len() != dynamic_manifest.optimizer_count + || dynamic_adam_v.len() != dynamic_manifest.optimizer_count + { + anyhow::bail!("dynamic adapter {id} checkpoint tensor count mismatch"); + } + dynamic_adapters.push(DynamicAdapterCheckpoint { + manifest: dynamic_manifest.clone(), + lora_a: dynamic_lora_a, + lora_b: dynamic_lora_b, + adam_m: dynamic_adam_m, + adam_v: dynamic_adam_v, + }); + } tracing::info!( step = manifest.step, @@ -99,59 +212,70 @@ pub fn load_checkpoint(dir: &Path) -> Result { lora_b, adam_m, adam_v, + dynamic_adapters, }) } -fn save_tensors(path: &Path, a: &[Tensor], b: &[Tensor]) -> Result<()> { +fn named_tensors( + a: &[Tensor], + b: &[Tensor], + a_prefix: &str, + b_prefix: &str, +) -> Vec<(String, Tensor)> { let mut named: Vec<(String, Tensor)> = Vec::new(); for (i, t) in a.iter().enumerate() { named.push(( - format!("a_{i}"), + format!("{a_prefix}{i}"), t.to_kind(tch::Kind::Float).to_device(tch::Device::Cpu), )); } for (i, t) in b.iter().enumerate() { named.push(( - format!("b_{i}"), + format!("{b_prefix}{i}"), t.to_kind(tch::Kind::Float).to_device(tch::Device::Cpu), )); } + named +} + +fn save_named_tensors(path: &Path, named: Vec<(String, Tensor)>) -> Result<()> { Tensor::write_safetensors(&named, path).with_context(|| format!("write {}", path.display()))?; Ok(()) } -fn load_tensors(path: &Path) -> Result<(Vec, Vec)> { +fn read_named_tensors(path: &Path) -> Result> { let named = Tensor::read_safetensors(path).with_context(|| format!("read {}", path.display()))?; let mut by_name = std::collections::BTreeMap::new(); for (name, tensor) in named { by_name.insert(name, tensor); } - fn collect_side( - tensors: &std::collections::BTreeMap, - prefix: &str, - ) -> Result> { - let mut indices = tensors - .keys() - .filter_map(|name| name.strip_prefix(prefix)?.parse::().ok()) - .collect::>(); - indices.sort_unstable(); - let mut result = Vec::with_capacity(indices.len()); - for (expected, index) in indices.into_iter().enumerate() { - if index != expected { - anyhow::bail!("checkpoint tensor indices for {prefix} are not contiguous"); - } - result.push( - tensors - .get(&format!("{prefix}{index}")) - .expect("index collected from map") - .shallow_clone(), - ); + Ok(by_name) +} + +fn collect_side( + tensors: &std::collections::BTreeMap, + prefix: &str, +) -> Result> { + let mut indices = tensors + .keys() + .filter_map(|name| name.strip_prefix(prefix)?.parse::().ok()) + .collect::>(); + indices.sort_unstable(); + let mut result = Vec::with_capacity(indices.len()); + for (expected, index) in indices.into_iter().enumerate() { + if index != expected { + anyhow::bail!("checkpoint tensor indices for {prefix} are not contiguous"); } - Ok(result) + result.push( + tensors + .get(&format!("{prefix}{index}")) + .expect("index collected from map") + .shallow_clone(), + ); } - Ok((collect_side(&by_name, "a_")?, collect_side(&by_name, "b_")?)) + Ok(result) } #[cfg(test)] @@ -188,4 +312,61 @@ mod tests { assert!(loaded.adam_m[0].allclose(&m, 1e-6, 1e-6, false)); assert!(loaded.adam_v[0].allclose(&v, 1e-6, 1e-6, false)); } + + #[test] + fn dynamic_checkpoint_roundtrip_preserves_metadata_and_state() { + let dir = tempfile::tempdir().unwrap(); + let dynamic = DynamicAdapterCheckpoint { + manifest: DynamicAdapterManifest { + id: 7, + rank: 3, + alpha: 6.0, + target_layers: vec![1, 3], + target_modules: vec!["q_proj".into(), "down_proj".into()], + parameter_count: 2, + optimizer_count: 4, + }, + lora_a: (0..2) + .map(|_| Tensor::ones([3, 8], (tch::Kind::Float, tch::Device::Cpu))) + .collect(), + lora_b: (0..2) + .map(|_| Tensor::full([16, 3], 2.0, (tch::Kind::Float, tch::Device::Cpu))) + .collect(), + adam_m: (0..4) + .map(|_| Tensor::full([3, 8], 3.0, (tch::Kind::Float, tch::Device::Cpu))) + .collect(), + adam_v: (0..4) + .map(|_| Tensor::full([3, 8], 4.0, (tch::Kind::Float, tch::Device::Cpu))) + .collect(), + }; + save_checkpoint_with_dynamic( + dir.path(), + 11, + 0.5, + "Qwen/test", + 2, + 4.0, + &[], + &[], + &[], + &[], + &[dynamic], + ) + .unwrap(); + let loaded = load_checkpoint(dir.path()).unwrap(); + assert_eq!(loaded.manifest.format, "rustrain-checkpoint-v2"); + assert_eq!(loaded.dynamic_adapters.len(), 1); + let loaded_dynamic = &loaded.dynamic_adapters[0]; + assert_eq!(loaded_dynamic.manifest.id, 7); + assert_eq!(loaded_dynamic.manifest.rank, 3); + assert_eq!(loaded_dynamic.manifest.target_layers, vec![1, 3]); + assert_eq!(loaded_dynamic.lora_a.len(), 2); + assert_eq!(loaded_dynamic.adam_m.len(), 4); + assert!(loaded_dynamic.adam_m[0].allclose( + &Tensor::full([3, 8], 3.0, (tch::Kind::Float, tch::Device::Cpu)), + 1e-6, + 1e-6, + false + )); + } } diff --git a/crates/rustrain-server/src/session.rs b/crates/rustrain-server/src/session.rs index d057293b..0fddea94 100644 --- a/crates/rustrain-server/src/session.rs +++ b/crates/rustrain-server/src/session.rs @@ -314,6 +314,7 @@ impl TrainingSession for Qwen36Session { .and_then(|s| s.parse::().ok()) .unwrap_or(0); let is_ep = ep_world_size > 1 && runtime_config.is_moe; + let is_data_parallel = ep_world_size > 1 && !runtime_config.is_moe; // Compute expert shard let (expert_start, expert_count) = if is_ep { @@ -329,8 +330,9 @@ impl TrainingSession for Qwen36Session { (0, runtime_config.num_experts) }; - // Set CUDA device for EP - if is_ep { + // Set CUDA device for any torchrun worker. Dense Qwen workers use + // replicated weights and NCCL gradient all-reduce (LoRA-only DP). + if is_ep || is_data_parallel { self.device = tch::Device::Cuda(local_rank); } @@ -394,8 +396,9 @@ impl TrainingSession for Qwen36Session { expert_count, )?; - // Initialize NCCL communicator for EP — directly in C++ - let nccl_ep = if is_ep { + // Initialize NCCL directly in C++. The same communicator handles EP + // output collectives and replicated-weight LoRA gradient all-reduce. + let nccl_ep = if is_ep || is_data_parallel { let ret = ctx.init_nccl(); if ret != 0 { return Err(anyhow!("C++ NCCL init failed (code {})", ret)); @@ -403,6 +406,7 @@ impl TrainingSession for Qwen36Session { tracing::info!( ep_rank, ep_world_size, + data_parallel = is_data_parallel, "NCCL communicator created in C++ for EP" ); true @@ -499,11 +503,6 @@ impl TrainingSession for Qwen36Session { .ctx .as_ref() .ok_or_else(|| anyhow!("LoRA not initialized"))?; - if !self.dynamic_lora_configs.is_empty() { - bail!( - "checkpoint-v1 stores only the fixed adapter; export dynamic adapters individually before checkpointing" - ); - } let lora_count = ctx.lora_count(); let mut lora_a = Vec::new(); @@ -518,7 +517,136 @@ impl TrainingSession for Qwen36Session { // Export Adam optimizer state let (adam_m, adam_v) = ctx.export_optimizer_state()?; - checkpoint::save_checkpoint( + let mut dynamic_adapters = Vec::new(); + if !self.dynamic_lora_configs.is_empty() { + let model_path = self + .model_path + .as_ref() + .ok_or_else(|| anyhow!("model path unavailable for dynamic LoRA checkpoint"))?; + let runtime_config = rustrain_qwen3_6::config::read_qwen36_runtime_config( + std::path::Path::new(model_path), + )?; + for (&adapter_id, lora_config) in &self.dynamic_lora_configs { + let slots = rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, lora_config); + let mut dynamic_a = Vec::new(); + let mut dynamic_b = Vec::new(); + let mut dynamic_m = Vec::new(); + let mut dynamic_v = Vec::new(); + for slot in slots.iter().filter(|slot| slot.active) { + let module = slot.module.cpp_name(); + dynamic_a.push( + ctx.get_adapter_lora_tensor( + adapter_id, + slot.layer as i64, + module, + false, + ) + .with_context(|| { + format!( + "dynamic LoRA A is missing: adapter={adapter_id} layer={} module={module}", + slot.layer + ) + })?, + ); + dynamic_b.push( + ctx.get_adapter_lora_tensor( + adapter_id, + slot.layer as i64, + module, + true, + ) + .with_context(|| { + format!( + "dynamic LoRA B is missing: adapter={adapter_id} layer={} module={module}", + slot.layer + ) + })?, + ); + // Keep one m/v entry for each A and B tensor, in slot order. + dynamic_m.push( + ctx.get_adapter_optimizer_tensor( + adapter_id, + slot.layer as i64, + module, + false, + false, + ) + .with_context(|| { + format!( + "dynamic LoRA m_a is missing: adapter={adapter_id} layer={} module={module}", + slot.layer + ) + })?, + ); + dynamic_m.push( + ctx.get_adapter_optimizer_tensor( + adapter_id, + slot.layer as i64, + module, + true, + false, + ) + .with_context(|| { + format!( + "dynamic LoRA m_b is missing: adapter={adapter_id} layer={} module={module}", + slot.layer + ) + })?, + ); + dynamic_v.push( + ctx.get_adapter_optimizer_tensor( + adapter_id, + slot.layer as i64, + module, + false, + true, + ) + .with_context(|| { + format!( + "dynamic LoRA v_a is missing: adapter={adapter_id} layer={} module={module}", + slot.layer + ) + })?, + ); + dynamic_v.push( + ctx.get_adapter_optimizer_tensor( + adapter_id, + slot.layer as i64, + module, + true, + true, + ) + .with_context(|| { + format!( + "dynamic LoRA v_b is missing: adapter={adapter_id} layer={} module={module}", + slot.layer + ) + })?, + ); + } + dynamic_adapters.push(checkpoint::DynamicAdapterCheckpoint { + manifest: checkpoint::DynamicAdapterManifest { + id: adapter_id, + rank: lora_config.rank, + alpha: lora_config.alpha, + target_layers: lora_config.target_layers.clone(), + target_modules: lora_config + .target_modules + .iter() + .map(|module| module.cpp_name().to_string()) + .collect(), + parameter_count: dynamic_a.len(), + optimizer_count: dynamic_m.len(), + }, + lora_a: dynamic_a, + lora_b: dynamic_b, + adam_m: dynamic_m, + adam_v: dynamic_v, + }); + } + } + + checkpoint::save_checkpoint_with_dynamic( std::path::Path::new(path), self.step, self.last_loss, @@ -529,6 +657,7 @@ impl TrainingSession for Qwen36Session { &lora_b, &adam_m, &adam_v, + &dynamic_adapters, )?; Ok((self.step, self.last_loss)) @@ -536,6 +665,135 @@ impl TrainingSession for Qwen36Session { fn load_checkpoint(&mut self, path: &str) -> Result<(u64, f64)> { let data = checkpoint::load_checkpoint(std::path::Path::new(path))?; + if !data.dynamic_adapters.is_empty() { + let model_path = self + .model_path + .as_ref() + .ok_or_else(|| anyhow!("model path unavailable for dynamic LoRA checkpoint"))?; + let runtime_config = rustrain_qwen3_6::config::read_qwen36_runtime_config( + std::path::Path::new(model_path), + )?; + if !self.dynamic_lora_configs.is_empty() { + bail!("cannot load dynamic LoRA checkpoint into a session with active adapters"); + } + let ctx = self + .ctx + .as_ref() + .ok_or_else(|| anyhow!("LoRA not initialized"))?; + for dynamic in &data.dynamic_adapters { + let target_modules = dynamic + .manifest + .target_modules + .iter() + .map(|name| Qwen36LoraTargetModule::parse(name)) + .collect::>>()?; + let lora_config = Qwen36LoraConfig { + rank: dynamic.manifest.rank, + alpha: dynamic.manifest.alpha, + target_layers: dynamic.manifest.target_layers.clone(), + target_modules, + }; + rustrain_qwen3_6::lora::validate_lora_targets(&runtime_config, &lora_config)?; + let layer_ids = lora_config + .target_layers + .iter() + .map(|&layer| layer as i64) + .collect::>(); + let module_csv = lora_config + .target_modules + .iter() + .map(Qwen36LoraTargetModule::cpp_name) + .collect::>() + .join(","); + let allocated_id = + ctx.add_lora(lora_config.rank, lora_config.alpha, &layer_ids, &module_csv)?; + if allocated_id != dynamic.manifest.id { + if let Err(error) = ctx.set_adapter_id(allocated_id, dynamic.manifest.id) { + let _ = ctx.remove_lora(allocated_id); + return Err(error); + } + } + let adapter_id = dynamic.manifest.id; + let load_result = (|| -> Result<()> { + let slots = + rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, &lora_config); + let mut optimizer_index = 0usize; + for (slot_index, slot) in slots.iter().filter(|slot| slot.active).enumerate() { + let module = slot.module.cpp_name(); + if slot_index >= dynamic.lora_a.len() + || slot_index >= dynamic.lora_b.len() + || optimizer_index + 1 >= dynamic.adam_m.len() + || optimizer_index + 1 >= dynamic.adam_v.len() + { + bail!( + "dynamic adapter {} tensor count mismatch", + dynamic.manifest.id + ); + } + // A/B vectors are ordered exactly like active native slots. + ctx.set_adapter_lora_tensor( + adapter_id, + slot.layer as i64, + module, + false, + &dynamic.lora_a[slot_index], + )?; + ctx.set_adapter_lora_tensor( + adapter_id, + slot.layer as i64, + module, + true, + &dynamic.lora_b[slot_index], + )?; + ctx.set_adapter_optimizer_tensor( + adapter_id, + slot.layer as i64, + module, + false, + false, + &dynamic.adam_m[optimizer_index], + )?; + ctx.set_adapter_optimizer_tensor( + adapter_id, + slot.layer as i64, + module, + true, + false, + &dynamic.adam_m[optimizer_index + 1], + )?; + ctx.set_adapter_optimizer_tensor( + adapter_id, + slot.layer as i64, + module, + false, + true, + &dynamic.adam_v[optimizer_index], + )?; + ctx.set_adapter_optimizer_tensor( + adapter_id, + slot.layer as i64, + module, + true, + true, + &dynamic.adam_v[optimizer_index + 1], + )?; + optimizer_index += 2; + } + if optimizer_index != dynamic.manifest.optimizer_count { + bail!( + "dynamic adapter {} optimizer count mismatch", + dynamic.manifest.id + ); + } + Ok(()) + })(); + if let Err(error) = load_result { + let _ = ctx.remove_lora(adapter_id); + return Err(error); + } + self.dynamic_lora_configs.insert(adapter_id, lora_config); + } + } // Import Adam optimizer state into C++ context if let Some(ctx) = &self.ctx { if data.lora_a.len() != data.lora_b.len() From ae7515264740b5a5459c4ebf712eeaa7117e8388 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 04:28:31 +0800 Subject: [PATCH 003/156] fix: make native accumulation and multi-lora updates logical-step aware --- .../kernels/qwen3_6_kernels.cpp | 269 +++++++++++------- crates/rustrain-qwen3-6/src/kernel.rs | 34 ++- crates/rustrain-qwen3-6/src/session.rs | 76 +++-- crates/rustrain-server/src/session.rs | 6 + 4 files changed, 261 insertions(+), 124 deletions(-) diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 38d2d247..446739a2 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -1527,6 +1527,7 @@ struct TrainingContext { cudaStream_t nccl_stream = nullptr; int ep_world_size = 1; int ep_rank = 0; + bool data_parallel = false; int cuda_device = 0; // ────────────────────────────────────────────────────────────────────── }; @@ -1581,10 +1582,15 @@ static ncclDataType_t nccl_dtype_for(const at::Tensor& tensor) { } } -static void allreduce_lora_grad(TrainingContext* ctx, at::Tensor& param) { +static void allreduce_lora_grad( + TrainingContext* ctx, at::Tensor& param, double local_token_scale +) { auto grad = param.grad(); if (!ctx->nccl_comm || !grad.defined()) return; auto contiguous = grad.contiguous(); + if (local_token_scale != 1.0) { + contiguous = contiguous * local_token_scale; + } auto reduced = at::empty_like(contiguous); int dev = contiguous.device().index(); cudaSetDevice(dev); @@ -1597,11 +1603,28 @@ static void allreduce_lora_grad(TrainingContext* ctx, at::Tensor& param) { param.mutable_grad() = reduced; } -// EP's routed expert output is summed across ranks in forward. The resulting -// upstream gradient must likewise be summed before replicated LoRA Adam steps; -// otherwise each rank updates attention/linear adapters from only its shard. -static void synchronize_lora_gradients(TrainingContext* ctx) { - if (!ctx->nccl_comm) return; +// Every rank evaluates the complete loss. Average replicated LoRA gradients +// across DP ranks with token-count weighting so their Adam update matches a +// single global batch. EP ranks already receive the complete routed activation +// in forward and keep replicated gradients local; routed expert adapters remain +// local because their parameter tensors are sharded. +static void synchronize_lora_gradients( + TrainingContext* ctx, const at::Tensor& target_mask +) { + if (!ctx->nccl_comm || !ctx->data_parallel) return; + auto shifted_mask = target_mask.narrow(1, 1, target_mask.size(1) - 1) + .to(at::kFloat).sum().reshape({1}); + auto global_mask = at::empty_like(shifted_mask); + auto stream = c10::cuda::getCurrentCUDAStream( + shifted_mask.device().index()).stream(); + auto err = ncclAllReduce( + shifted_mask.data_ptr(), global_mask.data_ptr(), 1, + ncclFloat, ncclSum, ctx->nccl_comm, stream); + TORCH_CHECK(err == ncclSuccess, "NCCL token-count all-reduce failed: ", + ncclGetErrorString(err)); + const double local_tokens = shifted_mask.item(); + const double global_tokens = global_mask.item(); + const double token_scale = local_tokens / std::max(global_tokens, 1.0); for (auto& adapter : ctx->adapters) { for (auto& [layer_idx, pairs] : adapter.params) { auto table = lora_projection_table(ctx->layer_configs[layer_idx]); @@ -1610,8 +1633,8 @@ static void synchronize_lora_gradients(TrainingContext* ctx) { // base experts; only replicated adapter tensors are reduced. if (table.entries[pair].grouped_expert) continue; auto& [a, b] = pairs[pair]; - allreduce_lora_grad(ctx, a); - allreduce_lora_grad(ctx, b); + allreduce_lora_grad(ctx, a, token_scale); + allreduce_lora_grad(ctx, b, token_scale); } } } @@ -1623,8 +1646,8 @@ static void synchronize_lora_gradients(TrainingContext* ctx) { // local gradients belong only to this EP rank and must not be // summed with a different expert shard on another rank. if (table.entries[pair].grouped_expert) continue; - allreduce_lora_grad(ctx, ctx->lora_a[offset + pair]); - allreduce_lora_grad(ctx, ctx->lora_b[offset + pair]); + allreduce_lora_grad(ctx, ctx->lora_a[offset + pair], token_scale); + allreduce_lora_grad(ctx, ctx->lora_b[offset + pair], token_scale); } } } @@ -3205,7 +3228,8 @@ static LossResult compute_loss( const at::Tensor& hidden, const at::Tensor& input_ids, const at::Tensor& target_mask, - int64_t vocab_size + int64_t vocab_size, + bool independent_samples = false ) { auto final_norm = *ctx->final_norm_ptr[0]; auto lm_head = *ctx->lm_head_ptr[0]; @@ -3236,6 +3260,13 @@ static LossResult compute_loss( int64_t num_chunks = (total_tokens + chunk_size - 1) / chunk_size; auto total_count = shifted_mask.sum().clamp_min(1.0); + at::Tensor token_denominators; + if (independent_samples) { + auto per_sample_count = target_mask.narrow(1, 1, seq_len - 1) + .sum(1, true).clamp_min(1.0); + token_denominators = per_sample_count + .expand({target_mask.size(0), seq_len - 1}).reshape({-1}); + } auto hidden_flat = shifted_hidden.reshape({-1, hidden_normed.size(2)}); double total_loss_val = 0.0; @@ -3276,7 +3307,9 @@ static LossResult compute_loss( // Normalize every chunk by the global response-token count. Backward // must match the mean returned to the caller, independent of sequence // length or chunk boundaries. - auto chunk_loss = masked_loss.sum() / total_count; + auto chunk_loss = independent_samples + ? (masked_loss / token_denominators.narrow(0, start, n)).sum() + : masked_loss.sum() / total_count; // Backward this chunk — each chunk creates an independent CE subgraph // because hidden_normed is a leaf tensor. retain_graph=false is safe @@ -3361,7 +3394,8 @@ static at::Tensor mtp_compute_loss( TrainingContext* ctx, const at::Tensor& mtp_hidden, const at::Tensor& input_ids, - const at::Tensor& target_mask + const at::Tensor& target_mask, + bool independent_samples = false ) { int64_t vocab_size = ctx->vocab_size; int64_t seq_len = input_ids.size(1); @@ -3380,6 +3414,13 @@ static at::Tensor mtp_compute_loss( auto total_loss = at::zeros({1}, at::TensorOptions().dtype(at::kFloat).device(mtp_hidden.device())); auto total_count = shifted_mask.sum().clamp_min(1.0); + at::Tensor token_denominators; + if (independent_samples) { + auto per_sample_count = target_mask.narrow(1, 2, n_tokens) + .sum(1, true).clamp_min(1.0); + token_denominators = per_sample_count + .expand({target_mask.size(0), n_tokens}).reshape({-1}); + } for (int64_t c = 0; c < num_chunks; c++) { int64_t start = c * chunk_size; @@ -3399,10 +3440,13 @@ static at::Tensor mtp_compute_loss( auto masked_loss = per_token_loss * chunk_mask.to(at::kFloat); // Avoid in-place accumulation into a non-grad leaf: out-of-place add // keeps the MTP loss connected to the frozen-head input graph. - total_loss = total_loss + masked_loss.sum(); + total_loss = total_loss + (independent_samples + ? (masked_loss / token_denominators.narrow(0, start, n)).sum() + : masked_loss.sum()); } - return (total_loss / total_count) * ctx->mtp_loss_scale; + if (!independent_samples) total_loss = total_loss / total_count; + return total_loss * ctx->mtp_loss_scale; } // ────────────────────────────────────────────────────────────────────── @@ -3412,7 +3456,7 @@ static at::Tensor mtp_compute_loss( extern "C" { __attribute__((visibility("default"))) int64_t qwen36_kernel_abi_version() { - return 6; + return 7; } // Create training context — called once at startup @@ -3612,16 +3656,20 @@ __attribute__((visibility("default"))) void qwen36_set_mtp_weights( (long)num_mtp_layers, (long)num_mtp_layer_weights); } -// Single training step: forward + loss + backward + Adam update -// Returns loss value, or -1 on error. -__attribute__((visibility("default"))) double qwen36_train_step( +// One training micro-step. Non-final micro-steps accumulate scaled leaf +// gradients; only the final micro-step synchronizes and updates parameters. +__attribute__((visibility("default"))) double qwen36_train_micro_step( void* ctx_ptr, void* input_ids_ptr, void* target_mask_ptr, - void* attention_mask_ptr + void* attention_mask_ptr, + double gradient_scale, + int32_t apply_optimizer ) { try { auto* ctx = reinterpret_cast(ctx_ptr); + TORCH_CHECK(gradient_scale > 0.0 && std::isfinite(gradient_scale), + "gradient_scale must be finite and positive"); // Set CUDA device for EP if (ctx->nccl_comm) { c10::cuda::set_device(ctx->cuda_device); @@ -3669,6 +3717,10 @@ __attribute__((visibility("default"))) double qwen36_train_step( loss_val += mtp_loss.item(); } + if (gradient_scale != 1.0) { + total_hidden_grad.mul_(gradient_scale); + } + // Trigger exactly one main-model backward with the combined hidden // gradient. Manual groups are the non-autograd checkpoint fallback; // normal and sub-checkpoint paths use the real graph. @@ -3678,9 +3730,13 @@ __attribute__((visibility("default"))) double qwen36_train_step( hidden.backward(total_hidden_grad); } + if (!apply_optimizer) { + return loss_val; + } + // EP produces a summed routed output, so synchronize replicated LoRA // gradients before every rank performs its local Adam update. - synchronize_lora_gradients(ctx); + synchronize_lora_gradients(ctx, target_mask); // ── Adam optimizer step — CUDA multi-tensor fused kernel ── at::AutoGradMode guard(false); @@ -3798,6 +3854,17 @@ __attribute__((visibility("default"))) double qwen36_train_step( } } +// Backward-compatible complete optimizer step. +__attribute__((visibility("default"))) double qwen36_train_step( + void* ctx_ptr, + void* input_ids_ptr, + void* target_mask_ptr, + void* attention_mask_ptr +) { + return qwen36_train_micro_step( + ctx_ptr, input_ids_ptr, target_mask_ptr, attention_mask_ptr, 1.0, 1); +} + // Get LoRA A tensor pointer by index __attribute__((visibility("default"))) void* qwen36_get_lora_a(void* ctx_ptr, int64_t index) { auto* ctx = reinterpret_cast(ctx_ptr); @@ -3997,6 +4064,16 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( double total_loss = 0.0; int64_t num_chunks = (total_adapters + n_max - 1) / n_max; + // Chunking is a memory scheduling detail, not an optimizer step. All + // adapters in this call must use the same Adam bias correction. + ctx->step_count++; + const double logical_step = (double)ctx->step_count; + const double bias_correction1 = 1.0 - std::pow(ctx->beta1, logical_step); + const double bias_correction2 = 1.0 - std::pow(ctx->beta2, logical_step); + const float lr_scaled = (float)(ctx->lr / bias_correction1); + const float eps_scaled = (float)(ctx->eps / std::sqrt(bias_correction2)); + const float one_minus_b1 = (float)(1.0 - ctx->beta1); + const float one_minus_b2 = (float)(1.0 - ctx->beta2); for (int64_t chunk = 0; chunk < num_chunks; chunk++) { int64_t start = chunk * n_max; @@ -4062,7 +4139,9 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( hidden = hidden.detach().set_requires_grad(true); TORCH_CHECK(!env_enabled("QWEN36_FUSED_CE"), "QWEN36_FUSED_CE is disabled until its tile gather and gradient normalization are validated"); - auto loss = compute_loss(ctx, hidden, input_ref, mask_ref, ctx->vocab_size); + auto loss = compute_loss( + ctx, hidden, input_ref, mask_ref, ctx->vocab_size, + /*independent_samples=*/true); loss_val = loss.value.item(); hidden_grad = loss.hidden_grad; } @@ -4074,7 +4153,9 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( if (ctx->has_mtp && !env_enabled("QWEN36_DISABLE_MTP")) { auto mtp_input = hidden.detach().set_requires_grad(true); auto mtp_hidden = mtp_forward(ctx, mtp_input, input_ref); - auto mtp_loss = mtp_compute_loss(ctx, mtp_hidden, input_ref, mask_ref); + auto mtp_loss = mtp_compute_loss( + ctx, mtp_hidden, input_ref, mask_ref, + /*independent_samples=*/true); mtp_loss.backward(); TORCH_CHECK(mtp_input.grad().defined(), "MTP did not produce a hidden gradient"); hidden_grad.add_(mtp_input.grad()); @@ -4095,88 +4176,81 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( fprintf(stderr, "[train_multi] chunk %ld/%ld: n=%ld loss=%f fwd=%.0fms loss=%.0fms bwd=%.0fms\n", (long)(chunk+1), (long)num_chunks, (long)n, loss_val, fwd_ms, loss_ms, bwd_ms); - // EP gradient synchronization must happen before this chunk's - // replicated adapter parameters are updated. - synchronize_lora_gradients(ctx); + // Restore the complete registry before the next chunk. Gradients + // remain attached to the intrusive tensor handles, so all chunks + // can accumulate and the optimizer runs exactly once below. + ctx->adapters.swap(all_adapters); - // Adam step - at::AutoGradMode guard(false); - ctx->step_count++; - ctx->lora_cache_valid = false; - ctx->lora_batch_valid = false; + if (chunk == num_chunks - 1) { + // DP gradient synchronization and Adam belong to the logical + // multi-tenant step, never to an activation-memory chunk. + synchronize_lora_gradients(ctx, target_mask); - double step_f = (double)ctx->step_count; - double bias_correction1 = 1.0 - std::pow(ctx->beta1, step_f); - double bias_correction2 = 1.0 - std::pow(ctx->beta2, step_f); - float lr_scaled = (float)(ctx->lr / bias_correction1); - float eps_scaled = (float)(ctx->eps / std::sqrt(bias_correction2)); - float one_minus_b1 = (float)(1.0 - ctx->beta1); - float one_minus_b2 = (float)(1.0 - ctx->beta2); + // Adam step + at::AutoGradMode guard(false); + ctx->lora_cache_valid = false; + ctx->lora_batch_valid = false; - std::vector h_params, h_grads; - std::vector h_m, h_v; - std::vector h_sizes; + std::vector h_params, h_grads; + std::vector h_m, h_v; + std::vector h_sizes; - for (auto& adapter : ctx->adapters) { - for (auto& [layer_idx, pairs] : adapter.params) { - auto& adam_states = adapter.adam_state[layer_idx]; - for (size_t i = 0; i < pairs.size(); i++) { - auto& [a, b] = pairs[i]; - auto& [m_a, v_a, m_b, v_b] = adam_states[i]; - if (a.grad().defined() && a.scalar_type() == at::kBFloat16) { - h_params.push_back(a.data_ptr()); - h_grads.push_back(a.grad().data_ptr()); - h_m.push_back((float*)m_a.data_ptr()); - h_v.push_back((float*)v_a.data_ptr()); - h_sizes.push_back((int)a.numel()); - } - if (b.grad().defined() && b.scalar_type() == at::kBFloat16) { - h_params.push_back(b.data_ptr()); - h_grads.push_back(b.grad().data_ptr()); - h_m.push_back((float*)m_b.data_ptr()); - h_v.push_back((float*)v_b.data_ptr()); - h_sizes.push_back((int)b.numel()); + for (auto& adapter : ctx->adapters) { + for (auto& [layer_idx, pairs] : adapter.params) { + auto& adam_states = adapter.adam_state[layer_idx]; + for (size_t i = 0; i < pairs.size(); i++) { + auto& [a, b] = pairs[i]; + auto& [m_a, v_a, m_b, v_b] = adam_states[i]; + if (a.grad().defined() && a.scalar_type() == at::kBFloat16) { + h_params.push_back(a.data_ptr()); + h_grads.push_back(a.grad().data_ptr()); + h_m.push_back((float*)m_a.data_ptr()); + h_v.push_back((float*)v_a.data_ptr()); + h_sizes.push_back((int)a.numel()); + } + if (b.grad().defined() && b.scalar_type() == at::kBFloat16) { + h_params.push_back(b.data_ptr()); + h_grads.push_back(b.grad().data_ptr()); + h_m.push_back((float*)m_b.data_ptr()); + h_v.push_back((float*)v_b.data_ptr()); + h_sizes.push_back((int)b.numel()); + } } } } - } - if (!h_params.empty()) { - int n_params = (int)h_params.size(); - auto opts_cpu_long = at::TensorOptions().dtype(at::kLong).device(at::kCPU); - auto opts_cpu_int = at::TensorOptions().dtype(at::kInt).device(at::kCPU); - auto params_cpu = at::from_blob(h_params.data(), {n_params}, opts_cpu_long); - auto grads_cpu = at::from_blob(h_grads.data(), {n_params}, opts_cpu_long); - auto m_cpu = at::from_blob(h_m.data(), {n_params}, opts_cpu_long); - auto v_cpu = at::from_blob(h_v.data(), {n_params}, opts_cpu_long); - auto sizes_cpu = at::from_blob(h_sizes.data(), {n_params}, opts_cpu_int); - ctx->adam_dev_bufs.ensure(n_params, ctx->adapters[0].params.begin()->second[0].first); - ctx->adam_dev_bufs.params_buf.narrow(0, 0, n_params).copy_(params_cpu); - ctx->adam_dev_bufs.grads_buf.narrow(0, 0, n_params).copy_(grads_cpu); - ctx->adam_dev_bufs.m_buf.narrow(0, 0, n_params).copy_(m_cpu); - ctx->adam_dev_bufs.v_buf.narrow(0, 0, n_params).copy_(v_cpu); - ctx->adam_dev_bufs.sizes_buf.narrow(0, 0, n_params).copy_(sizes_cpu); - - auto stream = c10::cuda::getCurrentCUDAStream().stream(); - launch_fused_adam_multi( - (void**)ctx->adam_dev_bufs.params_buf.data_ptr(), - (void**)ctx->adam_dev_bufs.grads_buf.data_ptr(), - (float**)ctx->adam_dev_bufs.m_buf.data_ptr(), - (float**)ctx->adam_dev_bufs.v_buf.data_ptr(), - (int*)ctx->adam_dev_bufs.sizes_buf.data_ptr(), - n_params, - (float)ctx->beta1, (float)ctx->beta2, - lr_scaled, eps_scaled, - one_minus_b1, one_minus_b2, - (void*)stream - ); + if (!h_params.empty()) { + int n_params = (int)h_params.size(); + auto opts_cpu_long = at::TensorOptions().dtype(at::kLong).device(at::kCPU); + auto opts_cpu_int = at::TensorOptions().dtype(at::kInt).device(at::kCPU); + auto params_cpu = at::from_blob(h_params.data(), {n_params}, opts_cpu_long); + auto grads_cpu = at::from_blob(h_grads.data(), {n_params}, opts_cpu_long); + auto m_cpu = at::from_blob(h_m.data(), {n_params}, opts_cpu_long); + auto v_cpu = at::from_blob(h_v.data(), {n_params}, opts_cpu_long); + auto sizes_cpu = at::from_blob(h_sizes.data(), {n_params}, opts_cpu_int); + ctx->adam_dev_bufs.ensure(n_params, ctx->adapters[0].params.begin()->second[0].first); + ctx->adam_dev_bufs.params_buf.narrow(0, 0, n_params).copy_(params_cpu); + ctx->adam_dev_bufs.grads_buf.narrow(0, 0, n_params).copy_(grads_cpu); + ctx->adam_dev_bufs.m_buf.narrow(0, 0, n_params).copy_(m_cpu); + ctx->adam_dev_bufs.v_buf.narrow(0, 0, n_params).copy_(v_cpu); + ctx->adam_dev_bufs.sizes_buf.narrow(0, 0, n_params).copy_(sizes_cpu); + + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + launch_fused_adam_multi( + (void**)ctx->adam_dev_bufs.params_buf.data_ptr(), + (void**)ctx->adam_dev_bufs.grads_buf.data_ptr(), + (float**)ctx->adam_dev_bufs.m_buf.data_ptr(), + (float**)ctx->adam_dev_bufs.v_buf.data_ptr(), + (int*)ctx->adam_dev_bufs.sizes_buf.data_ptr(), + n_params, + (float)ctx->beta1, (float)ctx->beta2, + lr_scaled, eps_scaled, + one_minus_b1, one_minus_b2, + (void*)stream + ); + } } - // Restore the complete registry. Adapter parameter and Adam-state - // tensors are intrusive handles, so the chunk copies above share - // the updated storage with all_adapters; no value merge is needed. - ctx->adapters.swap(all_adapters); - total_loss += loss_val; fprintf(stderr, "[train_multi] chunk %ld/%ld: n=%ld loss=%.6f\n", @@ -4184,7 +4258,7 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( } ctx->attention_mask = saved_attention_mask; - return total_loss / num_chunks; + return total_loss / total_adapters; } catch (const std::exception& e) { fprintf(stderr, "[train_multi] FAILED: %s\n", e.what()); return -1.0; @@ -4229,6 +4303,7 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( if (g_nccl_initialized) { ctx->nccl_comm = g_nccl_comm; ctx->nccl_stream = g_nccl_stream; + ctx->data_parallel = env_enabled("RUSTRAIN_DATA_PARALLEL"); // CRITICAL: also set ep_rank/ep_world_size — needed for new TrainingContext // created by subsequent CreateSession commands. The fast path previously // skipped this, leaving ep_rank=0 → cudaSetDevice(0) on all ranks → crash. @@ -4352,6 +4427,7 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( ctx->nccl_stream = nccl_stream; ctx->ep_rank = rank; ctx->ep_world_size = world_size; + ctx->data_parallel = env_enabled("RUSTRAIN_DATA_PARALLEL"); // Propagate to layer configs for (auto& lc : ctx->layer_configs) { @@ -4376,6 +4452,7 @@ __attribute__((visibility("default"))) void qwen36_set_nccl_comm( ctx->nccl_stream = reinterpret_cast(stream_ptr); ctx->ep_rank = ep_rank; ctx->ep_world_size = ep_world_size; + ctx->data_parallel = env_enabled("RUSTRAIN_DATA_PARALLEL"); int current_device = g_cuda_device; cudaGetDevice(¤t_device); ctx->cuda_device = current_device; diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs index 1554c23a..871b2830 100644 --- a/crates/rustrain-qwen3-6/src/kernel.rs +++ b/crates/rustrain-qwen3-6/src/kernel.rs @@ -4,7 +4,7 @@ //! Rust only handles: weight loading, data loading, training loop orchestration. use crate::lora::Qwen36LoraTargetModule; -use anyhow::{bail, Result}; +use anyhow::{Result, bail}; use std::ffi::c_void; use std::sync::OnceLock; use tch::{Kind, Tensor}; @@ -34,6 +34,8 @@ type FnCreateCtx = unsafe extern "C" fn( ) -> *mut c_void; type FnKernelAbiVersion = unsafe extern "C" fn() -> i64; type FnTrainStep = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void) -> f64; +type FnTrainMicroStep = + unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void, f64, i32) -> f64; type FnTrainMultiLora = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void, i32, i32) -> f64; type FnEvalStep = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void) -> f64; @@ -106,6 +108,7 @@ pub struct CppLayerConfig { struct KernelHandles { create_ctx: FnCreateCtx, train_step: FnTrainStep, + train_micro_step: FnTrainMicroStep, train_multi_lora: FnTrainMultiLora, eval_step: FnEvalStep, get_lora_count: FnGetLoraCount, @@ -171,12 +174,13 @@ unsafe fn load_kernels() -> Option { }}; } let abi_version: FnKernelAbiVersion = sym!("qwen36_kernel_abi_version"); - if abi_version() != 6 { + if abi_version() != 7 { return None; } Some(KernelHandles { create_ctx: sym!("qwen36_create_training_context"), train_step: sym!("qwen36_train_step"), + train_micro_step: sym!("qwen36_train_micro_step"), train_multi_lora: sym!("qwen36_train_multi_lora"), eval_step: sym!("qwen36_eval_step"), get_lora_count: sym!("qwen36_get_lora_count"), @@ -545,6 +549,32 @@ impl CppTrainingContext { Ok(loss) } + /// Run one micro-batch and optionally apply the synchronized Adam update. + pub fn train_micro_step( + &self, + input_ids: &Tensor, + target_mask: &Tensor, + attention_mask: &Tensor, + gradient_scale: f64, + apply_optimizer: bool, + ) -> Result { + let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; + let loss = unsafe { + (kh.train_micro_step)( + self.ptr, + input_ids.as_ptr() as *mut _, + target_mask.as_ptr() as *mut _, + attention_mask.as_ptr() as *mut _, + gradient_scale, + i32::from(apply_optimizer), + ) + }; + if loss < 0.0 { + bail!("C++ train_micro_step failed"); + } + Ok(loss) + } + /// Train all adapters in batched chunks. Each chunk runs independent /// forward → loss → backward → Adam. Input is expanded to [N, seq]. /// n_total: total number of adapters. lora_rank: LoRA rank for N_max calc. diff --git a/crates/rustrain-qwen3-6/src/session.rs b/crates/rustrain-qwen3-6/src/session.rs index 3fb07321..cc23a2ba 100644 --- a/crates/rustrain-qwen3-6/src/session.rs +++ b/crates/rustrain-qwen3-6/src/session.rs @@ -264,7 +264,9 @@ fn train_impl( let rank = shard_ref.map(|s| s.rank).unwrap_or(env_rank); let is_data_parallel = !is_ep && world_size > 1; if is_data_parallel && runtime_config.is_moe { - bail!("replicated Qwen data parallelism is only supported for dense/linear-attention models; use *_ep for MoE"); + bail!( + "replicated Qwen data parallelism is only supported for dense/linear-attention models; use *_ep for MoE" + ); } if is_data_parallel { crate::kernel::CppTrainingContext::set_cuda_device( @@ -355,6 +357,7 @@ fn train_impl( let batch_size = config.train.micro_batch_size; let max_steps = config.train.max_steps as usize; + let gradient_accumulation_steps = config.train.gradient_accumulation_steps; // ── C++ all-in-C++ training path (required) ── // LoRA A/B, Adam optimizer, forward, loss, backward all in C++. @@ -381,11 +384,23 @@ fn train_impl( )?; if world_size > 1 { + // Tell the native reducer whether this communicator is replicated DP + // or expert-parallel. EP already all-reduces routed activations and + // must not average replicated LoRA gradients a second time. + unsafe { + std::env::set_var( + "RUSTRAIN_DATA_PARALLEL", + if is_data_parallel { "1" } else { "0" }, + ); + } let ret = ctx.init_nccl(); if ret != 0 { bail!("C++ NCCL init failed (code {})", ret); } - info!(rank, world_size, is_ep, "NCCL communicator created for Qwen LoRA parallel training"); + info!( + rank, + world_size, is_ep, "NCCL communicator created for Qwen LoRA parallel training" + ); } info!("C++ TrainingContext: {} LoRA params", ctx.lora_count()); @@ -415,30 +430,39 @@ fn train_impl( let mut final_loss = 0.0_f64; for step in 0..max_steps { - let data_start = if is_data_parallel { - (step * batch_size * world_size + rank * batch_size) % data.len() - } else { - (step * batch_size) % data.len() - }; - let sft_batch = data.batch(data_start, batch_size); - let (input_ids, target_mask) = sft_batch.to_tensors(device, compute_kind); - - // Build attention mask: 1 for real tokens, 0 for padding - // target_mask > 0 means the token is either response (loss) or prompt (no loss but attend) - // We need: 1 for all non-padding tokens, 0 for padding tokens - // The SFT batch's target_mask is: 0=prompt, 1=response, but padding has mask=0 too. - // We need attention_mask = (target_mask >= 0).to(float) but that's all 1s. - // Actually: padding tokens have target_mask=0 AND are after EOS. - // The SftBatch already has padding at the end with mask=0. - // We need: attention_mask = 1 where token is NOT padding. - // Since prompt tokens have mask=0 and response tokens have mask=1, - // but padding also has mask=0, we can't distinguish prompt from padding using mask alone. - // Solution: use the pad_token_id to build attention mask from input_ids. - let pad_id = data.pad_token_id(); - let attention_mask = input_ids.ne(pad_id).to_kind(Kind::Float).unsqueeze(0); // [1, seq] - - // C++ all-in-C++ path: single call does forward + loss + backward + Adam - let loss_value = ctx.train_step(&input_ids, &target_mask, &attention_mask)?; + let mut loss_value = 0.0; + for accumulation_index in 0..gradient_accumulation_steps { + let micro_step = step * gradient_accumulation_steps + accumulation_index; + let data_start = if is_data_parallel { + (micro_step * batch_size * world_size + rank * batch_size) % data.len() + } else { + (micro_step * batch_size) % data.len() + }; + let sft_batch = data.batch(data_start, batch_size); + let (input_ids, target_mask) = sft_batch.to_tensors(device, compute_kind); + + // Build attention mask: 1 for real tokens, 0 for padding + // target_mask > 0 means the token is either response (loss) or prompt (no loss but attend) + // We need: 1 for all non-padding tokens, 0 for padding tokens + // The SFT batch's target_mask is: 0=prompt, 1=response, but padding has mask=0 too. + // We need attention_mask = (target_mask >= 0).to(float) but that's all 1s. + // Actually: padding tokens have target_mask=0 AND are after EOS. + // The SftBatch already has padding at the end with mask=0. + // We need: attention_mask = 1 where token is NOT padding. + // Since prompt tokens have mask=0 and response tokens have mask=1, + // but padding also has mask=0, we can't distinguish prompt from padding using mask alone. + // Solution: use the pad_token_id to build attention mask from input_ids. + let pad_id = data.pad_token_id(); + let attention_mask = input_ids.ne(pad_id).to_kind(Kind::Float).unsqueeze(0); // [1, seq] + + loss_value += ctx.train_micro_step( + &input_ids, + &target_mask, + &attention_mask, + 1.0 / gradient_accumulation_steps as f64, + accumulation_index + 1 == gradient_accumulation_steps, + )? / gradient_accumulation_steps as f64; + } if step == 0 { initial_loss = loss_value; } diff --git a/crates/rustrain-server/src/session.rs b/crates/rustrain-server/src/session.rs index 0fddea94..8db56469 100644 --- a/crates/rustrain-server/src/session.rs +++ b/crates/rustrain-server/src/session.rs @@ -399,6 +399,12 @@ impl TrainingSession for Qwen36Session { // Initialize NCCL directly in C++. The same communicator handles EP // output collectives and replicated-weight LoRA gradient all-reduce. let nccl_ep = if is_ep || is_data_parallel { + unsafe { + std::env::set_var( + "RUSTRAIN_DATA_PARALLEL", + if is_data_parallel { "1" } else { "0" }, + ); + } let ret = ctx.init_nccl(); if ret != 0 { return Err(anyhow!("C++ NCCL init failed (code {})", ret)); From 0794e6a01d8c7e5b10c808be5eddf56c94957423 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 04:32:41 +0800 Subject: [PATCH 004/156] test: cover native micro-step optimizer boundary --- crates/rustrain-qwen3-6/tests/native_smoke.cpp | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index ddebdfbb..89423b4f 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -29,6 +29,9 @@ extern "C" void* qwen36_get_lora_a(void*, int64_t); extern "C" void* qwen36_get_lora_b(void*, int64_t); extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*); extern "C" double qwen36_train_step(void*, void*, void*, void*); +extern "C" double qwen36_train_micro_step( + void*, void*, void*, void*, double, int32_t); +extern "C" int64_t qwen36_get_step_count(void*); extern "C" double qwen36_eval_step(void*, void*, void*, void*); extern "C" double qwen36_train_multi_lora( void*, void*, void*, void*, int32_t, int32_t); @@ -352,6 +355,20 @@ int main() { auto linear_b_value = at::ones(linear_b->sizes(), linear_b->options()); assert(qwen36_set_lora_tensor(ctx, 0, 1, &linear_b_value) == 0); auto linear_a_before = linear_a->clone(); + assert(qwen36_get_step_count(ctx) == 0); + const double accum_loss_0 = qwen36_train_micro_step( + ctx, &input_ids, &target_mask, &attention_mask, 0.5, 0); + c10::cuda::device_synchronize(); + assert(accum_loss_0 == accum_loss_0); + assert(qwen36_get_step_count(ctx) == 0); + assert((*linear_a - linear_a_before).abs().sum().item() == 0.0); + const double accum_loss_1 = qwen36_train_micro_step( + ctx, &input_ids, &target_mask, &attention_mask, 0.5, 1); + c10::cuda::device_synchronize(); + assert(accum_loss_1 == accum_loss_1); + assert(qwen36_get_step_count(ctx) == 1); + assert((*linear_a - linear_a_before).abs().sum().item() > 0.0); + linear_a_before = linear_a->clone(); const double linear_loss = qwen36_train_step( ctx, &input_ids, &target_mask, &attention_mask); c10::cuda::device_synchronize(); From 0530e79391b29af32c2fe463a729fec8fa6d08ac Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 04:33:27 +0800 Subject: [PATCH 005/156] fix: invalidate native lora graph between microbatches --- crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 446739a2..3f6594ad 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -3731,6 +3731,11 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( } if (!apply_optimizer) { + // The forward graph has been consumed, but parameters remain + // live for the next micro-batch. Never reuse cached LoRA deltas + // whose autograd nodes were freed by this backward. + ctx->lora_cache_valid = false; + ctx->lora_batch_valid = false; return loss_val; } From f65af15e1c590b5dff053908e8c31b038fbf7129 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 04:34:09 +0800 Subject: [PATCH 006/156] docs: clarify native loss and parallel reduction semantics --- crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 3f6594ad..896f32ac 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -3304,9 +3304,9 @@ static LossResult compute_loss( at::Tensor(), at::Reduction::None, -100, 0.0 ); auto masked_loss = per_token_loss * chunk_mask.to(at::kFloat); - // Normalize every chunk by the global response-token count. Backward - // must match the mean returned to the caller, independent of sequence - // length or chunk boundaries. + // Single-sample training uses the global response-token mean. Batched + // multi-LoRA instead divides each row by its own response-token count + // so tenant gradients do not depend on neighboring rows. auto chunk_loss = independent_samples ? (masked_loss / token_denominators.narrow(0, start, n)).sum() : masked_loss.sum() / total_count; @@ -3739,8 +3739,9 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( return loss_val; } - // EP produces a summed routed output, so synchronize replicated LoRA - // gradients before every rank performs its local Adam update. + // Replicated DP gradients are synchronized before the local Adam + // update. EP keeps replicated gradients local because its forward + // routed activation already contains the cross-rank sum. synchronize_lora_gradients(ctx, target_mask); // ── Adam optimizer step — CUDA multi-tensor fused kernel ── From 538cb5e5efc3737959bdfd05d2b752b5d7f8db1d Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 04:36:43 +0800 Subject: [PATCH 007/156] docs: describe multi-lora logical update --- crates/rustrain-qwen3-6/src/kernel.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs index 871b2830..ac8c6600 100644 --- a/crates/rustrain-qwen3-6/src/kernel.rs +++ b/crates/rustrain-qwen3-6/src/kernel.rs @@ -575,8 +575,8 @@ impl CppTrainingContext { Ok(loss) } - /// Train all adapters in batched chunks. Each chunk runs independent - /// forward → loss → backward → Adam. Input is expanded to [N, seq]. + /// Train all adapters in batched activation chunks. Chunks accumulate + /// gradients and share one logical Adam update. Input is expanded to [N, seq]. /// n_total: total number of adapters. lora_rank: LoRA rank for N_max calc. /// Returns average loss across chunks. pub fn train_multi_lora( From 4e242ee08c0386df0945f2de416822a5e020260e Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 04:44:43 +0800 Subject: [PATCH 008/156] feat: add Megatron-style parallel topology --- crates/rustrain-parallel/src/launcher.rs | 82 +++- crates/rustrain-parallel/src/lib.rs | 1 + crates/rustrain-parallel/src/topology.rs | 517 +++++++++++++++++++++++ 3 files changed, 597 insertions(+), 3 deletions(-) create mode 100644 crates/rustrain-parallel/src/topology.rs diff --git a/crates/rustrain-parallel/src/launcher.rs b/crates/rustrain-parallel/src/launcher.rs index 311002bd..e0b8530c 100644 --- a/crates/rustrain-parallel/src/launcher.rs +++ b/crates/rustrain-parallel/src/launcher.rs @@ -9,6 +9,8 @@ use std::{ use anyhow::{Context, Result, bail}; use serde::{Deserialize, Serialize}; +use crate::topology::ParallelTopology; + #[derive(Debug, Serialize)] struct LaunchSummary { nproc_per_node: usize, @@ -40,6 +42,12 @@ pub struct LaunchEnvSummary { pub cuda_visible_devices: Option, pub assigned_cuda_visible_device: Option, pub assigned_cuda_device_ordinal: Option, + pub tensor_model_parallel_size: usize, + pub pipeline_model_parallel_size: usize, + pub data_parallel_size: usize, + pub expert_model_parallel_size: usize, + pub context_parallel_size: usize, + pub parallel_rank_order: String, } pub fn launch( @@ -49,7 +57,15 @@ pub fn launch( master_port: u16, command: &[String], ) -> Result<()> { - launch_multi(nproc_per_node, 1, 0, output_dir, master_addr, master_port, command) + launch_multi( + nproc_per_node, + 1, + 0, + output_dir, + master_addr, + master_port, + command, + ) } pub fn launch_multi( @@ -81,6 +97,7 @@ pub fn launch_multi( .with_context(|| format!("failed to create {}", output_dir.display()))?; let current_exe = std::env::current_exe().context("failed to locate current executable")?; let timeout = launch_timeout()?; + let topology = ParallelTopology::from_env_with_world_size(world_size)?; let visible_cuda_devices = parse_visible_cuda_devices(std::env::var("CUDA_VISIBLE_DEVICES").ok()); validate_visible_cuda_devices(nproc_per_node, visible_cuda_devices.as_deref())?; @@ -107,6 +124,47 @@ pub fn launch_multi( .env("RUSTRAIN_LAUNCH_OUTPUT_DIR", output_dir) .env("NNODES", nnodes.to_string()) .env("NODE_RANK", node_rank.to_string()) + // Normalize topology variables for every child. Explicit + // TP_SIZE/PP_SIZE/... values come from the parent environment; + // unspecified axes default to replicated DP and are validated + // against the launcher's world size above. + .env("TP_SIZE", topology.tensor_model_parallel_size().to_string()) + .env( + "PP_SIZE", + topology.pipeline_model_parallel_size().to_string(), + ) + .env("DP_SIZE", topology.data_parallel_size().to_string()) + .env("EP_SIZE", topology.expert_model_parallel_size().to_string()) + .env("CP_SIZE", topology.context_parallel_size().to_string()) + .env( + "RUSTRAIN_TP_SIZE", + topology.tensor_model_parallel_size().to_string(), + ) + .env( + "RUSTRAIN_PP_SIZE", + topology.pipeline_model_parallel_size().to_string(), + ) + .env( + "RUSTRAIN_DP_SIZE", + topology.data_parallel_size().to_string(), + ) + .env( + "RUSTRAIN_EP_SIZE", + topology.expert_model_parallel_size().to_string(), + ) + .env( + "RUSTRAIN_CP_SIZE", + topology.context_parallel_size().to_string(), + ) + .env( + "RUSTRAIN_PARALLEL_ORDER", + topology + .order() + .into_iter() + .map(|axis| axis.name()) + .collect::>() + .join("-"), + ) .env( "PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True,max_split_size_mb:512", @@ -119,13 +177,19 @@ pub fn launch_multi( { child .env("RUSTRAIN_ASSIGNED_CUDA_VISIBLE_DEVICE", assigned_device) - .env("RUSTRAIN_ASSIGNED_CUDA_DEVICE_ORDINAL", local_rank.to_string()); + .env( + "RUSTRAIN_ASSIGNED_CUDA_DEVICE_ORDINAL", + local_rank.to_string(), + ); } children.push(( global_rank, log_path, child.spawn().with_context(|| { - format!("failed to spawn rank {global_rank} for command {:?}", command) + format!( + "failed to spawn rank {global_rank} for command {:?}", + command + ) })?, )); } @@ -188,6 +252,7 @@ pub fn print_launch_env() -> Result<()> { } fn read_launch_env() -> Result { + let topology = ParallelTopology::from_env()?; Ok(LaunchEnvSummary { rank: parse_env_usize("RANK")?, local_rank: parse_env_usize("LOCAL_RANK")?, @@ -204,6 +269,17 @@ fn read_launch_env() -> Result { .with_context(|| "RUSTRAIN_ASSIGNED_CUDA_DEVICE_ORDINAL must be a usize") }) .transpose()?, + tensor_model_parallel_size: topology.tensor_model_parallel_size(), + pipeline_model_parallel_size: topology.pipeline_model_parallel_size(), + data_parallel_size: topology.data_parallel_size(), + expert_model_parallel_size: topology.expert_model_parallel_size(), + context_parallel_size: topology.context_parallel_size(), + parallel_rank_order: topology + .order() + .into_iter() + .map(|axis| axis.name()) + .collect::>() + .join("-"), }) } diff --git a/crates/rustrain-parallel/src/lib.rs b/crates/rustrain-parallel/src/lib.rs index ad211184..e21308c5 100644 --- a/crates/rustrain-parallel/src/lib.rs +++ b/crates/rustrain-parallel/src/lib.rs @@ -2,3 +2,4 @@ pub mod dp_rank; pub mod launcher; pub mod parallel; pub mod parallel_modules; +pub mod topology; diff --git a/crates/rustrain-parallel/src/topology.rs b/crates/rustrain-parallel/src/topology.rs new file mode 100644 index 00000000..b0e59137 --- /dev/null +++ b/crates/rustrain-parallel/src/topology.rs @@ -0,0 +1,517 @@ +//! Rank topology for Megatron-style orthogonal parallel groups. +//! +//! This module only describes process topology. It deliberately does not +//! create NCCL process groups or shard model weights; those operations belong +//! to the model/runtime that owns each collective. Keeping the rank mapping in +//! one place prevents TP/PP/DP/EP/CP launchers from silently disagreeing. + +use std::{convert::TryInto, env, ops::Range}; + +use anyhow::{Context, Result, anyhow, bail}; +use serde::{Deserialize, Serialize}; + +use rustrain_core::runtime::ParallelConfig; + +/// Megatron's current default rank order for decoder groups. +pub const DEFAULT_RANK_ORDER: &str = "tp-cp-ep-dp-pp"; + +/// One orthogonal parallelism axis. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ParallelAxis { + Tensor, + Pipeline, + Data, + Expert, + Context, +} + +impl ParallelAxis { + pub const ALL: [Self; 5] = [ + Self::Tensor, + Self::Context, + Self::Expert, + Self::Data, + Self::Pipeline, + ]; + + pub const fn name(self) -> &'static str { + match self { + Self::Tensor => "tp", + Self::Pipeline => "pp", + Self::Data => "dp", + Self::Expert => "ep", + Self::Context => "cp", + } + } + + fn parse(raw: &str) -> Result { + match raw.trim().to_ascii_lowercase().as_str() { + "tp" | "tensor" => Ok(Self::Tensor), + "pp" | "pipeline" => Ok(Self::Pipeline), + "dp" | "data" => Ok(Self::Data), + "ep" | "expert" => Ok(Self::Expert), + "cp" | "context" => Ok(Self::Context), + other => bail!("unknown parallel axis '{other}' (expected tp/pp/dp/ep/cp)"), + } + } +} + +/// Coordinates of one global rank in the five-dimensional topology. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct RankCoordinates { + pub tensor: usize, + pub pipeline: usize, + pub data: usize, + pub expert: usize, + pub context: usize, +} + +impl RankCoordinates { + pub const ZERO: Self = Self { + tensor: 0, + pipeline: 0, + data: 0, + expert: 0, + context: 0, + }; + + pub const fn get(self, axis: ParallelAxis) -> usize { + match axis { + ParallelAxis::Tensor => self.tensor, + ParallelAxis::Pipeline => self.pipeline, + ParallelAxis::Data => self.data, + ParallelAxis::Expert => self.expert, + ParallelAxis::Context => self.context, + } + } + + pub fn set(&mut self, axis: ParallelAxis, value: usize) { + match axis { + ParallelAxis::Tensor => self.tensor = value, + ParallelAxis::Pipeline => self.pipeline = value, + ParallelAxis::Data => self.data = value, + ParallelAxis::Expert => self.expert = value, + ParallelAxis::Context => self.context = value, + } + } +} + +/// A validated orthogonal TP/PP/DP/EP/CP rank topology. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ParallelTopology { + sizes: RankCoordinates, + /// The first axis is the least-significant rank digit, matching + /// Megatron's `generate_masked_orthogonal_rank_groups` convention. + order: [ParallelAxis; 5], +} + +impl ParallelTopology { + /// Construct a topology using the Megatron default order. + pub fn new( + tensor: usize, + pipeline: usize, + data: usize, + expert: usize, + context: usize, + ) -> Result { + Self::with_order(tensor, pipeline, data, expert, context, DEFAULT_RANK_ORDER) + } + + /// Construct a topology with an explicit least-significant-first order. + pub fn with_order( + tensor: usize, + pipeline: usize, + data: usize, + expert: usize, + context: usize, + order: &str, + ) -> Result { + let sizes = RankCoordinates { + tensor, + pipeline, + data, + expert, + context, + }; + for axis in ParallelAxis::ALL { + if sizes.get(axis) == 0 { + bail!("{} parallel size must be greater than zero", axis.name()); + } + } + + let parsed: Vec = order + .split('-') + .filter(|token| !token.trim().is_empty()) + .map(ParallelAxis::parse) + .collect::>()?; + if parsed.is_empty() || parsed.len() > ParallelAxis::ALL.len() { + bail!("parallel rank order must list each axis at most once (got '{order}')"); + } + for axis in ParallelAxis::ALL { + let count = parsed + .iter() + .filter(|candidate| **candidate == axis) + .count(); + if count > 1 { + bail!( + "parallel rank order must contain {} at most once", + axis.name() + ); + } + if count == 0 && sizes.get(axis) != 1 { + bail!( + "parallel rank order omits non-singleton {} axis (size {})", + axis.name(), + sizes.get(axis) + ); + } + } + let mut order = parsed; + for axis in ParallelAxis::ALL { + if !order.contains(&axis) { + order.push(axis); + } + } + let order: [ParallelAxis; 5] = order + .try_into() + .map_err(|_| anyhow!("parallel rank order must contain five axes"))?; + Ok(Self { sizes, order }) + } + + /// Construct from the public runtime configuration and expected world size. + pub fn from_config(config: &ParallelConfig, world_size: usize) -> Result { + let topology = Self::new( + config.tensor_model_parallel_size, + config.pipeline_model_parallel_size, + config.data_parallel_size, + config.expert_model_parallel_size, + config.context_parallel_size, + )?; + topology.validate_world_size(world_size)?; + Ok(topology) + } + + /// Read axis sizes from launcher environment variables. + /// + /// `TP_SIZE`, `PP_SIZE`, `DP_SIZE`, `EP_SIZE`, and `CP_SIZE` are accepted; + /// each also has a `RUSTRAIN_`-prefixed alias. When no axis is specified, + /// all ranks are treated as data-parallel replicas. If DP is omitted while + /// other axes are specified, it is inferred from `WORLD_SIZE`. + pub fn from_env() -> Result { + let world_size = parse_env_usize("WORLD_SIZE")?; + Self::from_env_with_world_size(world_size) + } + + /// Read axis sizes from the launcher environment while taking world size + /// from the caller. This is used by the launcher before it spawns ranks, + /// when `WORLD_SIZE` is not yet present in the parent process. + pub fn from_env_with_world_size(world_size: usize) -> Result { + if world_size == 0 { + bail!("WORLD_SIZE must be greater than zero"); + } + let mut values = [None; 5]; + for (index, names) in [ + ["TP_SIZE", "RUSTRAIN_TP_SIZE", "TENSOR_MODEL_PARALLEL_SIZE"], + [ + "PP_SIZE", + "RUSTRAIN_PP_SIZE", + "PIPELINE_MODEL_PARALLEL_SIZE", + ], + ["DP_SIZE", "RUSTRAIN_DP_SIZE", "DATA_PARALLEL_SIZE"], + ["EP_SIZE", "RUSTRAIN_EP_SIZE", "EXPERT_MODEL_PARALLEL_SIZE"], + ["CP_SIZE", "RUSTRAIN_CP_SIZE", "CONTEXT_PARALLEL_SIZE"], + ] + .into_iter() + .enumerate() + { + values[index] = first_env_usize(&names)?; + } + + let any_explicit = values.iter().any(Option::is_some); + let tensor = values[0].unwrap_or(1); + let pipeline = values[1].unwrap_or(1); + let expert = values[3].unwrap_or(1); + let context = values[4].unwrap_or(1); + let data = match values[2] { + Some(value) => value, + None if !any_explicit => world_size, + None => { + let model_parallel = tensor + .checked_mul(pipeline) + .and_then(|value| value.checked_mul(expert)) + .and_then(|value| value.checked_mul(context)) + .ok_or_else(|| anyhow!("parallel axis product overflowed usize"))?; + if model_parallel == 0 || world_size % model_parallel != 0 { + bail!( + "WORLD_SIZE={world_size} is not divisible by specified model-parallel product {model_parallel}; set DP_SIZE explicitly" + ); + } + world_size / model_parallel + } + }; + let order = env::var("RUSTRAIN_PARALLEL_ORDER") + .or_else(|_| env::var("PARALLEL_ORDER")) + .unwrap_or_else(|_| DEFAULT_RANK_ORDER.to_string()); + let topology = Self::with_order(tensor, pipeline, data, expert, context, &order)?; + topology.validate_world_size(world_size)?; + Ok(topology) + } + + pub const fn order(&self) -> [ParallelAxis; 5] { + self.order + } + + pub const fn sizes(&self) -> RankCoordinates { + self.sizes + } + + pub const fn tensor_model_parallel_size(&self) -> usize { + self.sizes.tensor + } + + pub const fn pipeline_model_parallel_size(&self) -> usize { + self.sizes.pipeline + } + + pub const fn data_parallel_size(&self) -> usize { + self.sizes.data + } + + pub const fn expert_model_parallel_size(&self) -> usize { + self.sizes.expert + } + + pub const fn context_parallel_size(&self) -> usize { + self.sizes.context + } + + pub fn world_size(&self) -> usize { + self.order + .into_iter() + .map(|axis| self.sizes.get(axis)) + .product() + } + + pub fn validate_world_size(&self, world_size: usize) -> Result<()> { + let expected = self.world_size(); + if world_size != expected { + bail!( + "parallel topology expects world_size={expected} (tp={} pp={} dp={} ep={} cp={}), got {world_size}", + self.sizes.tensor, + self.sizes.pipeline, + self.sizes.data, + self.sizes.expert, + self.sizes.context, + ); + } + Ok(()) + } + + /// Convert a global rank to its five local coordinates. + pub fn coordinates(&self, rank: usize) -> Result { + if rank >= self.world_size() { + bail!( + "global rank {rank} is outside world_size {}", + self.world_size() + ); + } + let mut remainder = rank; + let mut coordinates = RankCoordinates::ZERO; + for axis in self.order { + let size = self.sizes.get(axis); + coordinates.set(axis, remainder % size); + remainder /= size; + } + debug_assert_eq!(remainder, 0); + Ok(coordinates) + } + + /// Convert local coordinates to a global rank. + pub fn rank(&self, coordinates: RankCoordinates) -> Result { + for axis in ParallelAxis::ALL { + if coordinates.get(axis) >= self.sizes.get(axis) { + bail!( + "{} coordinate {} is outside size {}", + axis.name(), + coordinates.get(axis), + self.sizes.get(axis) + ); + } + } + let mut rank = 0usize; + let mut stride = 1usize; + for axis in self.order { + rank = rank + .checked_add( + coordinates + .get(axis) + .checked_mul(stride) + .ok_or_else(|| anyhow!("parallel rank calculation overflowed usize"))?, + ) + .ok_or_else(|| anyhow!("parallel rank calculation overflowed usize"))?; + stride = stride + .checked_mul(self.sizes.get(axis)) + .ok_or_else(|| anyhow!("parallel rank calculation overflowed usize"))?; + } + Ok(rank) + } + + /// Return the global ranks in the process group for one axis. + pub fn group(&self, rank: usize, axis: ParallelAxis) -> Result> { + let coordinates = self.coordinates(rank)?; + let mut ranks = Vec::with_capacity(self.sizes.get(axis)); + for local_rank in 0..self.sizes.get(axis) { + let mut member = coordinates; + member.set(axis, local_rank); + ranks.push(self.rank(member)?); + } + Ok(ranks) + } + + pub fn tensor_group(&self, rank: usize) -> Result> { + self.group(rank, ParallelAxis::Tensor) + } + + pub fn pipeline_group(&self, rank: usize) -> Result> { + self.group(rank, ParallelAxis::Pipeline) + } + + pub fn data_group(&self, rank: usize) -> Result> { + self.group(rank, ParallelAxis::Data) + } + + pub fn expert_group(&self, rank: usize) -> Result> { + self.group(rank, ParallelAxis::Expert) + } + + pub fn context_group(&self, rank: usize) -> Result> { + self.group(rank, ParallelAxis::Context) + } + + pub fn tensor_rank(&self, rank: usize) -> Result { + Ok(self.coordinates(rank)?.tensor) + } + + pub fn pipeline_rank(&self, rank: usize) -> Result { + Ok(self.coordinates(rank)?.pipeline) + } + + pub fn data_rank(&self, rank: usize) -> Result { + Ok(self.coordinates(rank)?.data) + } + + pub fn expert_rank(&self, rank: usize) -> Result { + Ok(self.coordinates(rank)?.expert) + } + + pub fn context_rank(&self, rank: usize) -> Result { + Ok(self.coordinates(rank)?.context) + } + + pub fn is_first_pipeline_stage(&self, rank: usize) -> Result { + Ok(self.pipeline_rank(rank)? == 0) + } + + pub fn is_last_pipeline_stage(&self, rank: usize) -> Result { + Ok(self.pipeline_rank(rank)? + 1 == self.pipeline_model_parallel_size()) + } + + /// Return the contiguous layer range assigned to a pipeline stage. + pub fn layer_range(&self, rank: usize, num_layers: usize) -> Result> { + let stage = self.pipeline_rank(rank)?; + let stages = self.pipeline_model_parallel_size(); + if num_layers < stages { + bail!("num_layers={num_layers} must be >= pipeline size {stages}"); + } + let start = num_layers * stage / stages; + let end = num_layers * (stage + 1) / stages; + Ok(start..end) + } +} + +fn parse_env_usize(name: &str) -> Result { + env::var(name) + .with_context(|| format!("{name} is not set"))? + .parse::() + .with_context(|| format!("{name} must be a positive integer")) +} + +fn first_env_usize(names: &[&str]) -> Result> { + let Some((name, raw)) = names + .iter() + .find_map(|name| env::var(name).ok().map(|raw| (*name, raw))) + else { + return Ok(None); + }; + let value = raw + .parse::() + .with_context(|| format!("{name} must be a positive integer"))?; + if value == 0 { + bail!("{name} must be greater than zero"); + } + Ok(Some(value)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn megatron_like_tp_dp_pp_order_round_trips_and_builds_groups() { + let topology = ParallelTopology::with_order(2, 4, 3, 1, 1, "tp-dp-pp").unwrap(); + assert_eq!(topology.world_size(), 24); + let rank = topology + .rank(RankCoordinates { + tensor: 1, + pipeline: 2, + data: 1, + expert: 0, + context: 0, + }) + .unwrap(); + // tp is the least-significant digit: 1 + 1*2 + 2*2*3 = 15. + assert_eq!(rank, 15); + assert_eq!(topology.coordinates(rank).unwrap().tensor, 1); + assert_eq!(topology.coordinates(rank).unwrap().data, 1); + assert_eq!(topology.coordinates(rank).unwrap().pipeline, 2); + assert_eq!(topology.tensor_group(rank).unwrap(), vec![14, 15]); + assert_eq!(topology.data_group(rank).unwrap(), vec![13, 15, 17]); + assert_eq!(topology.pipeline_group(rank).unwrap(), vec![3, 9, 15, 21]); + } + + #[test] + fn five_dimensional_default_order_matches_orthogonal_mapping() { + let topology = ParallelTopology::new(2, 2, 2, 2, 2).unwrap(); + let coordinates = RankCoordinates { + tensor: 1, + pipeline: 0, + data: 1, + expert: 1, + context: 0, + }; + // tp-cp-ep-dp-pp: 1 + 1*(2*2) + 1*(2*2*2) = 13. + assert_eq!(topology.rank(coordinates).unwrap(), 13); + assert_eq!(topology.coordinates(13).unwrap(), coordinates); + assert_eq!(topology.expert_group(13).unwrap(), vec![9, 13]); + assert_eq!(topology.context_group(13).unwrap(), vec![13, 15]); + } + + #[test] + fn rejects_invalid_order_and_world_size() { + assert!(ParallelTopology::with_order(1, 1, 1, 1, 1, "tp-unknown").is_err()); + assert!(ParallelTopology::with_order(1, 1, 1, 1, 1, "tp-dp-dp-ep-cp").is_err()); + let topology = ParallelTopology::new(2, 2, 1, 1, 1).unwrap(); + assert!(topology.validate_world_size(3).is_err()); + assert!(topology.coordinates(4).is_err()); + } + + #[test] + fn layer_ranges_are_contiguous_and_cover_the_model() { + let topology = ParallelTopology::new(1, 3, 1, 1, 1).unwrap(); + assert_eq!(topology.layer_range(0, 10).unwrap(), 0..3); + assert_eq!(topology.layer_range(1, 10).unwrap(), 3..6); + assert_eq!(topology.layer_range(2, 10).unwrap(), 6..10); + assert!(topology.layer_range(0, 2).is_err()); + } +} From afcf091ff2889fe808157f1f5857ef682cd9bc42 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 04:45:24 +0800 Subject: [PATCH 009/156] fix: restore native optimizer step on checkpoint load --- .../kernels/qwen3_6_kernels.cpp | 12 +++- crates/rustrain-qwen3-6/src/kernel.rs | 15 ++++- .../rustrain-qwen3-6/tests/native_smoke.cpp | 6 ++ crates/rustrain-server/src/session.rs | 5 +- docs/qwen35-qwen36-megatron-audit.md | 57 +++++++++++++++++++ 5 files changed, 92 insertions(+), 3 deletions(-) create mode 100644 docs/qwen35-qwen36-megatron-audit.md diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 896f32ac..4c3cd3aa 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -3456,7 +3456,7 @@ static at::Tensor mtp_compute_loss( extern "C" { __attribute__((visibility("default"))) int64_t qwen36_kernel_abi_version() { - return 7; + return 8; } // Create training context — called once at startup @@ -4829,6 +4829,16 @@ int64_t qwen36_get_step_count(void* ctx_ptr) { return (int64_t)reinterpret_cast(ctx_ptr)->step_count; } +// Restore the Adam bias-correction clock independently from tensor state. +// Checkpoint loading imports m/v through a separate ABI, so omitting this +// value would resume the next update as step 1 even for a mature optimizer. +__attribute__((visibility("default"))) +int32_t qwen36_set_step_count(void* ctx_ptr, int64_t step_count) { + if (!ctx_ptr || step_count < 0) return -1; + reinterpret_cast(ctx_ptr)->step_count = step_count; + return 0; +} + __attribute__((visibility("default"))) int64_t qwen36_export_optimizer_state(void* ctx_ptr, void** m_ptrs, void** v_ptrs, int64_t max_count) { auto* ctx = reinterpret_cast(ctx_ptr); diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs index ac8c6600..a45b2018 100644 --- a/crates/rustrain-qwen3-6/src/kernel.rs +++ b/crates/rustrain-qwen3-6/src/kernel.rs @@ -44,6 +44,7 @@ type FnGetLoraA = unsafe extern "C" fn(*mut c_void, i64) -> *mut c_void; type FnGetLoraB = unsafe extern "C" fn(*mut c_void, i64) -> *mut c_void; type FnSetLoraTensor = unsafe extern "C" fn(*mut c_void, i64, i32, *mut c_void) -> i32; type FnGetStepCount = unsafe extern "C" fn(*mut c_void) -> i64; +type FnSetStepCount = unsafe extern "C" fn(*mut c_void, i64) -> i32; type FnExportOptimizer = unsafe extern "C" fn(*mut c_void, *mut *mut c_void, *mut *mut c_void, i64) -> i64; type FnImportOptimizer = @@ -116,6 +117,7 @@ struct KernelHandles { get_lora_b: FnGetLoraB, set_lora_tensor: FnSetLoraTensor, get_step_count: FnGetStepCount, + set_step_count: FnSetStepCount, export_optimizer: FnExportOptimizer, import_optimizer: FnImportOptimizer, free_ctx: FnFreeCtx, @@ -174,7 +176,7 @@ unsafe fn load_kernels() -> Option { }}; } let abi_version: FnKernelAbiVersion = sym!("qwen36_kernel_abi_version"); - if abi_version() != 7 { + if abi_version() != 8 { return None; } Some(KernelHandles { @@ -188,6 +190,7 @@ unsafe fn load_kernels() -> Option { get_lora_b: sym!("qwen36_get_lora_b"), set_lora_tensor: sym!("qwen36_set_lora_tensor"), get_step_count: sym!("qwen36_get_step_count"), + set_step_count: sym!("qwen36_set_step_count"), export_optimizer: sym!("qwen36_export_optimizer_state"), import_optimizer: sym!("qwen36_import_optimizer_state"), free_ctx: sym!("qwen36_free_training_context"), @@ -932,6 +935,16 @@ impl CppTrainingContext { unsafe { (kh.get_step_count)(self.ptr) } } + /// Restore the native Adam bias-correction step from a checkpoint. + pub fn set_step_count(&self, step_count: i64) -> Result<()> { + let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; + let status = unsafe { (kh.set_step_count)(self.ptr, step_count) }; + if status != 0 { + bail!("C++ set_step_count failed for step {step_count}"); + } + Ok(()) + } + /// Export Adam optimizer state (m and v vectors). /// Returns (m_tensors, v_tensors) — owned copies on CPU. pub fn export_optimizer_state(&self) -> Result<(Vec, Vec)> { diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index 89423b4f..bac57733 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -32,6 +32,7 @@ extern "C" double qwen36_train_step(void*, void*, void*, void*); extern "C" double qwen36_train_micro_step( void*, void*, void*, void*, double, int32_t); extern "C" int64_t qwen36_get_step_count(void*); +extern "C" int32_t qwen36_set_step_count(void*, int64_t); extern "C" double qwen36_eval_step(void*, void*, void*, void*); extern "C" double qwen36_train_multi_lora( void*, void*, void*, void*, int32_t, int32_t); @@ -356,6 +357,11 @@ int main() { assert(qwen36_set_lora_tensor(ctx, 0, 1, &linear_b_value) == 0); auto linear_a_before = linear_a->clone(); assert(qwen36_get_step_count(ctx) == 0); + assert(qwen36_set_step_count(ctx, -1) != 0); + assert(qwen36_get_step_count(ctx) == 0); + assert(qwen36_set_step_count(ctx, 7) == 0); + assert(qwen36_get_step_count(ctx) == 7); + assert(qwen36_set_step_count(ctx, 0) == 0); const double accum_loss_0 = qwen36_train_micro_step( ctx, &input_ids, &target_mask, &attention_mask, 0.5, 0); c10::cuda::device_synchronize(); diff --git a/crates/rustrain-server/src/session.rs b/crates/rustrain-server/src/session.rs index 8db56469..f5113c1b 100644 --- a/crates/rustrain-server/src/session.rs +++ b/crates/rustrain-server/src/session.rs @@ -1,6 +1,6 @@ //! Training session trait + Qwen3.6 implementation. -use anyhow::{anyhow, bail, Context, Result}; +use anyhow::{Context, Result, anyhow, bail}; use std::path::PathBuf; use std::sync::Arc; use tch::{Device, Kind, Tensor}; @@ -820,6 +820,9 @@ impl TrainingSession for Qwen36Session { ctx.import_optimizer_state(&data.adam_m, &data.adam_v)?; tracing::info!(imported = data.adam_m.len(), "optimizer state imported"); } + let native_step = i64::try_from(data.manifest.step) + .context("checkpoint step exceeds the native optimizer range")?; + ctx.set_step_count(native_step)?; } self.step = data.manifest.step; self.last_loss = data.manifest.loss; diff --git a/docs/qwen35-qwen36-megatron-audit.md b/docs/qwen35-qwen36-megatron-audit.md new file mode 100644 index 00000000..9515b925 --- /dev/null +++ b/docs/qwen35-qwen36-megatron-audit.md @@ -0,0 +1,57 @@ +# Qwen3.5/3.6 LoRA 并行与性能审计 + +本文记录当前 native Qwen3.5/3.6 LoRA 后端与 Megatron-LM 级训练栈的边界。结论按实际代码和 smoke/integration 结果整理，不把配置字段或通用拓扑类型当作已经实现的 kernel。 + +## 结论 + +- 模型语义：Qwen3.5 dense、Qwen3.6 dense/MoE 的 native forward/backward 路径已经覆盖 hybrid full attention、GDN、MoE、MTP 和 LoRA 目标模块；已有配置解析、集成测试及 H20 native smoke 证据。 +- 已实现并可验证的分布式子集：MoE expert parallel，以及 replicated LoRA 的 data parallel；梯度累积和 dynamic multi-LoRA 已有 logical-step 边界。 +- 性能：MoE grouped dispatch 相对逐 expert matmul 的已有 microbenchmark 为约 3.70x（E=32, N=4096, H=2048, I=768，结果误差为 0）；这不是端到端训练吞吐或 Megatron 对比。 +- 尚未实现：Qwen native 路径的 tensor parallel、pipeline parallel、context parallel，以及 TP/PP/CP 与 EP/DP 的组合。当前训练上下文仍由单个进程持有完整 dense 权重和完整层栈。 +- 因此当前实现不能宣称“Megatron-LM 级别”。它是一个计算集中在 C++ 的 LoRA/EP/DP 子集，离 Megatron 的完整并行和通信重叠仍有实质差距。 + +## 当前能力矩阵 + +| 能力 | 当前状态 | 证据/限制 | +| --- | --- | --- | +| Qwen3.5 full attention | 已实现 | native smoke、Qwen3.5 配置集成测试 | +| Qwen3.5 GDN/linear attention | 已实现 | CUDA delta-rule forward/backward 与 native smoke | +| Qwen3.6 MoE | 已实现 | grouped dispatch、EP smoke；完整模型仍需目标 GPU/权重运行 | +| MTP | 已实现 | C++ hidden gradient 检查和集成测试；可通过环境变量关闭 | +| fixed LoRA | 已实现 | attention/GDN/MLP/shared/routed expert 目标模块 | +| dynamic multi-LoRA | 已实现子集 | 请求按 adapter 分组，单个 logical step 统一 backward/Adam；adapter 仍共享 context optimizer step | +| microbatch accumulation | 已实现子集 | non-final microbatch 只 backward，final microbatch 才 optimizer；梯度仍累加在 BF16 leaf 上 | +| replicated data parallel | 已实现 | logical-step 边界同步 replicated LoRA；EP expert 参数不走该 reduction | +| expert parallel | 已实现子集 | 路由输出 all-reduce 和本地 expert 权重；没有 DeepEP 式 fused A2A/dispatch overlap | +| tensor parallel | 未实现于 Qwen native | 不切分 attention/MLP/LM-head 权重，也没有 Qwen TP communicator | +| pipeline parallel | 未实现于 Qwen native | 没有 stage 切分、microbatch scheduler 或 activation send/recv | +| context parallel | 未实现于 Qwen native | 没有 ring attention、跨 rank KV/索引合并 | +| TP/PP/CP 组合 checkpoint | 未实现 | 当前 checkpoint 不是 Megatron rank-sharded topology | + +## 与 Megatron-LM 的关键差距 + +### 并行语义 + +Megatron 的 TP 会按 head、hidden/intermediate 和 vocab 维度切分权重，并在线性层边界执行必要的 reduce-scatter/all-reduce；PP 会把层分到不同 stage 并使用 1F1B 等调度；CP 会在序列和 attention state 上做跨 rank 通信。当前 Qwen native `TrainingContext` 仍加载完整模型并在一个 C++ forward 中执行全部层，因此仅增加 `tensor_model_parallel_size` 等配置不能得到正确的 TP/PP/CP。 + +当前 DP/EP 也不是完整 Megatron 语义：DP 只同步 replicated LoRA 梯度，expert 参数留在 EP rank；EP 使用已有 all-reduce，但没有 fused token dispatch/combine、异步 A2A 和通信计算重叠。 + +### 优化器与恢复 + +固定 LoRA 的 Adam 状态可导出/导入，且 native context 的 logical step 需要与 checkpoint step 对齐。dynamic adapter 的请求频率不同，但目前仍共享 context-level step；尚无每租户独立 optimizer step、FP32 gradient accumulator 或 accumulation window abort/zero API。这些差距会影响长时间多租户训练的数值一致性和故障恢复。 + +### 性能工程 + +当前粗粒度 C++ FFI、grouped MoE 和 activation checkpoint/offload 是有效优化，但尚无 Megatron/Transformer Engine 级别的端到端数据：没有完整模型在同一 GPU、序列长度、microbatch、精度和通信配置下的 tokens/s、显存、扩展效率对照，也没有 FP8/FP4 参数与 fused attention/DeepEP 的 Qwen 路径。 + +## 验证边界 + +已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试、H20 ABI0 native smoke，以及已有 ABI1 环境中的单卡、EP 和 DP smoke。没有完成 Qwen3.5/3.6 完整大模型的长时间训练、跨节点通信、TP/PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖的模型/LoRA/EP/DP 子集，而不是所有并行配置。 + +## 继续达到 Megatron 级别所需的最小工作包 + +1. 建立 5D TP/PP/DP/EP/CP topology，并让 launcher、NCCL process groups 和 checkpoint 使用同一 rank 映射。 +2. 为 Qwen full/GDN attention、dense MLP、MoE、LM-head/CE 实现 TP shard 和对应 collective；为 PP 实现 stage forward/backward 与 1F1B scheduler；为 CP 实现 ring attention/state exchange。 +3. 将 EP dispatch/combine 替换为 fused/异步路径，并测量通信与计算重叠。 +4. 为 LoRA 增加 FP32 accumulation、每 adapter optimizer step、可恢复的 accumulation 状态和 rank-sharded checkpoint。 +5. 在固定硬件和 workload 上，与 Megatron-LM 记录 tokens/s、step time、峰值显存、通信占比和 loss 曲线。 From 5a58d6f7449e50e4f6463ff4327a69141364b24b Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 05:12:53 +0800 Subject: [PATCH 010/156] feat: isolate dynamic lora optimizer state --- crates/rustrain-ipc/src/command.rs | 2 + .../kernels/qwen3_6_kernels.cpp | 250 ++++++++++++++---- crates/rustrain-qwen3-6/src/kernel.rs | 71 ++++- .../rustrain-qwen3-6/tests/native_smoke.cpp | 56 ++++ crates/rustrain-server/src/api.rs | 15 ++ crates/rustrain-server/src/checkpoint.rs | 22 ++ crates/rustrain-server/src/ep.rs | 13 +- crates/rustrain-server/src/session.rs | 49 +++- docs/plans/qwen-lora-megatron-progress.md | 38 +++ docs/plans/qwen-lora-megatron-spec.md | 58 ++++ 10 files changed, 510 insertions(+), 64 deletions(-) create mode 100644 docs/plans/qwen-lora-megatron-progress.md create mode 100644 docs/plans/qwen-lora-megatron-spec.md diff --git a/crates/rustrain-ipc/src/command.rs b/crates/rustrain-ipc/src/command.rs index fd4d8e3e..4c77edc7 100644 --- a/crates/rustrain-ipc/src/command.rs +++ b/crates/rustrain-ipc/src/command.rs @@ -67,6 +67,8 @@ pub enum EpCommand { seq_len: usize, n_total: i32, lora_rank: i32, + #[serde(default)] + adapter_ids: Vec, }, EvalStep { session_id: String, diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 4c3cd3aa..94963ead 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -1463,6 +1464,9 @@ struct TrainingContext { struct LoRAAdapter { int64_t id; int64_t rank; + // Each tenant owns an independent Adam bias-correction clock. The + // session-wide step_count remains a transport/metric clock only. + int64_t optimizer_step = 0; double alpha; std::set target_layers; std::set target_modules; @@ -3456,7 +3460,7 @@ static at::Tensor mtp_compute_loss( extern "C" { __attribute__((visibility("default"))) int64_t qwen36_kernel_abi_version() { - return 8; + return 9; } // Create training context — called once at startup @@ -4028,6 +4032,40 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( "attention_mask must match input_ids shape"); } const at::Tensor saved_attention_mask = ctx->attention_mask; + struct AttentionMaskGuard { + TrainingContext* ctx; + at::Tensor saved; + ~AttentionMaskGuard() { ctx->attention_mask = saved; } + } attention_mask_guard{ctx, saved_attention_mask}; + + struct AdapterRegistryChunkGuard { + TrainingContext* ctx; + std::vector all; + bool active = false; + + AdapterRegistryChunkGuard( + TrainingContext* context, int64_t start, int64_t end) + : ctx(context) { + all.swap(ctx->adapters); + try { + ctx->adapters.assign(all.begin() + start, all.begin() + end); + active = true; + } catch (...) { + ctx->adapters.clear(); + ctx->adapters.swap(all); + throw; + } + } + + void restore() { + if (!active) return; + ctx->adapters.clear(); + ctx->adapters.swap(all); + active = false; + } + + ~AdapterRegistryChunkGuard() { restore(); } + }; // Compute N_max from available GPU memory. // CRITICAL: all workers must agree on n_max to keep NCCL all-reduce in sync. @@ -4070,16 +4108,10 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( double total_loss = 0.0; int64_t num_chunks = (total_adapters + n_max - 1) / n_max; - // Chunking is a memory scheduling detail, not an optimizer step. All - // adapters in this call must use the same Adam bias correction. + // Chunking is a memory scheduling detail, not an optimizer step. The + // session clock is kept for backwards-compatible status reporting; + // Adam bias correction below uses each adapter's own clock. ctx->step_count++; - const double logical_step = (double)ctx->step_count; - const double bias_correction1 = 1.0 - std::pow(ctx->beta1, logical_step); - const double bias_correction2 = 1.0 - std::pow(ctx->beta2, logical_step); - const float lr_scaled = (float)(ctx->lr / bias_correction1); - const float eps_scaled = (float)(ctx->eps / std::sqrt(bias_correction2)); - const float one_minus_b1 = (float)(1.0 - ctx->beta1); - const float one_minus_b2 = (float)(1.0 - ctx->beta2); for (int64_t chunk = 0; chunk < num_chunks; chunk++) { int64_t start = chunk * n_max; @@ -4090,13 +4122,9 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( ctx->lora_batch_valid = false; ctx->lora_cache_valid = false; - // Temporarily set lora_batch_valid so prepare_lora_batch runs - // We need to select only adapters[start:end] - // HACK: move non-chunk adapters to a temp vector, run, then restore - std::vector all_adapters; - all_adapters.swap(ctx->adapters); - ctx->adapters.assign( - all_adapters.begin() + start, all_adapters.begin() + end); + // Scope the registry to this activation-memory chunk. The guard + // restores the full registry even if forward/backward fails. + AdapterRegistryChunkGuard registry_guard(ctx, start, end); // Mark batched mode active ctx->lora_batch_valid = true; // triggers prepare_lora_batch in forward @@ -4185,47 +4213,58 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( // Restore the complete registry before the next chunk. Gradients // remain attached to the intrusive tensor handles, so all chunks // can accumulate and the optimizer runs exactly once below. - ctx->adapters.swap(all_adapters); + registry_guard.restore(); if (chunk == num_chunks - 1) { // DP gradient synchronization and Adam belong to the logical // multi-tenant step, never to an activation-memory chunk. synchronize_lora_gradients(ctx, target_mask); - // Adam step + // Adam step. Group tenants by their own logical clock so + // newly-added or resumed tenants do not inherit another + // tenant's bias correction. Adapters with the same clock + // still share one fused multi-tensor launch. at::AutoGradMode guard(false); ctx->lora_cache_valid = false; ctx->lora_batch_valid = false; - - std::vector h_params, h_grads; - std::vector h_m, h_v; - std::vector h_sizes; - + for (auto& adapter : ctx->adapters) adapter.optimizer_step++; + std::map> groups; for (auto& adapter : ctx->adapters) { - for (auto& [layer_idx, pairs] : adapter.params) { - auto& adam_states = adapter.adam_state[layer_idx]; - for (size_t i = 0; i < pairs.size(); i++) { - auto& [a, b] = pairs[i]; - auto& [m_a, v_a, m_b, v_b] = adam_states[i]; - if (a.grad().defined() && a.scalar_type() == at::kBFloat16) { - h_params.push_back(a.data_ptr()); - h_grads.push_back(a.grad().data_ptr()); - h_m.push_back((float*)m_a.data_ptr()); - h_v.push_back((float*)v_a.data_ptr()); - h_sizes.push_back((int)a.numel()); - } - if (b.grad().defined() && b.scalar_type() == at::kBFloat16) { - h_params.push_back(b.data_ptr()); - h_grads.push_back(b.grad().data_ptr()); - h_m.push_back((float*)m_b.data_ptr()); - h_v.push_back((float*)v_b.data_ptr()); - h_sizes.push_back((int)b.numel()); + groups[adapter.optimizer_step].push_back(&adapter); + } + for (auto& [logical_step, adapters] : groups) { + std::vector h_params, h_grads; + std::vector h_m, h_v; + std::vector h_sizes; + for (auto* adapter : adapters) { + for (auto& [layer_idx, pairs] : adapter->params) { + auto& adam_states = adapter->adam_state[layer_idx]; + for (size_t i = 0; i < pairs.size(); i++) { + auto& [a, b] = pairs[i]; + auto& [m_a, v_a, m_b, v_b] = adam_states[i]; + if (a.grad().defined() && a.scalar_type() == at::kBFloat16) { + h_params.push_back(a.data_ptr()); + h_grads.push_back(a.grad().data_ptr()); + h_m.push_back((float*)m_a.data_ptr()); + h_v.push_back((float*)v_a.data_ptr()); + h_sizes.push_back((int)a.numel()); + } + if (b.grad().defined() && b.scalar_type() == at::kBFloat16) { + h_params.push_back(b.data_ptr()); + h_grads.push_back(b.grad().data_ptr()); + h_m.push_back((float*)m_b.data_ptr()); + h_v.push_back((float*)v_b.data_ptr()); + h_sizes.push_back((int)b.numel()); + } } } } - } - - if (!h_params.empty()) { + if (h_params.empty()) continue; + const double step_f = (double)logical_step; + const float lr_scaled = (float)(ctx->lr / (1.0 - std::pow(ctx->beta1, step_f))); + const float eps_scaled = (float)(ctx->eps / std::sqrt(1.0 - std::pow(ctx->beta2, step_f))); + const float one_minus_b1 = (float)(1.0 - ctx->beta1); + const float one_minus_b2 = (float)(1.0 - ctx->beta2); int n_params = (int)h_params.size(); auto opts_cpu_long = at::TensorOptions().dtype(at::kLong).device(at::kCPU); auto opts_cpu_int = at::TensorOptions().dtype(at::kInt).device(at::kCPU); @@ -4234,13 +4273,12 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( auto m_cpu = at::from_blob(h_m.data(), {n_params}, opts_cpu_long); auto v_cpu = at::from_blob(h_v.data(), {n_params}, opts_cpu_long); auto sizes_cpu = at::from_blob(h_sizes.data(), {n_params}, opts_cpu_int); - ctx->adam_dev_bufs.ensure(n_params, ctx->adapters[0].params.begin()->second[0].first); + ctx->adam_dev_bufs.ensure(n_params, adapters[0]->params.begin()->second[0].first); ctx->adam_dev_bufs.params_buf.narrow(0, 0, n_params).copy_(params_cpu); ctx->adam_dev_bufs.grads_buf.narrow(0, 0, n_params).copy_(grads_cpu); ctx->adam_dev_bufs.m_buf.narrow(0, 0, n_params).copy_(m_cpu); ctx->adam_dev_bufs.v_buf.narrow(0, 0, n_params).copy_(v_cpu); ctx->adam_dev_bufs.sizes_buf.narrow(0, 0, n_params).copy_(sizes_cpu); - auto stream = c10::cuda::getCurrentCUDAStream().stream(); launch_fused_adam_multi( (void**)ctx->adam_dev_bufs.params_buf.data_ptr(), @@ -4248,12 +4286,9 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( (float**)ctx->adam_dev_bufs.m_buf.data_ptr(), (float**)ctx->adam_dev_bufs.v_buf.data_ptr(), (int*)ctx->adam_dev_bufs.sizes_buf.data_ptr(), - n_params, - (float)ctx->beta1, (float)ctx->beta2, - lr_scaled, eps_scaled, - one_minus_b1, one_minus_b2, - (void*)stream - ); + n_params, (float)ctx->beta1, (float)ctx->beta2, + lr_scaled, eps_scaled, one_minus_b1, one_minus_b2, + (void*)stream); } } @@ -4263,7 +4298,6 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( (long)(chunk + 1), (long)num_chunks, (long)n, loss_val); } - ctx->attention_mask = saved_attention_mask; return total_loss / total_adapters; } catch (const std::exception& e) { fprintf(stderr, "[train_multi] FAILED: %s\n", e.what()); @@ -4274,6 +4308,89 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( } } +// Train only the requested dynamic tenants. The existing train_multi_lora +// implementation already batches activation-level LoRA projections and owns +// the logical Adam boundary; this wrapper scopes its adapter registry to the +// selected IDs and restores the original order on every exit. +__attribute__((visibility("default"))) double qwen36_train_multi_lora_selected( + void* ctx_ptr, + void* input_ids_ptr, + void* target_mask_ptr, + void* attention_mask_ptr, + const int64_t* adapter_ids, + int32_t n_adapters, + int32_t lora_rank +) { + auto* ctx = reinterpret_cast(ctx_ptr); + std::vector original; + std::vector selected; + std::vector merged; + std::vector selected_indexes; + std::vector moved; + bool registry_detached = false; + bool selected_installed = false; + auto restore_registry = [&]() { + if (!ctx || !registry_detached) return; + if (selected_installed) { + selected.swap(ctx->adapters); + selected_installed = false; + } + for (size_t i = 0; i < original.size(); ++i) { + if (!moved[i]) { + merged.push_back(std::move(original[i])); + continue; + } + auto selected_it = std::find( + selected_indexes.begin(), selected_indexes.end(), i); + TORCH_CHECK(selected_it != selected_indexes.end(), + "selected adapter index disappeared: ", i); + const auto selected_index = static_cast( + selected_it - selected_indexes.begin()); + merged.push_back(std::move(selected[selected_index])); + } + ctx->adapters.swap(merged); + registry_detached = false; + }; + try { + TORCH_CHECK(ctx && adapter_ids && n_adapters > 0, + "selected multi-LoRA requires at least one adapter ID"); + const auto original_count = ctx->adapters.size(); + original.reserve(original_count); + selected.reserve(n_adapters); + merged.reserve(original_count); + selected_indexes.reserve(n_adapters); + moved.resize(original_count, 0); + original.swap(ctx->adapters); + registry_detached = true; + for (int32_t i = 0; i < n_adapters; ++i) { + TORCH_CHECK(adapter_ids[i] > 0, "selected adapter IDs must be positive"); + auto it = std::find_if(original.begin(), original.end(), + [&](const auto& adapter) { return adapter.id == adapter_ids[i]; }); + TORCH_CHECK(it != original.end(), "unknown selected adapter ID: ", adapter_ids[i]); + const auto index = static_cast(it - original.begin()); + TORCH_CHECK(!moved[index], "duplicate selected adapter ID: ", adapter_ids[i]); + moved[index] = 1; + selected_indexes.push_back(index); + selected.push_back(std::move(*it)); + } + ctx->adapters.swap(selected); + selected_installed = true; + const double loss = qwen36_train_multi_lora( + ctx_ptr, input_ids_ptr, target_mask_ptr, attention_mask_ptr, + n_adapters, lora_rank); + restore_registry(); + return loss; + } catch (const std::exception& e) { + try { restore_registry(); } catch (...) {} + fprintf(stderr, "[train_multi_selected] FAILED: %s\n", e.what()); + return -1.0; + } catch (...) { + try { restore_registry(); } catch (...) {} + fprintf(stderr, "[train_multi_selected] FAILED: unknown exception\n"); + return -1.0; + } +} + // Set NCCL communicator for Expert Parallel all-reduce // Creates NCCL communicator directly in C++ using env vars RANK/WORLD_SIZE. // Rank 0 generates unique ID and writes to /tmp/rustrain-nccl/nccl-id.bin @@ -4796,6 +4913,31 @@ int32_t qwen36_set_adapter_optimizer_tensor( } } +__attribute__((visibility("default"))) +int64_t qwen36_get_adapter_step_count(void* ctx_ptr, int64_t adapter_id) { + auto* ctx = reinterpret_cast(ctx_ptr); + if (!ctx || adapter_id <= 0) return -1; + for (const auto& adapter : ctx->adapters) { + if (adapter.id == adapter_id) return adapter.optimizer_step; + } + return -1; +} + +__attribute__((visibility("default"))) +int32_t qwen36_set_adapter_step_count( + void* ctx_ptr, int64_t adapter_id, int64_t step_count +) { + auto* ctx = reinterpret_cast(ctx_ptr); + if (!ctx || adapter_id <= 0 || step_count < 0) return -1; + for (auto& adapter : ctx->adapters) { + if (adapter.id == adapter_id) { + adapter.optimizer_step = step_count; + return 0; + } + } + return -1; +} + __attribute__((visibility("default"))) int64_t qwen36_get_lora_count(void* ctx_ptr) { auto* ctx = reinterpret_cast(ctx_ptr); diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs index a45b2018..c9e7ff98 100644 --- a/crates/rustrain-qwen3-6/src/kernel.rs +++ b/crates/rustrain-qwen3-6/src/kernel.rs @@ -38,6 +38,15 @@ type FnTrainMicroStep = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void, f64, i32) -> f64; type FnTrainMultiLora = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void, i32, i32) -> f64; +type FnTrainMultiLoraSelected = unsafe extern "C" fn( + *mut c_void, + *mut c_void, + *mut c_void, + *mut c_void, + *const i64, + i32, + i32, +) -> f64; type FnEvalStep = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void) -> f64; type FnGetLoraCount = unsafe extern "C" fn(*mut c_void) -> i64; type FnGetLoraA = unsafe extern "C" fn(*mut c_void, i64) -> *mut c_void; @@ -79,6 +88,8 @@ type FnGetAdapterOptimizerTensor = unsafe extern "C" fn(*mut c_void, i64, i64, *const i8, i32, i32) -> *mut c_void; type FnSetAdapterOptimizerTensor = unsafe extern "C" fn(*mut c_void, i64, i64, *const i8, i32, i32, *mut c_void) -> i32; +type FnGetAdapterStepCount = unsafe extern "C" fn(*mut c_void, i64) -> i64; +type FnSetAdapterStepCount = unsafe extern "C" fn(*mut c_void, i64, i64) -> i32; #[repr(C)] pub struct CppLayerConfig { @@ -111,6 +122,7 @@ struct KernelHandles { train_step: FnTrainStep, train_micro_step: FnTrainMicroStep, train_multi_lora: FnTrainMultiLora, + train_multi_lora_selected: FnTrainMultiLoraSelected, eval_step: FnEvalStep, get_lora_count: FnGetLoraCount, get_lora_a: FnGetLoraA, @@ -136,6 +148,8 @@ struct KernelHandles { set_adapter_id: FnSetAdapterId, get_adapter_optimizer_tensor: FnGetAdapterOptimizerTensor, set_adapter_optimizer_tensor: FnSetAdapterOptimizerTensor, + get_adapter_step_count: FnGetAdapterStepCount, + set_adapter_step_count: FnSetAdapterStepCount, } static KERNELS: OnceLock> = OnceLock::new(); @@ -176,7 +190,7 @@ unsafe fn load_kernels() -> Option { }}; } let abi_version: FnKernelAbiVersion = sym!("qwen36_kernel_abi_version"); - if abi_version() != 8 { + if abi_version() != 9 { return None; } Some(KernelHandles { @@ -184,6 +198,7 @@ unsafe fn load_kernels() -> Option { train_step: sym!("qwen36_train_step"), train_micro_step: sym!("qwen36_train_micro_step"), train_multi_lora: sym!("qwen36_train_multi_lora"), + train_multi_lora_selected: sym!("qwen36_train_multi_lora_selected"), eval_step: sym!("qwen36_eval_step"), get_lora_count: sym!("qwen36_get_lora_count"), get_lora_a: sym!("qwen36_get_lora_a"), @@ -209,6 +224,8 @@ unsafe fn load_kernels() -> Option { set_adapter_id: sym!("qwen36_set_adapter_id"), get_adapter_optimizer_tensor: sym!("qwen36_get_adapter_optimizer_tensor"), set_adapter_optimizer_tensor: sym!("qwen36_set_adapter_optimizer_tensor"), + get_adapter_step_count: sym!("qwen36_get_adapter_step_count"), + set_adapter_step_count: sym!("qwen36_set_adapter_step_count"), }) } @@ -607,6 +624,38 @@ impl CppTrainingContext { Ok(loss) } + /// Train only the requested dynamic adapters. The native wrapper scopes + /// its registry to these IDs so unselected tenants keep their parameters, + /// gradients, and optimizer clocks untouched. + pub fn train_multi_lora_selected( + &self, + input_ids: &Tensor, + target_mask: &Tensor, + attention_mask: &Tensor, + adapter_ids: &[i64], + lora_rank: i32, + ) -> Result { + if adapter_ids.is_empty() { + bail!("selected multi-LoRA requires at least one adapter ID"); + } + let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; + let loss = unsafe { + (kh.train_multi_lora_selected)( + self.ptr, + input_ids.as_ptr() as *mut _, + target_mask.as_ptr() as *mut _, + attention_mask.as_ptr() as *mut _, + adapter_ids.as_ptr(), + adapter_ids.len() as i32, + lora_rank, + ) + }; + if loss < 0.0 { + bail!("C++ train_multi_lora_selected failed"); + } + Ok(loss) + } + /// Get LoRA A tensor by index (for saving). pub fn get_lora_a(&self, index: i64) -> Option { let kh = get_kernels()?; @@ -904,6 +953,26 @@ impl CppTrainingContext { Ok(()) } + /// Return the independent Adam bias-correction clock for a dynamic tenant. + pub fn get_adapter_step_count(&self, adapter_id: i64) -> Result { + let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; + let step = unsafe { (kh.get_adapter_step_count)(self.ptr, adapter_id) }; + if step < 0 { + bail!("C++ get_adapter_step_count failed for adapter {adapter_id}"); + } + Ok(step) + } + + /// Restore a dynamic tenant's Adam bias-correction clock. + pub fn set_adapter_step_count(&self, adapter_id: i64, step_count: i64) -> Result<()> { + let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; + let status = unsafe { (kh.set_adapter_step_count)(self.ptr, adapter_id, step_count) }; + if status != 0 { + bail!("C++ set_adapter_step_count failed for adapter {adapter_id}"); + } + Ok(()) + } + /// Eval step: forward + loss, no backward, no Adam update. pub fn eval_step( &self, diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index bac57733..4747bb64 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -23,6 +23,7 @@ extern "C" void* qwen36_create_training_context( void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, double, double, double, double, double, int64_t, double, int64_t, const int64_t*, int64_t, const char*); +extern "C" int64_t qwen36_kernel_abi_version(); extern "C" int32_t qwen36_init_nccl(void*); extern "C" int64_t qwen36_get_lora_count(void*); extern "C" void* qwen36_get_lora_a(void*, int64_t); @@ -33,9 +34,13 @@ extern "C" double qwen36_train_micro_step( void*, void*, void*, void*, double, int32_t); extern "C" int64_t qwen36_get_step_count(void*); extern "C" int32_t qwen36_set_step_count(void*, int64_t); +extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t); +extern "C" int32_t qwen36_set_adapter_step_count(void*, int64_t, int64_t); extern "C" double qwen36_eval_step(void*, void*, void*, void*); extern "C" double qwen36_train_multi_lora( void*, void*, void*, void*, int32_t, int32_t); +extern "C" double qwen36_train_multi_lora_selected( + void*, void*, void*, void*, const int64_t*, int32_t, int32_t); extern "C" int64_t qwen36_add_lora( void*, int64_t, double, const int64_t*, int64_t, const char*); extern "C" void* qwen36_get_adapter_lora_tensor( @@ -49,6 +54,7 @@ static at::Tensor cuda_rand(std::initializer_list shape) { } int main() { + assert(qwen36_kernel_abi_version() == 9); const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); const int process_rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0"); const int local_rank = std::atoi(std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); @@ -195,6 +201,12 @@ int main() { const int64_t adapter_two = qwen36_add_lora( ctx, rank, 1.0, &target_layer, 1, shared_targets); assert(adapter_one > 0 && adapter_two > adapter_one); + assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 0); + assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 0); + assert(qwen36_set_adapter_step_count(ctx, adapter_one, -1) != 0); + assert(qwen36_set_adapter_step_count(ctx, adapter_one, 4) == 0); + assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 4); + assert(qwen36_set_adapter_step_count(ctx, adapter_one, 0) == 0); auto* dynamic_b = reinterpret_cast( qwen36_get_adapter_lora_tensor( ctx, adapter_one, 0, "shared_gate_proj", 1)); @@ -225,6 +237,8 @@ int main() { const double multi_loss = qwen36_train_multi_lora( ctx, &multi_input_ids, &multi_target_mask, &multi_attention_mask, 2, rank); c10::cuda::device_synchronize(); + assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 1); + assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 1); dynamic_b = reinterpret_cast( qwen36_get_adapter_lora_tensor( ctx, adapter_one, 0, "shared_gate_proj", 1)); @@ -251,6 +265,48 @@ int main() { assert(dynamic_update > 0.0); assert(dynamic_expert_update > 0.0); assert(dynamic_two_update > 0.0); + + auto adapter_one_before_selected = dynamic_b->clone(); + auto adapter_two_before_selected = dynamic_b_two->clone(); + auto selected_input_ids = multi_input_ids.narrow(0, 0, 1).contiguous(); + auto selected_target_mask = multi_target_mask.narrow(0, 0, 1).contiguous(); + auto selected_attention_mask = multi_attention_mask.narrow(0, 0, 1).contiguous(); + const int64_t selected_adapter_ids[] = {adapter_one}; + const double selected_loss = qwen36_train_multi_lora_selected( + ctx, &selected_input_ids, &selected_target_mask, + &selected_attention_mask, selected_adapter_ids, 1, rank); + c10::cuda::device_synchronize(); + assert(selected_loss == selected_loss && selected_loss > 0.0); + assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 2); + assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 1); + dynamic_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_one, 0, "shared_gate_proj", 1)); + dynamic_b_two = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_two, 0, "shared_gate_proj", 1)); + assert(dynamic_b && dynamic_b_two); + const double selected_update = + (*dynamic_b - adapter_one_before_selected).abs().sum().item(); + const double unselected_update = + (*dynamic_b_two - adapter_two_before_selected).abs().sum().item(); + std::printf( + "native_qwen36_selected_multi_lora_smoke loss=%0.8f " + "selected_update=%0.8e unselected_update=%0.8e\n", + selected_loss, selected_update, unselected_update); + assert(selected_update > 0.0); + assert(unselected_update == 0.0); + + const int64_t unknown_adapter_ids[] = {adapter_two + 1000}; + assert(qwen36_train_multi_lora_selected( + ctx, &selected_input_ids, &selected_target_mask, + &selected_attention_mask, unknown_adapter_ids, 1, rank) < 0.0); + assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 2); + assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 1); + assert(qwen36_get_adapter_lora_tensor( + ctx, adapter_one, 0, "shared_gate_proj", 1) != nullptr); + assert(qwen36_get_adapter_lora_tensor( + ctx, adapter_two, 0, "shared_gate_proj", 1) != nullptr); qwen36_free_training_context(ctx); // Dense Qwen3.5 variants use the same per-sample activation path for diff --git a/crates/rustrain-server/src/api.rs b/crates/rustrain-server/src/api.rs index 4896f450..21d820ae 100644 --- a/crates/rustrain-server/src/api.rs +++ b/crates/rustrain-server/src/api.rs @@ -247,6 +247,8 @@ struct TrainMultiLoraHttp { attention_mask: TensorHttp, n_total: i32, lora_rank: i32, + #[serde(default)] + adapter_ids: Vec, } #[derive(Serialize)] struct TrainStepResponse { @@ -817,6 +819,18 @@ async fn ep_train_multi_lora( req.n_total, ) .map_err(|e| err_resp(&e))?; + if !req.adapter_ids.is_empty() { + if req.adapter_ids.len() != req.n_total as usize { + return Err(err_resp(&format!( + "adapter_ids length {} must match n_total={}", + req.adapter_ids.len(), + req.n_total + ))); + } + if req.adapter_ids.iter().any(|id| *id <= 0) { + return Err(err_resp("adapter_ids must contain only positive IDs")); + } + } let cmd = rustrain_ipc::EpCommand::TrainMultiLora { session_id: id, @@ -826,6 +840,7 @@ async fn ep_train_multi_lora( seq_len, n_total: req.n_total, lora_rank: req.lora_rank, + adapter_ids: req.adapter_ids, }; match state.coordinator.dispatch(&cmd) { rustrain_ipc::EpResult::Loss(loss) => Ok(Json(TrainStepResponse { loss, step: 0 })), diff --git a/crates/rustrain-server/src/checkpoint.rs b/crates/rustrain-server/src/checkpoint.rs index 2042255e..3e265585 100644 --- a/crates/rustrain-server/src/checkpoint.rs +++ b/crates/rustrain-server/src/checkpoint.rs @@ -23,6 +23,11 @@ pub struct DynamicAdapterManifest { pub id: i64, pub rank: i64, pub alpha: f64, + /// Adam bias-correction clock for this tenant. It is independent from + /// the session-wide `CheckpointManifest::step` and other adapters. + /// Missing values in older v2 manifests default to zero. + #[serde(default)] + pub optimizer_step: u64, pub target_layers: Vec, pub target_modules: Vec, pub parameter_count: usize, @@ -321,6 +326,7 @@ mod tests { id: 7, rank: 3, alpha: 6.0, + optimizer_step: 19, target_layers: vec![1, 3], target_modules: vec!["q_proj".into(), "down_proj".into()], parameter_count: 2, @@ -359,6 +365,7 @@ mod tests { let loaded_dynamic = &loaded.dynamic_adapters[0]; assert_eq!(loaded_dynamic.manifest.id, 7); assert_eq!(loaded_dynamic.manifest.rank, 3); + assert_eq!(loaded_dynamic.manifest.optimizer_step, 19); assert_eq!(loaded_dynamic.manifest.target_layers, vec![1, 3]); assert_eq!(loaded_dynamic.lora_a.len(), 2); assert_eq!(loaded_dynamic.adam_m.len(), 4); @@ -369,4 +376,19 @@ mod tests { false )); } + + #[test] + fn old_dynamic_manifest_defaults_optimizer_step_to_zero() { + let json = r#"{ + "id": 7, + "rank": 3, + "alpha": 6.0, + "target_layers": [1], + "target_modules": ["q_proj"], + "parameter_count": 1, + "optimizer_count": 2 + }"#; + let manifest: DynamicAdapterManifest = serde_json::from_str(json).unwrap(); + assert_eq!(manifest.optimizer_step, 0); + } } diff --git a/crates/rustrain-server/src/ep.rs b/crates/rustrain-server/src/ep.rs index 3b5bf7a7..70e6b865 100644 --- a/crates/rustrain-server/src/ep.rs +++ b/crates/rustrain-server/src/ep.rs @@ -262,7 +262,16 @@ fn execute_command(session: &mut Qwen36Session, cmd: &EpCommand) -> EpResult { Err(e) => EpResult::Error(e.to_string()), } } - EpCommand::TrainMultiLora { input_ids, target_mask, attention_mask, seq_len, n_total, lora_rank, .. } => { + EpCommand::TrainMultiLora { + input_ids, + target_mask, + attention_mask, + seq_len, + n_total, + lora_rank, + adapter_ids, + .. + } => { let sl = *seq_len as i64; let batch = if *n_total > 0 && input_ids.len() == (*n_total as usize).saturating_mul(*seq_len) @@ -293,7 +302,7 @@ fn execute_command(session: &mut Qwen36Session, cmd: &EpCommand) -> EpResult { input_ids: input_ids_tensor, target_mask: target_mask_tensor, attention_mask: attention_mask_tensor, - }, *n_total, *lora_rank) { + }, *n_total, *lora_rank, adapter_ids) { Ok(TrainOutput { loss, .. }) => EpResult::Loss(loss), Err(e) => EpResult::Error(e.to_string()), } diff --git a/crates/rustrain-server/src/session.rs b/crates/rustrain-server/src/session.rs index f5113c1b..89ad821f 100644 --- a/crates/rustrain-server/src/session.rs +++ b/crates/rustrain-server/src/session.rs @@ -83,6 +83,7 @@ pub trait TrainingSession: Send { input: TrainInput, n_total: i32, rank: i32, + adapter_ids: &[i64], ) -> Result; fn eval_step(&self, input: TrainInput) -> Result; fn save_checkpoint(&self, path: &str) -> Result<(u64, f64)>; @@ -472,19 +473,44 @@ impl TrainingSession for Qwen36Session { input: TrainInput, n_total: i32, rank: i32, + adapter_ids: &[i64], ) -> Result { let ctx = self .ctx .as_ref() .ok_or_else(|| anyhow!("LoRA not initialized"))?; - let loss = ctx.train_multi_lora( - &input.input_ids, - &input.target_mask, - &input.attention_mask, - n_total, - rank, - )?; + if n_total <= 0 { + return Err(anyhow!("n_total must be positive, got {n_total}")); + } + if !adapter_ids.is_empty() { + if adapter_ids.len() != n_total as usize { + return Err(anyhow!( + "selected adapter count {} does not match n_total={n_total}", + adapter_ids.len() + )); + } + if adapter_ids.iter().any(|id| *id <= 0) { + return Err(anyhow!("selected adapter IDs must be positive")); + } + } + let loss = if adapter_ids.is_empty() { + ctx.train_multi_lora( + &input.input_ids, + &input.target_mask, + &input.attention_mask, + n_total, + rank, + )? + } else { + ctx.train_multi_lora_selected( + &input.input_ids, + &input.target_mask, + &input.attention_mask, + adapter_ids, + rank, + )? + }; self.step += 1; self.last_loss = loss; self.state = SessionState::Training { step: self.step }; @@ -630,11 +656,14 @@ impl TrainingSession for Qwen36Session { })?, ); } + let optimizer_step = u64::try_from(ctx.get_adapter_step_count(adapter_id)?) + .context("native dynamic adapter optimizer step is negative")?; dynamic_adapters.push(checkpoint::DynamicAdapterCheckpoint { manifest: checkpoint::DynamicAdapterManifest { id: adapter_id, rank: lora_config.rank, alpha: lora_config.alpha, + optimizer_step, target_layers: lora_config.target_layers.clone(), target_modules: lora_config .target_modules @@ -797,6 +826,12 @@ impl TrainingSession for Qwen36Session { let _ = ctx.remove_lora(adapter_id); return Err(error); } + let optimizer_step = i64::try_from(dynamic.manifest.optimizer_step) + .context("dynamic adapter optimizer step exceeds native range")?; + if let Err(error) = ctx.set_adapter_step_count(adapter_id, optimizer_step) { + let _ = ctx.remove_lora(adapter_id); + return Err(error); + } self.dynamic_lora_configs.insert(adapter_id, lora_config); } } diff --git a/docs/plans/qwen-lora-megatron-progress.md b/docs/plans/qwen-lora-megatron-progress.md new file mode 100644 index 00000000..94a92c1e --- /dev/null +++ b/docs/plans/qwen-lora-megatron-progress.md @@ -0,0 +1,38 @@ +--- +type: ProgressRecord +title: Qwen LoRA Megatron runtime progress +description: Verified milestones and open acceptance gaps for the Qwen native LoRA runtime. +tags: [qwen3.5, qwen3.6, lora, progress] +timestamp: 2026-07-17T00:00:00Z +--- + +# Source Spec + +[qwen-lora-megatron-spec.md](/docs/plans/qwen-lora-megatron-spec.md) + +# Current State + +Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP smoke, dense replicated-DP smoke, dynamic batch logical-step update, ABI9 fixed and per-tenant optimizer-step restore, selected-tenant isolation, and 5D topology mapping. + +Not yet verified or implemented: Qwen native TP/PP/CP, FP32 accumulation/abort, rank-sharded checkpoint topology, DeepEP/TE prebuilt integration, and matched Megatron throughput. + +# Durable Milestones + +- `4e242ee`: added Megatron-style 5D topology contract and launcher normalization; `cargo test -p rustrain-parallel --lib` passed 14/14. +- `afcf091`: restored native C++ Adam step on checkpoint load and added ABI8 smoke assertions; server checkpoint tests passed 2/2. +- ABI9 working tree: selected adapter IDs flow through HTTP, IPC, Rust, and C++; each dynamic tenant owns its Adam clock, checkpoint metadata preserves it, and failed selection restores the complete registry. +- H20 `123.57.26.97:28004`: ABI8 native smoke passed grouped/fallback parity, GDN, dense/MoE LoRA, dynamic adapters, and step setter validation. +- H20 `123.57.26.97:28004`: ABI9 native smoke passed selected-tenant training with a positive selected update, exactly zero unselected update, independent clocks (`2` vs `1`), and registry preservation after an unknown ID. +- Target runtime probe: PyTorch 2.5.1+cu121, ABI0; Transformer Engine, flash-attn, DeepEP, Triton, and DeepSpeed are not importable. + +# Decisions During Execution + +- Keep TP and EP communicators separate; do not reuse the existing EP `LayerConfig.nccl_comm` for LoRA TP deltas. +- Do not enable Qwen TP/PP/CP by merely relaxing runtime validation. +- Treat Exa/Jina dependency search failures as missing evidence, not as proof that a package is compatible. + +# Verification + +Passed: `cargo test -p rustrain-core --lib` (8), `cargo test -p rustrain-parallel --lib` (14), `cargo test -p rustrain-server --lib` (3), Qwen integration (6), remote ABI8 smoke, and remote ABI9 selected-tenant native smoke. + +Not run: full Qwen TP/PP/CP smoke, FP32 accumulation equivalence, rank-sharded checkpoint resume, and matched Megatron performance benchmark. diff --git a/docs/plans/qwen-lora-megatron-spec.md b/docs/plans/qwen-lora-megatron-spec.md new file mode 100644 index 00000000..b961932c --- /dev/null +++ b/docs/plans/qwen-lora-megatron-spec.md @@ -0,0 +1,58 @@ +--- +type: ChangeSpec +title: Qwen3.5/3.6 LoRA Megatron-level runtime +description: Contract for LoRA-only distributed parallelism, MoE performance, GDN, and independent multi-tenant training. +tags: [qwen3.5, qwen3.6, lora, megatron, distributed] +timestamp: 2026-07-17T00:00:00Z +--- + +# Problem + +The native Qwen3.5/3.6 path has correct and tested single-rank, EP, and dense-DP slices, but it is not yet a Megatron-LM-level LoRA runtime. In particular, Qwen native TP/PP/CP are not implemented, dynamic tenants share optimizer state semantics, and the MoE path lacks fused asynchronous token dispatch. + +# Target Outcome + +Provide a Qwen3.5/3.6 LoRA-only training backend whose parallel groups, rank-local adapter parameters, optimizer/checkpoint state, and runtime scheduling are correct for the supported TP/PP/DP/EP/CP topology. Base-model parameters remain frozen, but every trainable LoRA projection and every tenant update must be independent and reproducible. + +# Contract + +- Rust owns configuration, scheduling, data partitioning, checkpoint orchestration, and transport. C++/CUDA owns all training-path math and collectives. +- No Python JIT or ad-hoc dependency builds. A dependency is used only after a compatible prebuilt package is found and verified on the target runtime. +- TP and EP communicators are distinct. EP routed-output collectives must never be reused for LoRA TP output reduction. +- A global LoRA rank is partitioned over TP ranks (`rank % tp_size == 0`); rank-local A/B shapes and checkpoint metadata identify the topology. +- PP owns disjoint layer ranges and a real microbatch schedule. CP owns sequence/KV exchange; configuration fields alone do not satisfy either contract. +- Each dynamic adapter has independent optimizer step, m/v, gradient accumulation state, and checkpoint metadata. A request may train a selected subset of tenants without mutating others. + +# Requirements + +1. Correct Qwen3.5 full-attention and Qwen3.6 hybrid GDN/full-attention forward/backward with LoRA target projections. +2. Correct LoRA-only TP, DP, EP, PP, and CP group semantics, with explicit unsupported combinations rejected before model allocation. +3. MoE grouped/fused dispatch and communication overlap where a verified prebuilt dependency exists; otherwise retain a correct native fallback and report the gap. +4. FP32 gradient accumulation and logical optimizer boundaries, including abort/zero behavior. +5. Dynamic multi-LoRA selected-tenant scheduling, independent optimizer clocks, and no cross-tenant parameter or state updates. +6. Rank-aware checkpoint save/load and resume equivalence. +7. Target-machine smoke, numerical parity, throughput, peak-memory, and scaling evidence against a Megatron-LM reference under matched settings. + +# Out Of Scope + +Full-parameter training, unfrozen base weights, unverified third-party packages, Python JIT extensions, and claiming Megatron parity from topology metadata or unit tests alone. + +# Acceptance Evidence + +| Criterion | Direct evidence | +| --- | --- | +| TP LoRA rank sharding | two-rank native smoke checks local A/B shapes, TP delta all-reduce, and finite loss/update | +| DP/EP separation | multi-rank smoke checks replicated LoRA reduction and local expert ownership separately | +| PP/CP | multi-rank stage/ring smoke with layer ownership and sequence/KV exchange assertions | +| Tenant independence | selected-adapter test proves untouched tenant tensors, m/v, and step are unchanged | +| Accumulation | FP32 accumulation test matches concatenated-batch reference and abort clears pending state | +| Checkpoint resume | save/load continuation matches uninterrupted optimizer state and bias correction | +| Performance | target-machine matched benchmark records tokens/s, step time, peak memory, and communication share | + +# Dependency Constraints + +The target host currently exposes PyTorch 2.5.1+cu121 with ABI0 and no importable Transformer Engine, flash-attn, DeepEP, or Triton package. Public prebuilt availability must be rechecked before any dependency change; no package is added based on an unavailable search backend. + +# Known Baseline + +The repository currently has a working native EP/dense-DP subset, grouped MoE microbenchmark evidence, and a separate 5D topology helper. These are prerequisites, not proof of the target outcome. From e75b3c51e085da1a0a5274d9cfa193ba25a26a89 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 05:49:42 +0800 Subject: [PATCH 011/156] feat: add latent-rank tensor parallel lora --- .../kernels/qwen3_6_kernels.cpp | 289 ++++++++++++++---- crates/rustrain-qwen3-6/src/kernel.rs | 2 +- crates/rustrain-qwen3-6/src/lora.rs | 10 +- crates/rustrain-qwen3-6/src/session.rs | 51 +++- .../rustrain-qwen3-6/tests/native_smoke.cpp | 31 +- crates/rustrain-server/src/session.rs | 33 +- docs/plans/qwen-lora-megatron-progress.md | 11 +- 7 files changed, 331 insertions(+), 96 deletions(-) diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 94963ead..cf2c2774 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -40,9 +40,12 @@ struct TrainingContext; // forward declaration (defined below) struct LayerConfig; +static int64_t g_context_sequence = 0; // Forward declarations for functions defined after TrainingContext at::Tensor apply_multi_lora(TrainingContext* ctx, int64_t layer_idx, int64_t pair_idx, const at::Tensor& base_weight); +static at::Tensor tp_allreduce_lora_delta( + TrainingContext* ctx, const at::Tensor& local_delta); // ────────────────────────────────────────────────────────────────────── // Forward declarations @@ -868,16 +871,17 @@ static at::Tensor dense_mlp_forward_batched( const at::Tensor& down_proj, at::ScalarType compute_type); static at::Tensor lora_activation_delta( - const at::Tensor& x, const at::Tensor& A, const at::Tensor& B, + TrainingContext* ctx, const at::Tensor& x, + const at::Tensor& A, const at::Tensor& B, const at::Tensor& scaling); static at::Tensor add_batched_lora( - const at::Tensor& base, const at::Tensor& input, + TrainingContext* ctx, const at::Tensor& base, const at::Tensor& input, const LoraBatchEntry* entry ) { if (!entry) return base; return base + lora_activation_delta( - input, entry->a_stack, entry->b_stack, entry->scaling); + ctx, input, entry->a_stack, entry->b_stack, entry->scaling); } // Per-token routed-expert LoRA. Dynamic adapters add a leading sample axis to @@ -886,9 +890,11 @@ static at::Tensor add_batched_lora( // index_select + bmm operations select the correct adapter and expert without // materializing a full-rank delta weight. static at::Tensor dynamic_expert_lora_delta( + TrainingContext* ctx, const at::Tensor& input, const at::Tensor& token_indices, const at::Tensor& local_expert_indices, + int64_t batch, int64_t seq, const LoraBatchEntry* entry ) { @@ -898,18 +904,31 @@ static at::Tensor dynamic_expert_lora_delta( const int64_t local_experts = entry->a_stack.size(1); auto sample_indices = at::floor_divide(token_indices, seq); auto pair_indices = sample_indices * local_experts + local_expert_indices; - auto a = entry->a_stack.flatten(0, 1) + auto a_stack = entry->a_stack; + auto b_stack = entry->b_stack; + if (a_stack.size(0) == 1 && batch > 1) { + a_stack = a_stack.expand( + {batch, a_stack.size(1), a_stack.size(2), a_stack.size(3)}); + b_stack = b_stack.expand( + {batch, b_stack.size(1), b_stack.size(2), b_stack.size(3)}); + } + auto a = a_stack.flatten(0, 1) .index_select(0, pair_indices).to(input.scalar_type()); - auto b = entry->b_stack.flatten(0, 1) + auto b = b_stack.flatten(0, 1) .index_select(0, pair_indices).to(input.scalar_type()); auto low_rank = at::bmm(a, input.unsqueeze(-1)).squeeze(-1); auto delta = at::bmm(b, low_rank.unsqueeze(-1)).squeeze(-1); - auto scaling = entry->scaling.index_select(0, sample_indices) + auto scaling_stack = entry->scaling; + if (scaling_stack.size(0) == 1 && batch > 1) { + scaling_stack = scaling_stack.expand({batch, 1, 1}); + } + auto scaling = scaling_stack.index_select(0, sample_indices) .reshape({-1, 1}).to(input.scalar_type()); - return delta * scaling; + return tp_allreduce_lora_delta(ctx, delta * scaling); } static at::Tensor moe_forward( + TrainingContext* training_ctx, void* nccl_comm_v, void* nccl_stream_v, const at::Tensor& hidden, const at::Tensor& gate_w, const at::Tensor& shared_expert_gate_w, @@ -1036,12 +1055,13 @@ static at::Tensor moe_forward( selected, expert_lora.gate_up_a->transpose(1, 2), offsets); auto delta = at::_grouped_mm( low_rank, expert_lora.gate_up_b->transpose(1, 2), offsets); - gu = gu + delta * expert_lora.scaling; + gu = gu + tp_allreduce_lora_delta( + training_ctx, delta * expert_lora.scaling); } if (expert_gate_up_lora) { gu = gu + dynamic_expert_lora_delta( - selected, token_indices, local_expert_indices, - seq, expert_gate_up_lora); + training_ctx, selected, token_indices, local_expert_indices, + batch, seq, expert_gate_up_lora); } auto activated = fused_swiglu_op( gu.narrow(-1, 0, intermediate), @@ -1053,12 +1073,13 @@ static at::Tensor moe_forward( activated, expert_lora.down_a->transpose(1, 2), offsets); auto delta = at::_grouped_mm( low_rank, expert_lora.down_b->transpose(1, 2), offsets); - expert_out = expert_out + delta * expert_lora.scaling; + expert_out = expert_out + tp_allreduce_lora_delta( + training_ctx, delta * expert_lora.scaling); } if (expert_down_lora) { expert_out = expert_out + dynamic_expert_lora_delta( - activated, token_indices, local_expert_indices, - seq, expert_down_lora); + training_ctx, activated, token_indices, local_expert_indices, + batch, seq, expert_down_lora); } auto weights = gathered_weights.narrow( 0, local_start, local_tokens); @@ -1089,12 +1110,14 @@ static at::Tensor moe_forward( if (expert_lora.gate_up_a && expert_lora.gate_up_b) { auto a = expert_lora.gate_up_a->select(0, e_local); auto b = expert_lora.gate_up_b->select(0, e_local); - gu = gu + at::matmul(at::matmul(selected, a.t()), b.t()) * expert_lora.scaling; + auto delta = at::matmul(at::matmul(selected, a.t()), b.t()) + * expert_lora.scaling; + gu = gu + tp_allreduce_lora_delta(training_ctx, delta); } if (expert_gate_up_lora) { gu = gu + dynamic_expert_lora_delta( - selected, token_indices, local_expert_indices, - seq, expert_gate_up_lora); + training_ctx, selected, token_indices, local_expert_indices, + batch, seq, expert_gate_up_lora); } auto gate_part = gu.narrow(-1, 0, intermediate); auto up_part = gu.narrow(-1, intermediate, intermediate); @@ -1103,13 +1126,15 @@ static at::Tensor moe_forward( if (expert_lora.down_a && expert_lora.down_b) { auto a = expert_lora.down_a->select(0, e_local); auto b = expert_lora.down_b->select(0, e_local); + auto delta = at::matmul(at::matmul(activated, a.t()), b.t()) + * expert_lora.scaling; expert_out = expert_out - + at::matmul(at::matmul(activated, a.t()), b.t()) * expert_lora.scaling; + + tp_allreduce_lora_delta(training_ctx, delta); } if (expert_down_lora) { expert_out = expert_out + dynamic_expert_lora_delta( - activated, token_indices, local_expert_indices, - seq, expert_down_lora); + training_ctx, activated, token_indices, local_expert_indices, + batch, seq, expert_down_lora); } auto weights = gathered_weights.narrow(0, offset, n_tokens); routed_output = routed_output.index_add_(0, token_indices, expert_out * weights); @@ -1161,12 +1186,14 @@ static at::Tensor moe_forward( auto shared_up = at::matmul(flat, shared_up_proj.t()); if (shared_gate_lora) { shared_gate = add_batched_lora( - shared_gate.reshape({batch, seq, -1}), hidden, shared_gate_lora) + training_ctx, shared_gate.reshape({batch, seq, -1}), hidden, + shared_gate_lora) .reshape({batch * seq, -1}); } if (shared_up_lora) { shared_up = add_batched_lora( - shared_up.reshape({batch, seq, -1}), hidden, shared_up_lora) + training_ctx, shared_up.reshape({batch, seq, -1}), hidden, + shared_up_lora) .reshape({batch * seq, -1}); } auto shared_hidden = fused_swiglu_op( @@ -1175,7 +1202,8 @@ static at::Tensor moe_forward( auto shared_out = at::matmul(shared_hidden.reshape({batch * seq, -1}), shared_down_proj.t()); if (shared_down_lora) { shared_out = add_batched_lora( - shared_out.reshape({batch, seq, -1}), shared_hidden, shared_down_lora) + training_ctx, shared_out.reshape({batch, seq, -1}), shared_hidden, + shared_down_lora) .reshape({batch * seq, -1}); } auto seg = at::sigmoid(at::matmul(flat, shared_expert_gate_w.t())).to(compute_type); @@ -1333,7 +1361,7 @@ static at::Tensor forward_single_layer( auto shared_down = use_batched ? *w[12] : apply_multi_lora(ctx, layer_idx, shared_down_pair, *w[12]); auto expert_lora = routed_expert_lora(ctx, layer_idx, *cfg); - auto mlp_out = moe_forward(cfg->nccl_comm, cfg->nccl_stream, post_attn, + auto mlp_out = moe_forward(ctx, cfg->nccl_comm, cfg->nccl_stream, post_attn, *w[8], *w[9], shared_gate, shared_up, shared_down, *w[13], *w[14], expert_lora, cfg->num_experts, cfg->top_k, cfg->moe_intermediate, @@ -1395,7 +1423,7 @@ static at::Tensor forward_single_layer( auto shared_down = use_batched ? *w[15] : apply_multi_lora(ctx, layer_idx, shared_down_pair, *w[15]); auto expert_lora = routed_expert_lora(ctx, layer_idx, *cfg); - auto mlp_out = moe_forward(cfg->nccl_comm, cfg->nccl_stream, post_attn, + auto mlp_out = moe_forward(ctx, cfg->nccl_comm, cfg->nccl_stream, post_attn, *w[11], *w[12], shared_gate, shared_up, shared_down, *w[16], *w[17], expert_lora, cfg->num_experts, cfg->top_k, cfg->moe_intermediate, @@ -1449,6 +1477,9 @@ struct AdamDevBuffers { }; struct TrainingContext { + // All ranks create contexts in the same order. This sequence isolates + // per-context rendezvous files when a worker reuses one NCCL process. + int64_t context_sequence = 0; // Model weights (frozen, no grad) — pointers to external tensors std::vector weight_ptrs; // flat array, 15 or 18 per layer std::vector embed_ptr; @@ -1532,10 +1563,42 @@ struct TrainingContext { int ep_world_size = 1; int ep_rank = 0; bool data_parallel = false; + // LoRA-only tensor parallelism keeps frozen base weights replicated and + // shards the latent rank. Only the local LoRA delta uses this communicator. + ncclComm_t tp_comm = nullptr; + cudaStream_t tp_stream = nullptr; + int tp_world_size = 1; + int tp_rank = 0; int cuda_device = 0; // ────────────────────────────────────────────────────────────────────── }; +static at::Tensor tp_allreduce_lora_delta( + TrainingContext* ctx, const at::Tensor& local_delta +) { + if (!ctx || ctx->tp_world_size <= 1) return local_delta; + TORCH_CHECK(ctx->tp_comm, + "LoRA TP communicator is not initialized for TP_SIZE=", + ctx->tp_world_size); + return NcclAllReduceFunction::apply( + local_delta, (int64_t)ctx->tp_comm, + (int64_t)reinterpret_cast(ctx->tp_stream)); +} + +static at::Tensor initialize_lora_a( + TrainingContext* ctx, const at::TensorOptions& options, + int64_t experts, int64_t global_rank, int64_t in_features +) { + const int64_t local_rank = global_rank / ctx->tp_world_size; + const int64_t rank_start = ctx->tp_rank * local_rank; + if (experts > 0) { + auto global = at::randn({experts, global_rank, in_features}, options); + return global.narrow(1, rank_start, local_rank).contiguous() * 0.01; + } + auto global = at::randn({global_rank, in_features}, options); + return global.narrow(0, rank_start, local_rank).contiguous() * 0.01; +} + static const char* lora_pair_name(const LayerConfig& cfg, int64_t pair_idx) { auto table = lora_projection_table(cfg); TORCH_CHECK(pair_idx >= 0 && pair_idx < table.count, "invalid LoRA projection index"); @@ -1842,10 +1905,35 @@ static void prepare_lora_batch(TrainingContext* ctx) { ctx->lora_batch_valid = true; } +/// Build activation-level entries for the fixed adapter in TP mode. The +/// singleton adapter dimension is expanded lazily to the input batch by +/// lora_activation_delta/dynamic_expert_lora_delta. +static void prepare_fixed_lora_batch(TrainingContext* ctx) { + ctx->lora_batch_cache.clear(); + ctx->lora_batch_n = 1; + for (int64_t layer_idx = 0; layer_idx < ctx->num_layers; ++layer_idx) { + const int64_t pair_count = lora_pair_count(ctx->layer_configs[layer_idx]); + const int64_t offset = ctx->lora_layer_offset[layer_idx]; + for (int64_t pair_idx = 0; pair_idx < pair_count; ++pair_idx) { + const int64_t slot = offset + pair_idx; + if (!legacy_lora_slot_active(ctx, slot)) continue; + auto& a = ctx->lora_a[slot]; + auto& b = ctx->lora_b[slot]; + auto scaling = at::full( + {1, 1, 1}, ctx->lora_scaling, + at::TensorOptions().dtype(a.scalar_type()).device(a.device())); + ctx->lora_batch_cache[lora_cache_key(layer_idx, pair_idx)] = { + a.unsqueeze(0), b.unsqueeze(0), scaling}; + } + } + ctx->lora_batch_valid = true; +} + /// Compute B@(A@x) * scaling — never materializes B@A. /// x: [N, seq, in], A: [N, rank, in], B: [N, out, rank], scaling: [N, 1, 1] /// returns: [N, seq, out] static at::Tensor lora_activation_delta( + TrainingContext* ctx, const at::Tensor& x, // [N, seq, in] const at::Tensor& A, // [N, rank, in] const at::Tensor& B, // [N, out, rank] @@ -1856,11 +1944,16 @@ static at::Tensor lora_activation_delta( auto A_c = A.to(kind); auto B_c = B.to(kind); auto s_c = scaling.to(kind); + if (A_c.size(0) == 1 && x.size(0) > 1) { + A_c = A_c.expand({x.size(0), A_c.size(1), A_c.size(2)}); + B_c = B_c.expand({x.size(0), B_c.size(1), B_c.size(2)}); + s_c = s_c.expand({x.size(0), 1, 1}); + } // Ax = A @ x^T → [N, rank, seq] auto Ax = at::bmm(A_c, x.transpose(-2, -1)); // delta = B @ Ax → [N, out, seq] → transpose → [N, seq, out] auto delta = at::bmm(B_c, Ax).transpose(-2, -1); - return delta * s_c; + return tp_allreduce_lora_delta(ctx, delta * s_c); } static const LoraBatchEntry* lora_batch_entry( @@ -1884,13 +1977,13 @@ static at::Tensor dense_mlp_forward_batched( auto gate_out = at::matmul(hidden, gate_proj.t()); auto up_out = at::matmul(hidden, up_proj.t()); gate_out = add_batched_lora( - gate_out, hidden, lora_batch_entry(ctx, layer_idx, gate_pair)); + ctx, gate_out, hidden, lora_batch_entry(ctx, layer_idx, gate_pair)); up_out = add_batched_lora( - up_out, hidden, lora_batch_entry(ctx, layer_idx, up_pair)); + ctx, up_out, hidden, lora_batch_entry(ctx, layer_idx, up_pair)); auto activated = fused_swiglu_op(gate_out, up_out, 0.0); auto result = at::matmul(activated, down_proj.t()); result = add_batched_lora( - result, activated, lora_batch_entry(ctx, layer_idx, down_pair)); + ctx, result, activated, lora_batch_entry(ctx, layer_idx, down_pair)); return result.to(compute_type); } @@ -2014,7 +2107,7 @@ at::Tensor compute_mlp_only( : apply_multi_lora(ctx, layer_idx, shared_down_pair, *ctx->weight_ptrs[w_offset+mlp_start+4]); auto expert_lora = routed_expert_lora(ctx, layer_idx, cfg); - return moe_forward(cfg.nccl_comm, cfg.nccl_stream, post_attn, + return moe_forward(ctx, cfg.nccl_comm, cfg.nccl_stream, post_attn, *ctx->weight_ptrs[w_offset+mlp_start], *ctx->weight_ptrs[w_offset+mlp_start+1], shared_gate, shared_up, shared_down, *ctx->weight_ptrs[w_offset+mlp_start+5], @@ -2093,6 +2186,7 @@ struct SubLayerCkpt : public torch::autograd::Function { tc->lora_batch_valid = false; tc->lora_cache_valid = false; if (!tc->adapters.empty()) prepare_lora_batch(tc); + else if (tc->tp_world_size > 1) prepare_fixed_lora_batch(tc); else precompute_lora_cache(tc); auto output = is_attn ? compute_attn_only(tc, input, layer, tc->compute_type) : compute_mlp_only(tc, input, layer, tc->compute_type); @@ -2207,15 +2301,15 @@ static at::Tensor full_attention_batched( // Apply activation-level LoRA: q += B@(A@hidden) * scaling auto it_q = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 0)); if (it_q != ctx->lora_batch_cache.end()) { - q = q + lora_activation_delta(hidden, it_q->second.a_stack, it_q->second.b_stack, it_q->second.scaling); + q = q + lora_activation_delta(ctx, hidden, it_q->second.a_stack, it_q->second.b_stack, it_q->second.scaling); } auto it_k = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 1)); if (it_k != ctx->lora_batch_cache.end()) { - k = k + lora_activation_delta(hidden, it_k->second.a_stack, it_k->second.b_stack, it_k->second.scaling); + k = k + lora_activation_delta(ctx, hidden, it_k->second.a_stack, it_k->second.b_stack, it_k->second.scaling); } auto it_v = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 2)); if (it_v != ctx->lora_batch_cache.end()) { - v = v + lora_activation_delta(hidden, it_v->second.a_stack, it_v->second.b_stack, it_v->second.scaling); + v = v + lora_activation_delta(ctx, hidden, it_v->second.a_stack, it_v->second.b_stack, it_v->second.scaling); } // Reshape Q: [batch, seq, num_heads, head_dim*2] → split into q and gate @@ -2284,7 +2378,7 @@ static at::Tensor full_attention_batched( // Apply LoRA delta on o_proj output auto it_o = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 3)); if (it_o != ctx->lora_batch_cache.end()) { - result = result + lora_activation_delta(attn_flat, + result = result + lora_activation_delta(ctx, attn_flat, it_o->second.a_stack, it_o->second.b_stack, it_o->second.scaling); } return result; @@ -2311,7 +2405,7 @@ static at::Tensor linear_attention_batched( auto qkv = at::matmul(hidden, in_proj_qkv.t()); auto it_qkv = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 0)); if (it_qkv != ctx->lora_batch_cache.end()) { - qkv = qkv + lora_activation_delta(hidden, it_qkv->second.a_stack, it_qkv->second.b_stack, it_qkv->second.scaling); + qkv = qkv + lora_activation_delta(ctx, hidden, it_qkv->second.a_stack, it_qkv->second.b_stack, it_qkv->second.scaling); } // DIAG: dump after QKV projection @@ -2383,18 +2477,18 @@ static at::Tensor linear_attention_batched( auto b = at::matmul(hidden, in_proj_b.t()); auto it_a = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 2)); if (it_a != ctx->lora_batch_cache.end()) { - a = a + lora_activation_delta(hidden, it_a->second.a_stack, it_a->second.b_stack, it_a->second.scaling); + a = a + lora_activation_delta(ctx, hidden, it_a->second.a_stack, it_a->second.b_stack, it_a->second.scaling); } auto it_b = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 3)); if (it_b != ctx->lora_batch_cache.end()) { - b = b + lora_activation_delta(hidden, it_b->second.a_stack, it_b->second.b_stack, it_b->second.scaling); + b = b + lora_activation_delta(ctx, hidden, it_b->second.a_stack, it_b->second.b_stack, it_b->second.scaling); } // Z projection + LoRA delta auto z = at::matmul(hidden, in_proj_z.t()); auto it_z = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 1)); if (it_z != ctx->lora_batch_cache.end()) { - z = z + lora_activation_delta(hidden, it_z->second.a_stack, it_z->second.b_stack, it_z->second.scaling); + z = z + lora_activation_delta(ctx, hidden, it_z->second.a_stack, it_z->second.b_stack, it_z->second.scaling); } z = z.reshape({batch, seq, num_v_heads, head_v_dim}); @@ -2506,7 +2600,7 @@ static at::Tensor linear_attention_batched( // out_proj LoRA delta: result += B@(A@gated) * scaling auto it_op = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 4)); if (it_op != ctx->lora_batch_cache.end()) { - result = result + lora_activation_delta(gated, it_op->second.a_stack, it_op->second.b_stack, it_op->second.scaling); + result = result + lora_activation_delta(ctx, gated, it_op->second.a_stack, it_op->second.b_stack, it_op->second.scaling); } return result; @@ -2532,7 +2626,8 @@ static at::Tensor forward_full( TrainingContext* ctx, const at::Tensor& input_ids ) { - if (ctx->lora_batch_valid) prepare_lora_batch(ctx); + if (!ctx->adapters.empty()) prepare_lora_batch(ctx); + else if (ctx->tp_world_size > 1) prepare_fixed_lora_batch(ctx); else precompute_lora_cache(ctx); auto kind = ctx->compute_type; auto embed = *ctx->embed_ptr[0]; @@ -2768,7 +2863,8 @@ static at::Tensor forward_full_fused( TrainingContext* ctx, const at::Tensor& input_ids ) { - if (ctx->lora_batch_valid) prepare_lora_batch(ctx); + if (!ctx->adapters.empty()) prepare_lora_batch(ctx); + else if (ctx->tp_world_size > 1) prepare_fixed_lora_batch(ctx); else precompute_lora_cache(ctx); auto embed = *ctx->embed_ptr[0]; at::Tensor hidden = at::embedding(embed, input_ids); @@ -2792,9 +2888,12 @@ static at::Tensor forward_full_checkpoint( TrainingContext* ctx, const at::Tensor& input_ids ) { - // Use batched path if multiple adapters, else legacy weight-level - if (ctx->lora_batch_valid) { + // Use activation-level paths for dynamic adapters and LoRA TP; the + // legacy weight cache remains the single-rank fast path. + if (!ctx->adapters.empty()) { prepare_lora_batch(ctx); + } else if (ctx->tp_world_size > 1) { + prepare_fixed_lora_batch(ctx); } else { precompute_lora_cache(ctx); } @@ -2895,6 +2994,7 @@ static void manual_group_backward( ctx->lora_batch_valid = false; ctx->lora_cache_valid = false; if (!ctx->adapters.empty()) prepare_lora_batch(ctx); + else if (ctx->tp_world_size > 1) prepare_fixed_lora_batch(ctx); else precompute_lora_cache(ctx); // Recompute forward with grad for this group only @@ -2906,7 +3006,7 @@ static void manual_group_backward( // LoRA params are shared across groups, so we accumulate their gradients. std::vector grad_inputs = {input}; - if (ctx->lora_batch_valid) { + if (!ctx->adapters.empty()) { // Multi-LoRA: collect A/B from ctx->adapters for (int64_t l = start; l < end; l++) { int64_t lora_count = lora_pair_count(ctx->layer_configs[l]); @@ -2945,7 +3045,7 @@ static void manual_group_backward( // Manually accumulate LoRA param gradients - if (ctx->lora_batch_valid) { + if (!ctx->adapters.empty()) { // Multi-LoRA: accumulate into ctx->adapters int64_t gi = 1; // skip input grad (index 0) for (int64_t l = start; l < end; l++) { @@ -3460,7 +3560,7 @@ static at::Tensor mtp_compute_loss( extern "C" { __attribute__((visibility("default"))) int64_t qwen36_kernel_abi_version() { - return 9; + return 10; } // Create training context — called once at startup @@ -3480,12 +3580,24 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( ) { try { auto* ctx = new TrainingContext(); + ctx->context_sequence = ++g_context_sequence; ctx->compute_type = static_cast(compute_type); ctx->lr = lr; ctx->beta1 = beta1; ctx->beta2 = beta2; ctx->eps = eps; ctx->vocab_size = vocab_size; ctx->rms_eps = rms_eps; ctx->step_count = 0; ctx->lora_scaling = lora_scaling; ctx->num_layers = num_layers; ctx->use_checkpoint = false; ctx->group_size = 4; + const char* tp_size_env = getenv("TP_SIZE"); + if (!tp_size_env) tp_size_env = getenv("RUSTRAIN_TP_SIZE"); + ctx->tp_world_size = tp_size_env ? atoi(tp_size_env) : 1; + TORCH_CHECK(ctx->tp_world_size > 0, "TP_SIZE must be positive"); + const char* rank_env = getenv("RANK"); + const int global_rank = rank_env ? atoi(rank_env) : 0; + ctx->tp_rank = global_rank % ctx->tp_world_size; + TORCH_CHECK(lora_rank > 0 && lora_rank % ctx->tp_world_size == 0, + "LoRA rank ", lora_rank, " must be divisible by TP_SIZE=", + ctx->tp_world_size); + const int64_t local_lora_rank = lora_rank / ctx->tp_world_size; if (const char* mtp_scale = getenv("QWEN36_MTP_LOSS_SCALE")) { ctx->mtp_loss_scale = std::strtod(mtp_scale, nullptr); } @@ -3591,12 +3703,12 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( } else if (projection.grouped_expert) { int64_t experts = base->size(0); int64_t out_f = base->size(1), in_f = base->size(2); - a = at::randn({experts, lora_rank, in_f}, opts) * 0.01; - b = at::zeros({experts, out_f, lora_rank}, opts); + a = initialize_lora_a(ctx, opts, experts, lora_rank, in_f); + b = at::zeros({experts, out_f, local_lora_rank}, opts); } else { int64_t out_f = base->size(0), in_f = base->size(1); - a = at::randn({lora_rank, in_f}, opts) * 0.01; - b = at::zeros({out_f, lora_rank}, opts); + a = initialize_lora_a(ctx, opts, 0, lora_rank, in_f); + b = at::zeros({out_f, local_lora_rank}, opts); } a.set_requires_grad(active); b.set_requires_grad(active); @@ -4074,7 +4186,8 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( cudaMemGetInfo(&free_mem, &total_mem); int64_t n_max; if (ctx->nccl_comm && ctx->ep_world_size > 1) { - const std::string sync_path = nccl_sync_dir() + "/nmax_sync.txt"; + const std::string sync_path = nccl_sync_dir() + "/nmax_sync_" + + std::to_string(ctx->context_sequence) + ".txt"; if (ctx->ep_rank == 0) { n_max = compute_n_max( (int64_t)free_mem, lora_rank, @@ -4399,6 +4512,8 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora_selected( // Process-level NCCL singleton — created once, reused across sessions. static ncclComm_t g_nccl_comm = nullptr; static cudaStream_t g_nccl_stream = nullptr; +static ncclComm_t g_tp_comm = nullptr; +static cudaStream_t g_tp_stream = nullptr; static bool g_nccl_initialized = false; static int g_cuda_device = 0; @@ -4436,8 +4551,25 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( if (world_str2) ctx->ep_world_size = atoi(world_str2); const char* local_rank_str2 = getenv("LOCAL_RANK"); ctx->cuda_device = local_rank_str2 ? atoi(local_rank_str2) : g_cuda_device; - for (auto& lc : ctx->layer_configs) { lc.nccl_comm = (void*)g_nccl_comm; lc.nccl_stream = (void*)g_nccl_stream; } - for (auto& lc : ctx->mtp_layer_configs) { lc.nccl_comm = (void*)g_nccl_comm; lc.nccl_stream = (void*)g_nccl_stream; } + const char* tp_size_str2 = getenv("TP_SIZE"); + if (!tp_size_str2) tp_size_str2 = getenv("RUSTRAIN_TP_SIZE"); + ctx->tp_world_size = tp_size_str2 ? atoi(tp_size_str2) : 1; + ctx->tp_rank = ctx->tp_world_size > 0 + ? ctx->ep_rank % ctx->tp_world_size : 0; + ctx->tp_comm = ctx->tp_world_size > 1 ? g_tp_comm : nullptr; + ctx->tp_stream = ctx->tp_world_size > 1 ? g_tp_stream : nullptr; + // In TP-only mode the parent communicator is reserved for the TP + // split; EP layer collectives must remain disabled on replicated MoE. + if (ctx->tp_world_size <= 1) { + for (auto& lc : ctx->layer_configs) { + lc.nccl_comm = (void*)g_nccl_comm; + lc.nccl_stream = (void*)g_nccl_stream; + } + for (auto& lc : ctx->mtp_layer_configs) { + lc.nccl_comm = (void*)g_nccl_comm; + lc.nccl_stream = (void*)g_nccl_stream; + } + } return 0; } @@ -4544,6 +4676,27 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( // Store as process-level singleton g_nccl_comm = comm; g_nccl_stream = nccl_stream; + const char* tp_size_str = getenv("TP_SIZE"); + if (!tp_size_str) tp_size_str = getenv("RUSTRAIN_TP_SIZE"); + const int tp_size = tp_size_str ? atoi(tp_size_str) : 1; + if (tp_size <= 0 || world_size % tp_size != 0) { + fprintf(stderr, "[tp_nccl] invalid TP_SIZE=%d for WORLD_SIZE=%d\n", + tp_size, world_size); + return -1; + } + if (tp_size > 1) { + // The default rank order makes TP the least-significant axis. A + // future multi-axis implementation must pass an explicit color/key + // mapping instead of reusing this world-rank split. + ncclResult_t tp_err = ncclCommSplit( + comm, rank / tp_size, rank % tp_size, &g_tp_comm, nullptr); + if (tp_err != ncclSuccess) { + fprintf(stderr, "[tp_nccl] ncclCommSplit failed: %d (%s)\n", + tp_err, ncclGetErrorString(tp_err)); + return -1; + } + g_tp_stream = nccl_stream; + } g_nccl_initialized = true; ctx->nccl_comm = comm; @@ -4551,15 +4704,21 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( ctx->ep_rank = rank; ctx->ep_world_size = world_size; ctx->data_parallel = env_enabled("RUSTRAIN_DATA_PARALLEL"); + ctx->tp_world_size = tp_size; + ctx->tp_rank = rank % tp_size; + ctx->tp_comm = tp_size > 1 ? g_tp_comm : nullptr; + ctx->tp_stream = tp_size > 1 ? g_tp_stream : nullptr; // Propagate to layer configs - for (auto& lc : ctx->layer_configs) { - lc.nccl_comm = (void*)comm; - lc.nccl_stream = (void*)nccl_stream; - } - for (auto& lc : ctx->mtp_layer_configs) { - lc.nccl_comm = (void*)comm; - lc.nccl_stream = (void*)nccl_stream; + if (tp_size <= 1) { + for (auto& lc : ctx->layer_configs) { + lc.nccl_comm = (void*)comm; + lc.nccl_stream = (void*)nccl_stream; + } + for (auto& lc : ctx->mtp_layer_configs) { + lc.nccl_comm = (void*)comm; + lc.nccl_stream = (void*)nccl_stream; + } } return 0; @@ -4632,6 +4791,10 @@ int64_t qwen36_add_lora( try { auto* ctx = reinterpret_cast(ctx_ptr); TORCH_CHECK(rank > 0, "LoRA rank must be positive"); + TORCH_CHECK(rank % ctx->tp_world_size == 0, + "dynamic LoRA rank ", rank, " must be divisible by TP_SIZE=", + ctx->tp_world_size); + const int64_t local_rank = rank / ctx->tp_world_size; TORCH_CHECK(alpha > 0.0, "LoRA alpha must be positive"); TrainingContext::LoRAAdapter adapter; adapter.id = ++ctx->next_adapter_id; @@ -4720,14 +4883,14 @@ int64_t qwen36_add_lora( projection.name); int64_t experts = base->size(0); int64_t out_f = base->size(1), in_f = base->size(2); - a = at::randn({experts, rank, in_f}, opts) * 0.01; - b = at::zeros({experts, out_f, rank}, opts); + a = initialize_lora_a(ctx, opts, experts, rank, in_f); + b = at::zeros({experts, out_f, local_rank}, opts); } else { TORCH_CHECK(base->dim() == 2, "dynamic LoRA projection must be a matrix: ", projection.name); int64_t out_f = base->size(0), in_f = base->size(1); - a = at::randn({rank, in_f}, opts) * 0.01; - b = at::zeros({out_f, rank}, opts); + a = initialize_lora_a(ctx, opts, 0, rank, in_f); + b = at::zeros({out_f, local_rank}, opts); } } else { a = at::zeros({}, opts); diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs index c9e7ff98..42f33e1a 100644 --- a/crates/rustrain-qwen3-6/src/kernel.rs +++ b/crates/rustrain-qwen3-6/src/kernel.rs @@ -190,7 +190,7 @@ unsafe fn load_kernels() -> Option { }}; } let abi_version: FnKernelAbiVersion = sym!("qwen36_kernel_abi_version"); - if abi_version() != 9 { + if abi_version() != 10 { return None; } Some(KernelHandles { diff --git a/crates/rustrain-qwen3-6/src/lora.rs b/crates/rustrain-qwen3-6/src/lora.rs index a3396164..9d6d70ed 100644 --- a/crates/rustrain-qwen3-6/src/lora.rs +++ b/crates/rustrain-qwen3-6/src/lora.rs @@ -967,9 +967,17 @@ mod tests { "k_proj", "v_proj", "o_proj", + "gate_proj", + "up_proj", + "down_proj", "in_proj_qkv", "in_proj_z", - "out_proj" + "in_proj_a", + "in_proj_b", + "out_proj", + "gate_proj", + "up_proj", + "down_proj" ] ); } diff --git a/crates/rustrain-qwen3-6/src/session.rs b/crates/rustrain-qwen3-6/src/session.rs index cc23a2ba..04513ee5 100644 --- a/crates/rustrain-qwen3-6/src/session.rs +++ b/crates/rustrain-qwen3-6/src/session.rs @@ -262,13 +262,48 @@ fn train_impl( .unwrap_or(0); let world_size = shard_ref.map(|s| s.world_size).unwrap_or(env_world_size); let rank = shard_ref.map(|s| s.rank).unwrap_or(env_rank); - let is_data_parallel = !is_ep && world_size > 1; + let tp_size = config.parallel.tensor_model_parallel_size; + let env_tp_size = std::env::var("TP_SIZE") + .or_else(|_| std::env::var("RUSTRAIN_TP_SIZE")) + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(1); + if env_tp_size > 1 && env_tp_size != tp_size { + bail!( + "TP_SIZE environment ({env_tp_size}) does not match config tensor_model_parallel_size ({tp_size})" + ); + } + if tp_size > 1 { + if world_size != tp_size + || config.parallel.pipeline_model_parallel_size != 1 + || config.parallel.data_parallel_size != 1 + || config.parallel.expert_model_parallel_size != 1 + || config.parallel.context_parallel_size != 1 + { + bail!( + "native Qwen LoRA currently supports TP-only topology: TP={} WORLD_SIZE={} PP={} DP={} EP={} CP={}", + tp_size, + world_size, + config.parallel.pipeline_model_parallel_size, + config.parallel.data_parallel_size, + config.parallel.expert_model_parallel_size, + config.parallel.context_parallel_size + ); + } + if lora_config.rank % tp_size as i64 != 0 { + bail!("LoRA rank {} must be divisible by TP_SIZE={tp_size}", lora_config.rank); + } + unsafe { + std::env::set_var("TP_SIZE", tp_size.to_string()); + } + } + let is_data_parallel = !is_ep && world_size > 1 && tp_size == 1; if is_data_parallel && runtime_config.is_moe { bail!( "replicated Qwen data parallelism is only supported for dense/linear-attention models; use *_ep for MoE" ); } - if is_data_parallel { + if is_data_parallel || tp_size > 1 { crate::kernel::CppTrainingContext::set_cuda_device( std::env::var("LOCAL_RANK") .ok() @@ -367,6 +402,10 @@ fn train_impl( ); } + let expert_start = shard_ref.map(|s| s.expert_start).unwrap_or(0); + let expert_count = shard_ref + .map(|s| s.experts_per_rank) + .unwrap_or(runtime_config.num_experts); let ctx = crate::kernel::CppTrainingContext::new( &weights_gpu, &runtime_config, @@ -379,8 +418,8 @@ fn train_impl( lora_config.rank as i64, &lora_config.target_layers, &lora_config.target_modules, - shard_ref.map(|s| s.expert_start).unwrap_or(0), - shard_ref.map(|s| s.experts_per_rank).unwrap_or(0), + expert_start, + expert_count, )?; if world_size > 1 { @@ -409,8 +448,8 @@ fn train_impl( ctx.set_mtp_weights( &weights_gpu, &runtime_config, - shard_ref.map(|s| s.expert_start).unwrap_or(0), - shard_ref.map(|s| s.experts_per_rank).unwrap_or(0), + expert_start, + expert_count, )?; info!( "C++ TrainingContext: MTP weights set ({} layers)", diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index 4747bb64..141ea426 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -54,11 +54,13 @@ static at::Tensor cuda_rand(std::initializer_list shape) { } int main() { - assert(qwen36_kernel_abi_version() == 9); + assert(qwen36_kernel_abi_version() == 10); const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); const int process_rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0"); const int local_rank = std::atoi(std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); assert(world == 1 || (world == 2 && process_rank >= 0 && process_rank < world)); + const int tp_size = std::atoi(std::getenv("TP_SIZE") ? std::getenv("TP_SIZE") : "1"); + assert(tp_size == 1 || (tp_size == world && tp_size == 2)); c10::cuda::CUDAGuard guard(local_rank); at::manual_seed(7); @@ -68,6 +70,7 @@ int main() { constexpr int64_t head_dim = 8; constexpr int64_t intermediate = 8; constexpr int64_t rank = 8; + const int64_t local_lora_rank = rank / tp_size; // One full-attention MoE layer. Shapes intentionally match the native // weight order used by build_weight_ptrs/kernel.cpp. @@ -130,6 +133,7 @@ int main() { 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank, &target_layer, 1, "experts_gate_up_proj,experts_down_proj"); if (!ctx) return 2; + if (world > 1) assert(qwen36_init_nccl(ctx) == 0); const int64_t count = qwen36_get_lora_count(ctx); assert(count == 9); @@ -138,10 +142,10 @@ int main() { auto* down_a = reinterpret_cast(qwen36_get_lora_a(ctx, 8)); auto* down_b = reinterpret_cast(qwen36_get_lora_b(ctx, 8)); assert(expert_a && expert_b && down_a && down_b); - assert(expert_a->sizes() == at::IntArrayRef({experts, rank, hidden})); - assert(expert_b->sizes() == at::IntArrayRef({experts, 2 * intermediate, rank})); - assert(down_a->sizes() == at::IntArrayRef({experts, rank, intermediate})); - assert(down_b->sizes() == at::IntArrayRef({experts, hidden, rank})); + assert(expert_a->sizes() == at::IntArrayRef({experts, local_lora_rank, hidden})); + assert(expert_b->sizes() == at::IntArrayRef({experts, 2 * intermediate, local_lora_rank})); + assert(down_a->sizes() == at::IntArrayRef({experts, local_lora_rank, intermediate})); + assert(down_b->sizes() == at::IntArrayRef({experts, hidden, local_lora_rank})); // Make both B tensors nonzero so the step exercises the LoRA branches. auto expert_b_value = at::ones(expert_b->sizes(), expert_b->options()); @@ -210,7 +214,7 @@ int main() { auto* dynamic_b = reinterpret_cast( qwen36_get_adapter_lora_tensor( ctx, adapter_one, 0, "shared_gate_proj", 1)); - assert(dynamic_b && dynamic_b->sizes() == at::IntArrayRef({intermediate, rank})); + assert(dynamic_b && dynamic_b->sizes() == at::IntArrayRef({intermediate, local_lora_rank})); auto dynamic_b_value = at::ones(dynamic_b->sizes(), dynamic_b->options()); assert(qwen36_set_adapter_lora_tensor( ctx, adapter_one, 0, "shared_gate_proj", 1, &dynamic_b_value) == 0); @@ -227,7 +231,7 @@ int main() { qwen36_get_adapter_lora_tensor( ctx, adapter_one, 0, "experts_gate_up_proj", 1)); assert(dynamic_expert_b && dynamic_expert_b->sizes() == - at::IntArrayRef({experts, 2 * intermediate, rank})); + at::IntArrayRef({experts, 2 * intermediate, local_lora_rank})); auto dynamic_expert_b_value = at::ones( dynamic_expert_b->sizes(), dynamic_expert_b->options()); assert(qwen36_set_adapter_lora_tensor( @@ -330,9 +334,7 @@ int main() { 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank, &target_layer, 1, "q_proj"); assert(ctx); - if (world > 1) { - assert(qwen36_init_nccl(ctx) == 0); - } + if (world > 1) assert(qwen36_init_nccl(ctx) == 0); const char* dense_targets = "gate_proj,up_proj,down_proj"; const int64_t dense_one = qwen36_add_lora( ctx, rank, 1.0, &target_layer, 1, dense_targets); @@ -341,7 +343,7 @@ int main() { assert(dense_one > 0 && dense_two > dense_one); auto* dense_b = reinterpret_cast( qwen36_get_adapter_lora_tensor(ctx, dense_one, 0, "gate_proj", 1)); - assert(dense_b && dense_b->sizes() == at::IntArrayRef({intermediate, rank})); + assert(dense_b && dense_b->sizes() == at::IntArrayRef({intermediate, local_lora_rank})); auto dense_b_value = at::ones(dense_b->sizes(), dense_b->options()); assert(qwen36_set_adapter_lora_tensor( ctx, dense_one, 0, "gate_proj", 1, &dense_b_value) == 0); @@ -403,12 +405,13 @@ int main() { &target_layer, 1, "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj"); assert(ctx); + if (world > 1) assert(qwen36_init_nccl(ctx) == 0); assert(qwen36_get_lora_count(ctx) == 8); auto* linear_a = reinterpret_cast(qwen36_get_lora_a(ctx, 0)); auto* linear_b = reinterpret_cast(qwen36_get_lora_b(ctx, 0)); assert(linear_a && linear_b); - assert(linear_a->sizes() == at::IntArrayRef({rank, hidden})); - assert(linear_b->sizes() == at::IntArrayRef({linear_qkv, rank})); + assert(linear_a->sizes() == at::IntArrayRef({local_lora_rank, hidden})); + assert(linear_b->sizes() == at::IntArrayRef({linear_qkv, local_lora_rank})); auto linear_b_value = at::ones(linear_b->sizes(), linear_b->options()); assert(qwen36_set_lora_tensor(ctx, 0, 1, &linear_b_value) == 0); auto linear_a_before = linear_a->clone(); @@ -451,7 +454,7 @@ int main() { qwen36_get_adapter_lora_tensor( ctx, linear_adapter_one, 0, "in_proj_qkv", 1)); assert(dynamic_linear_b && dynamic_linear_b->sizes() == - at::IntArrayRef({linear_qkv, rank})); + at::IntArrayRef({linear_qkv, local_lora_rank})); auto dynamic_linear_b_value = at::ones( dynamic_linear_b->sizes(), dynamic_linear_b->options()); assert(qwen36_set_adapter_lora_tensor( diff --git a/crates/rustrain-server/src/session.rs b/crates/rustrain-server/src/session.rs index 89ad821f..878b2c13 100644 --- a/crates/rustrain-server/src/session.rs +++ b/crates/rustrain-server/src/session.rs @@ -314,8 +314,26 @@ impl TrainingSession for Qwen36Session { .ok() .and_then(|s| s.parse::().ok()) .unwrap_or(0); - let is_ep = ep_world_size > 1 && runtime_config.is_moe; - let is_data_parallel = ep_world_size > 1 && !runtime_config.is_moe; + let tp_size = std::env::var("TP_SIZE") + .or_else(|_| std::env::var("RUSTRAIN_TP_SIZE")) + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(1); + if tp_size > 1 { + if ep_world_size != tp_size || req.rank < 0 || req.rank as usize >= ep_world_size { + return Err(anyhow!( + "native Qwen server TP-only mode requires WORLD_SIZE=TP_SIZE and a valid global rank (world={}, tp={}, rank={})", + ep_world_size, + tp_size, + req.rank + )); + } + unsafe { + std::env::set_var("TP_SIZE", tp_size.to_string()); + } + } + let is_ep = ep_world_size > 1 && runtime_config.is_moe && tp_size == 1; + let is_data_parallel = ep_world_size > 1 && !runtime_config.is_moe && tp_size == 1; // Compute expert shard let (expert_start, expert_count) = if is_ep { @@ -333,7 +351,7 @@ impl TrainingSession for Qwen36Session { // Set CUDA device for any torchrun worker. Dense Qwen workers use // replicated weights and NCCL gradient all-reduce (LoRA-only DP). - if is_ep || is_data_parallel { + if is_ep || is_data_parallel || tp_size > 1 { self.device = tch::Device::Cuda(local_rank); } @@ -397,9 +415,9 @@ impl TrainingSession for Qwen36Session { expert_count, )?; - // Initialize NCCL directly in C++. The same communicator handles EP - // output collectives and replicated-weight LoRA gradient all-reduce. - let nccl_ep = if is_ep || is_data_parallel { + // Initialize NCCL directly in C++. The parent communicator handles + // EP/DP, while TP-only LoRA uses a split communicator for deltas. + let nccl_ep = if ep_world_size > 1 { unsafe { std::env::set_var( "RUSTRAIN_DATA_PARALLEL", @@ -414,7 +432,8 @@ impl TrainingSession for Qwen36Session { ep_rank, ep_world_size, data_parallel = is_data_parallel, - "NCCL communicator created in C++ for EP" + tp_size, + "NCCL communicator created in C++ for Qwen parallel training" ); true } else { diff --git a/docs/plans/qwen-lora-megatron-progress.md b/docs/plans/qwen-lora-megatron-progress.md index 94a92c1e..b69b3765 100644 --- a/docs/plans/qwen-lora-megatron-progress.md +++ b/docs/plans/qwen-lora-megatron-progress.md @@ -12,27 +12,30 @@ timestamp: 2026-07-17T00:00:00Z # Current State -Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP smoke, dense replicated-DP smoke, dynamic batch logical-step update, ABI9 fixed and per-tenant optimizer-step restore, selected-tenant isolation, and 5D topology mapping. +Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP smoke, dense replicated-DP smoke, TP-only latent-rank-sharded LoRA smoke, dynamic batch logical-step update, ABI10 fixed and per-tenant optimizer-step restore, selected-tenant isolation, and 5D topology mapping. -Not yet verified or implemented: Qwen native TP/PP/CP, FP32 accumulation/abort, rank-sharded checkpoint topology, DeepEP/TE prebuilt integration, and matched Megatron throughput. +Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axis TP+DP/EP, PP/CP, FP32 accumulation/abort, rank-sharded checkpoint topology, DeepEP/TE prebuilt integration, and matched Megatron throughput. # Durable Milestones - `4e242ee`: added Megatron-style 5D topology contract and launcher normalization; `cargo test -p rustrain-parallel --lib` passed 14/14. - `afcf091`: restored native C++ Adam step on checkpoint load and added ABI8 smoke assertions; server checkpoint tests passed 2/2. - ABI9 working tree: selected adapter IDs flow through HTTP, IPC, Rust, and C++; each dynamic tenant owns its Adam clock, checkpoint metadata preserves it, and failed selection restores the complete registry. +- ABI10 working tree: TP-only mode shards each fixed and dynamic adapter's latent rank, all-reduces only activation-level LoRA deltas on a split TP communicator, and keeps replicated MoE layers off the EP communicator. - H20 `123.57.26.97:28004`: ABI8 native smoke passed grouped/fallback parity, GDN, dense/MoE LoRA, dynamic adapters, and step setter validation. - H20 `123.57.26.97:28004`: ABI9 native smoke passed selected-tenant training with a positive selected update, exactly zero unselected update, independent clocks (`2` vs `1`), and registry preservation after an unknown ID. +- H20 `123.57.26.97:28004`: ABI10 two-rank TP native smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, and selected-tenant isolation. Rank-local LoRA tensors used distinct rank `4` slices for global rank `8`; losses matched and both shards had positive updates. - Target runtime probe: PyTorch 2.5.1+cu121, ABI0; Transformer Engine, flash-attn, DeepEP, Triton, and DeepSpeed are not importable. # Decisions During Execution - Keep TP and EP communicators separate; do not reuse the existing EP `LayerConfig.nccl_comm` for LoRA TP deltas. +- Scope multi-LoRA `n_max` rendezvous files per native context; reusing one filename across sessions can give ranks different chunk schedules and deadlock TP collectives. - Do not enable Qwen TP/PP/CP by merely relaxing runtime validation. - Treat Exa/Jina dependency search failures as missing evidence, not as proof that a package is compatible. # Verification -Passed: `cargo test -p rustrain-core --lib` (8), `cargo test -p rustrain-parallel --lib` (14), `cargo test -p rustrain-server --lib` (3), Qwen integration (6), remote ABI8 smoke, and remote ABI9 selected-tenant native smoke. +Passed: `cargo test -p rustrain-core --lib` (8), `cargo test -p rustrain-parallel --lib` (14), `cargo test -p rustrain-server --lib` (3), Qwen unit tests (3), Qwen integration (6), remote ABI8 smoke, remote ABI9 selected-tenant native smoke, remote ABI10 single-rank native smoke, and remote ABI10 two-rank TP native smoke. -Not run: full Qwen TP/PP/CP smoke, FP32 accumulation equivalence, rank-sharded checkpoint resume, and matched Megatron performance benchmark. +Not run: Megatron-style base-model TP, multi-axis TP+DP/EP, PP/CP, FP32 accumulation equivalence, rank-sharded checkpoint resume, and matched Megatron performance benchmark. From 5aa4ea4169e24a4dcfd0dbb124663bd077a2a5cb Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 06:07:18 +0800 Subject: [PATCH 012/156] feat: accumulate native lora gradients in fp32 --- .../rustrain-qwen3-6/kernels/fused_kernels.cu | 10 +- .../kernels/qwen3_6_kernels.cpp | 389 +++++++++++++++--- crates/rustrain-qwen3-6/src/kernel.rs | 31 +- .../rustrain-qwen3-6/tests/native_smoke.cpp | 43 +- 4 files changed, 409 insertions(+), 64 deletions(-) diff --git a/crates/rustrain-qwen3-6/kernels/fused_kernels.cu b/crates/rustrain-qwen3-6/kernels/fused_kernels.cu index c3894d05..56abc241 100644 --- a/crates/rustrain-qwen3-6/kernels/fused_kernels.cu +++ b/crates/rustrain-qwen3-6/kernels/fused_kernels.cu @@ -187,7 +187,7 @@ __global__ void fused_rmsnorm_matmul_kernel( // 4. Multi-tensor Fused Adam // ────────────────────────────────────────────────────────────────────── // One block per param tensor. Each thread processes multiple elements. -// Handles BF16 params/grads with FP32 m/v. +// Handles BF16 params with FP32 accumulated grads and FP32 m/v. // // Replaces 7 ATen ops per param with 1 kernel launch for ALL params: // m = m * beta1 + grad * (1 - beta1) @@ -197,7 +197,7 @@ __global__ void fused_rmsnorm_matmul_kernel( __global__ void fused_adam_multi_kernel( void** __restrict__ param_ptrs, // [n_params] BF16 - void** __restrict__ grad_ptrs, // [n_params] BF16 + void** __restrict__ grad_ptrs, // [n_params] FP32 float** __restrict__ m_ptrs, // [n_params] FP32 float** __restrict__ v_ptrs, // [n_params] FP32 const int* __restrict__ sizes, // [n_params] @@ -211,12 +211,12 @@ __global__ void fused_adam_multi_kernel( int size = sizes[pidx]; __nv_bfloat16* param = (__nv_bfloat16*)param_ptrs[pidx]; - __nv_bfloat16* grad = (__nv_bfloat16*)grad_ptrs[pidx]; + float* grad = (float*)grad_ptrs[pidx]; float* m = m_ptrs[pidx]; float* v = v_ptrs[pidx]; for (int i = threadIdx.x; i < size; i += blockDim.x) { - float g = __bfloat162float(grad[i]); + float g = grad[i]; float m_new = m[i] * beta1 + g * one_minus_beta1; float v_new = v[i] * beta2 + g * g * one_minus_beta2; m[i] = m_new; @@ -224,7 +224,7 @@ __global__ void fused_adam_multi_kernel( float p = __bfloat162float(param[i]); p -= lr_scaled * m_new / (sqrtf(v_new) + eps_scaled); param[i] = __float2bfloat16_rn(p); - grad[i] = __float2bfloat16_rn(0.0f); + grad[i] = 0.0f; } } diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index cf2c2774..ef8047c0 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -1503,6 +1503,10 @@ struct TrainingContext { std::set target_modules; std::map>> params; std::map>> adam_state; + // Gradients are harvested after each backward into FP32 tensors. The + // accumulator tensors are intentionally shared by value when chunk + // registry guards copy adapters, so their contents survive restore. + std::map>> grad_accum; }; std::vector adapters; @@ -1522,6 +1526,10 @@ struct TrainingContext { // Legacy single-LoRA (backward compat) std::vector lora_a; std::vector lora_b; + // Fixed-slot FP32 gradient accumulators. Inactive slots are undefined + // tensors to preserve the positional LoRA ABI without extra allocation. + std::vector grad_accum_a; + std::vector grad_accum_b; std::vector lora_active; std::vector lora_layer_offset; double lora_scaling; @@ -1532,6 +1540,13 @@ struct TrainingContext { std::vector adam_v; double lr, beta1, beta2, eps; int64_t step_count; + // A failed native call aborts the in-flight gradient window. This flag is + // consumed by the scoped accumulation guard on stack unwinding. + bool accumulation_active = false; + // Sum of (micro weight * supervised tokens) for the fixed-LoRA window. + // Fixed gradients are accumulated as token-weighted numerators and + // normalized exactly once at the optimizer boundary. + double accumulated_token_weight = 0.0; // Device buffer cache for multi-tensor fused Adam AdamDevBuffers adam_dev_bufs; @@ -1573,6 +1588,119 @@ struct TrainingContext { // ────────────────────────────────────────────────────────────────────── }; +static bool harvest_leaf_grad( + at::Tensor& param, at::Tensor& fp32_accumulator +) { + auto grad = param.grad(); + if (!grad.defined()) return false; + TORCH_CHECK(fp32_accumulator.defined(), + "active LoRA parameter is missing its FP32 gradient accumulator"); + TORCH_CHECK(fp32_accumulator.scalar_type() == at::kFloat, + "LoRA gradient accumulator must be FP32"); + TORCH_CHECK(fp32_accumulator.sizes() == param.sizes(), + "LoRA gradient accumulator shape mismatch"); + at::NoGradGuard guard; + fp32_accumulator.add_(grad.to(at::kFloat)); + // Never retain BF16 leaf gradients across micro-batches. The next + // backward starts with a fresh leaf grad and is harvested again. + param.mutable_grad() = at::Tensor(); + return true; +} + +static bool harvest_adapter_gradients(TrainingContext::LoRAAdapter& adapter) { + bool harvested = false; + for (auto& [layer_idx, pairs] : adapter.params) { + auto accum_it = adapter.grad_accum.find(layer_idx); + TORCH_CHECK(accum_it != adapter.grad_accum.end() && + accum_it->second.size() == pairs.size(), + "dynamic LoRA gradient accumulator layout mismatch"); + for (size_t i = 0; i < pairs.size(); ++i) { + auto& [a, b] = pairs[i]; + auto& accum = accum_it->second[i]; + if (a.requires_grad()) harvested |= harvest_leaf_grad(a, accum[0]); + if (b.requires_grad()) harvested |= harvest_leaf_grad(b, accum[1]); + } + } + return harvested; +} + +static bool harvest_gradient_accumulators(TrainingContext* ctx) { + bool harvested = false; + for (auto& adapter : ctx->adapters) { + harvested |= harvest_adapter_gradients(adapter); + } + TORCH_CHECK(ctx->grad_accum_a.size() == ctx->lora_a.size() && + ctx->grad_accum_b.size() == ctx->lora_b.size(), + "fixed LoRA gradient accumulator layout mismatch"); + for (size_t i = 0; i < ctx->lora_a.size(); ++i) { + if (!ctx->lora_active[i]) continue; + harvested |= harvest_leaf_grad(ctx->lora_a[i], ctx->grad_accum_a[i]); + harvested |= harvest_leaf_grad(ctx->lora_b[i], ctx->grad_accum_b[i]); + } + ctx->accumulation_active |= harvested; + return harvested; +} + +static void clear_adapter_gradient_accumulators( + TrainingContext::LoRAAdapter& adapter +) { + at::NoGradGuard guard; + for (auto& [layer_idx, pairs] : adapter.params) { + auto accum_it = adapter.grad_accum.find(layer_idx); + for (size_t i = 0; i < pairs.size(); ++i) { + auto& [a, b] = pairs[i]; + if (a.grad().defined()) a.mutable_grad() = at::Tensor(); + if (b.grad().defined()) b.mutable_grad() = at::Tensor(); + if (accum_it != adapter.grad_accum.end() && + i < accum_it->second.size()) { + if (accum_it->second[i][0].defined()) accum_it->second[i][0].zero_(); + if (accum_it->second[i][1].defined()) accum_it->second[i][1].zero_(); + } + } + } +} + +static void clear_gradient_accumulators(TrainingContext* ctx) { + if (!ctx) return; + at::NoGradGuard guard; + for (auto& adapter : ctx->adapters) { + clear_adapter_gradient_accumulators(adapter); + } + for (size_t i = 0; i < ctx->lora_a.size(); ++i) { + if (ctx->lora_a[i].grad().defined()) + ctx->lora_a[i].mutable_grad() = at::Tensor(); + if (ctx->lora_b[i].grad().defined()) + ctx->lora_b[i].mutable_grad() = at::Tensor(); + if (i < ctx->grad_accum_a.size() && ctx->grad_accum_a[i].defined()) + ctx->grad_accum_a[i].zero_(); + if (i < ctx->grad_accum_b.size() && ctx->grad_accum_b[i].defined()) + ctx->grad_accum_b[i].zero_(); + } + ctx->group_inputs.clear(); + ctx->group_outputs.clear(); + ctx->lora_cache_valid = false; + ctx->lora_batch_valid = false; + ctx->accumulation_active = false; + ctx->accumulated_token_weight = 0.0; +} + +struct GradientAccumulationFailureGuard { + TrainingContext* ctx; + bool disarmed = false; + ~GradientAccumulationFailureGuard() noexcept { + if (disarmed) return; + try { + clear_gradient_accumulators(ctx); + } catch (const std::exception& e) { + fprintf(stderr, + "[q36] secondary gradient cleanup failure: %s\n", e.what()); + } catch (...) { + fprintf(stderr, + "[q36] secondary gradient cleanup failure: unknown error\n"); + } + } +}; + static at::Tensor tp_allreduce_lora_delta( TrainingContext* ctx, const at::Tensor& local_delta ) { @@ -1649,15 +1777,23 @@ static ncclDataType_t nccl_dtype_for(const at::Tensor& tensor) { } } -static void allreduce_lora_grad( - TrainingContext* ctx, at::Tensor& param, double local_token_scale +static void reduce_lora_accumulator( + TrainingContext* ctx, at::Tensor& accumulator, double scale, + bool allreduce ) { - auto grad = param.grad(); - if (!ctx->nccl_comm || !grad.defined()) return; - auto contiguous = grad.contiguous(); - if (local_token_scale != 1.0) { - contiguous = contiguous * local_token_scale; + if (!accumulator.defined()) return; + TORCH_CHECK(accumulator.scalar_type() == at::kFloat, + "LoRA DP gradient accumulator must be FP32"); + auto contiguous = accumulator.contiguous(); + if (scale != 1.0) { + contiguous = contiguous * scale; + } + if (!allreduce) { + at::NoGradGuard guard; + accumulator.copy_(contiguous); + return; } + TORCH_CHECK(ctx->nccl_comm, "LoRA gradient all-reduce has no communicator"); auto reduced = at::empty_like(contiguous); int dev = contiguous.device().index(); cudaSetDevice(dev); @@ -1667,7 +1803,8 @@ static void allreduce_lora_grad( nccl_dtype_for(contiguous), ncclSum, ctx->nccl_comm, stream); TORCH_CHECK(err == ncclSuccess, "NCCL LoRA gradient all-reduce failed: ", ncclGetErrorString(err)); - param.mutable_grad() = reduced; + at::NoGradGuard guard; + accumulator.copy_(reduced); } // Every rank evaluates the complete loss. Average replicated LoRA gradients @@ -1676,32 +1813,73 @@ static void allreduce_lora_grad( // in forward and keep replicated gradients local; routed expert adapters remain // local because their parameter tensors are sharded. static void synchronize_lora_gradients( - TrainingContext* ctx, const at::Tensor& target_mask + TrainingContext* ctx, const at::Tensor& target_mask, + double accumulated_token_weight = 0.0 ) { - if (!ctx->nccl_comm || !ctx->data_parallel) return; - auto shifted_mask = target_mask.narrow(1, 1, target_mask.size(1) - 1) - .to(at::kFloat).sum().reshape({1}); - auto global_mask = at::empty_like(shifted_mask); - auto stream = c10::cuda::getCurrentCUDAStream( - shifted_mask.device().index()).stream(); - auto err = ncclAllReduce( - shifted_mask.data_ptr(), global_mask.data_ptr(), 1, - ncclFloat, ncclSum, ctx->nccl_comm, stream); - TORCH_CHECK(err == ncclSuccess, "NCCL token-count all-reduce failed: ", + const bool allreduce = ctx->nccl_comm && ctx->data_parallel; + double scale = 1.0; + if (accumulated_token_weight > 0.0) { + double global_weight = accumulated_token_weight; + if (allreduce) { + auto local = at::full({1}, accumulated_token_weight, + at::TensorOptions().dtype(at::kFloat).device(target_mask.device())); + auto global = at::empty_like(local); + auto stream = c10::cuda::getCurrentCUDAStream( + target_mask.device().index()).stream(); + auto err = ncclAllReduce( + local.data_ptr(), global.data_ptr(), 1, + ncclFloat, ncclSum, ctx->nccl_comm, stream); + TORCH_CHECK(err == ncclSuccess, + "NCCL accumulated token-count all-reduce failed: ", ncclGetErrorString(err)); - const double local_tokens = shifted_mask.item(); - const double global_tokens = global_mask.item(); - const double token_scale = local_tokens / std::max(global_tokens, 1.0); + global_weight = global.item(); + } + scale = 1.0 / std::max(global_weight, 1.0); + } else { + // Dynamic multi-LoRA currently contributes one independently-normalized + // row per tenant. Preserve that contract while weighting replicated DP + // ranks by the selected batch's token count. + if (!allreduce) return; + auto shifted_mask = target_mask.narrow(1, 1, target_mask.size(1) - 1) + .to(at::kFloat).sum().reshape({1}); + auto global_mask = at::empty_like(shifted_mask); + auto stream = c10::cuda::getCurrentCUDAStream( + shifted_mask.device().index()).stream(); + auto err = ncclAllReduce( + shifted_mask.data_ptr(), global_mask.data_ptr(), 1, + ncclFloat, ncclSum, ctx->nccl_comm, stream); + TORCH_CHECK(err == ncclSuccess, "NCCL token-count all-reduce failed: ", + ncclGetErrorString(err)); + const double local_tokens = shifted_mask.item(); + const double global_tokens = global_mask.item(); + scale = local_tokens / std::max(global_tokens, 1.0); + } for (auto& adapter : ctx->adapters) { for (auto& [layer_idx, pairs] : adapter.params) { auto table = lora_projection_table(ctx->layer_configs[layer_idx]); for (int64_t pair = 0; pair < (int64_t)pairs.size(); ++pair) { + auto accum_it = adapter.grad_accum.find(layer_idx); + TORCH_CHECK(accum_it != adapter.grad_accum.end() && + pair < (int64_t)accum_it->second.size(), + "dynamic LoRA gradient accumulator layout mismatch"); // Dynamic routed-expert tensors are sharded exactly like the - // base experts; only replicated adapter tensors are reduced. - if (table.entries[pair].grouped_expert) continue; - auto& [a, b] = pairs[pair]; - allreduce_lora_grad(ctx, a, token_scale); - allreduce_lora_grad(ctx, b, token_scale); + // base experts. In a fixed accumulation window they still + // need local token normalization, but must never be reduced + // across EP ranks; the legacy per-row path leaves them + // untouched to preserve its independent-sample contract. + if (table.entries[pair].grouped_expert) { + if (accumulated_token_weight > 0.0) { + reduce_lora_accumulator( + ctx, accum_it->second[pair][0], scale, false); + reduce_lora_accumulator( + ctx, accum_it->second[pair][1], scale, false); + } + continue; + } + reduce_lora_accumulator( + ctx, accum_it->second[pair][0], scale, allreduce); + reduce_lora_accumulator( + ctx, accum_it->second[pair][1], scale, allreduce); } } } @@ -1712,9 +1890,19 @@ static void synchronize_lora_gradients( // Routed expert LoRA is sharded with the base expert weights. Its // local gradients belong only to this EP rank and must not be // summed with a different expert shard on another rank. - if (table.entries[pair].grouped_expert) continue; - allreduce_lora_grad(ctx, ctx->lora_a[offset + pair], token_scale); - allreduce_lora_grad(ctx, ctx->lora_b[offset + pair], token_scale); + if (table.entries[pair].grouped_expert) { + if (accumulated_token_weight > 0.0) { + reduce_lora_accumulator( + ctx, ctx->grad_accum_a[offset + pair], scale, false); + reduce_lora_accumulator( + ctx, ctx->grad_accum_b[offset + pair], scale, false); + } + continue; + } + reduce_lora_accumulator( + ctx, ctx->grad_accum_a[offset + pair], scale, allreduce); + reduce_lora_accumulator( + ctx, ctx->grad_accum_b[offset + pair], scale, allreduce); } } } @@ -3560,7 +3748,7 @@ static at::Tensor mtp_compute_loss( extern "C" { __attribute__((visibility("default"))) int64_t qwen36_kernel_abi_version() { - return 10; + return 11; } // Create training context — called once at startup @@ -3712,6 +3900,11 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( } a.set_requires_grad(active); b.set_requires_grad(active); + auto grad_opts = at::TensorOptions().dtype(at::kFloat).device(base->device()); + ctx->grad_accum_a.push_back( + active ? at::zeros(a.sizes(), grad_opts) : at::Tensor()); + ctx->grad_accum_b.push_back( + active ? at::zeros(b.sizes(), grad_opts) : at::Tensor()); ctx->lora_a.push_back(std::move(a)); ctx->lora_b.push_back(std::move(b)); ctx->lora_active.push_back(active ? 1 : 0); @@ -3784,6 +3977,7 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( ) { try { auto* ctx = reinterpret_cast(ctx_ptr); + GradientAccumulationFailureGuard accumulation_guard{ctx}; TORCH_CHECK(gradient_scale > 0.0 && std::isfinite(gradient_scale), "gradient_scale must be finite and positive"); // Set CUDA device for EP @@ -3833,9 +4027,15 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( loss_val += mtp_loss.item(); } - if (gradient_scale != 1.0) { - total_hidden_grad.mul_(gradient_scale); - } + const double supervised_tokens = target_mask + .narrow(1, 1, target_mask.size(1) - 1) + .to(at::kFloat).sum().item(); + const double micro_token_weight = + gradient_scale * std::max(supervised_tokens, 1.0); + // compute_loss returns a local token mean. Convert it to a weighted + // numerator before backward; the FP32 window is divided by the global + // accumulated token weight exactly once at the optimizer boundary. + total_hidden_grad.mul_(micro_token_weight); // Trigger exactly one main-model backward with the combined hidden // gradient. Manual groups are the non-autograd checkpoint fallback; @@ -3846,26 +4046,33 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( hidden.backward(total_hidden_grad); } + // Consume every BF16 leaf gradient immediately. Only the FP32 buffers + // survive the micro-step boundary. + harvest_gradient_accumulators(ctx); + ctx->accumulated_token_weight += micro_token_weight; + if (!apply_optimizer) { // The forward graph has been consumed, but parameters remain // live for the next micro-batch. Never reuse cached LoRA deltas // whose autograd nodes were freed by this backward. ctx->lora_cache_valid = false; ctx->lora_batch_valid = false; + accumulation_guard.disarmed = true; return loss_val; } // Replicated DP gradients are synchronized before the local Adam // update. EP keeps replicated gradients local because its forward // routed activation already contains the cross-rank sum. - synchronize_lora_gradients(ctx, target_mask); + synchronize_lora_gradients( + ctx, target_mask, ctx->accumulated_token_weight); // ── Adam optimizer step — CUDA multi-tensor fused kernel ── at::AutoGradMode guard(false); - ctx->step_count++; ctx->lora_cache_valid = false; ctx->lora_batch_valid = false; - double step_f = (double)ctx->step_count; + const int64_t next_step = ctx->step_count + 1; + double step_f = (double)next_step; double bias_correction1 = 1.0 - std::pow(ctx->beta1, step_f); double bias_correction2 = 1.0 - std::pow(ctx->beta2, step_f); float lr_scaled = (float)(ctx->lr / bias_correction1); @@ -3882,19 +4089,23 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( for (auto& adapter : ctx->adapters) { for (auto& [layer_idx, pairs] : adapter.params) { auto& adam_states = adapter.adam_state[layer_idx]; + auto& accumulators = adapter.grad_accum[layer_idx]; for (size_t i = 0; i < pairs.size(); i++) { auto& [a, b] = pairs[i]; auto& [m_a, v_a, m_b, v_b] = adam_states[i]; - if (a.grad().defined() && a.scalar_type() == at::kBFloat16) { + auto& [accum_a, accum_b] = accumulators[i]; + if (a.requires_grad() && accum_a.defined() && + a.scalar_type() == at::kBFloat16) { h_params.push_back(a.data_ptr()); - h_grads.push_back(a.grad().data_ptr()); + h_grads.push_back(accum_a.data_ptr()); h_m.push_back((float*)m_a.data_ptr()); h_v.push_back((float*)v_a.data_ptr()); h_sizes.push_back((int)a.numel()); } - if (b.grad().defined() && b.scalar_type() == at::kBFloat16) { + if (b.requires_grad() && accum_b.defined() && + b.scalar_type() == at::kBFloat16) { h_params.push_back(b.data_ptr()); - h_grads.push_back(b.grad().data_ptr()); + h_grads.push_back(accum_b.data_ptr()); h_m.push_back((float*)m_b.data_ptr()); h_v.push_back((float*)v_b.data_ptr()); h_sizes.push_back((int)b.numel()); @@ -3907,10 +4118,11 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( for (size_t i = 0; i < ctx->lora_a.size(); i++) { { auto& param = ctx->lora_a[i]; - auto& grad = param.grad(); - if (grad.defined() && param.scalar_type() == at::kBFloat16) { + auto& accum = ctx->grad_accum_a[i]; + if (ctx->lora_active[i] && accum.defined() && + param.scalar_type() == at::kBFloat16) { h_params.push_back(param.data_ptr()); - h_grads.push_back(grad.data_ptr()); + h_grads.push_back(accum.data_ptr()); h_m.push_back((float*)ctx->adam_m[adam_idx].data_ptr()); h_v.push_back((float*)ctx->adam_v[adam_idx].data_ptr()); h_sizes.push_back((int)param.numel()); @@ -3919,10 +4131,11 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( adam_idx++; { auto& param = ctx->lora_b[i]; - auto& grad = param.grad(); - if (grad.defined() && param.scalar_type() == at::kBFloat16) { + auto& accum = ctx->grad_accum_b[i]; + if (ctx->lora_active[i] && accum.defined() && + param.scalar_type() == at::kBFloat16) { h_params.push_back(param.data_ptr()); - h_grads.push_back(grad.data_ptr()); + h_grads.push_back(accum.data_ptr()); h_m.push_back((float*)ctx->adam_m[adam_idx].data_ptr()); h_v.push_back((float*)ctx->adam_v[adam_idx].data_ptr()); h_sizes.push_back((int)param.numel()); @@ -3967,8 +4180,15 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( one_minus_b1, one_minus_b2, (void*)stream ); + auto launch_error = cudaGetLastError(); + TORCH_CHECK(launch_error == cudaSuccess, + "fused FP32-gradient Adam launch failed: ", + cudaGetErrorString(launch_error)); + ctx->step_count = next_step; } + clear_gradient_accumulators(ctx); + accumulation_guard.disarmed = true; return loss_val; } catch (const std::exception& e) { fprintf(stderr, "[q36] train_step FAILED: %s\n", e.what()); @@ -4001,6 +4221,35 @@ __attribute__((visibility("default"))) void* qwen36_get_lora_b(void* ctx_ptr, in return &ctx->lora_b[index]; } +// Read-only diagnostic accessor used by native validation and runtime +// observability. The returned tensor is owned by the training context. +__attribute__((visibility("default"))) void* qwen36_get_lora_grad_accumulator( + void* ctx_ptr, int64_t index, int32_t is_b +) { + auto* ctx = reinterpret_cast(ctx_ptr); + if (!ctx || index < 0 || index >= (int64_t)ctx->lora_a.size()) return nullptr; + auto& accumulators = is_b ? ctx->grad_accum_b : ctx->grad_accum_a; + if (index >= (int64_t)accumulators.size() || !accumulators[index].defined()) + return nullptr; + return &accumulators[index]; +} + +// Explicitly abort an incomplete accumulation window. The operation is +// idempotent and leaves parameters, Adam state, and optimizer clocks intact. +__attribute__((visibility("default"))) int32_t qwen36_abort_gradient_accumulation( + void* ctx_ptr +) { + try { + auto* ctx = reinterpret_cast(ctx_ptr); + TORCH_CHECK(ctx, "null training context"); + clear_gradient_accumulators(ctx); + return 0; + } catch (const std::exception& e) { + fprintf(stderr, "[q36] abort_gradient_accumulation FAILED: %s\n", e.what()); + return -1; + } +} + // Copy one exported LoRA tensor back into the native leaf parameter. This is // used by checkpoint resume and adapter import; derived delta caches are // invalidated so the next forward rebuilds the graph from the new leaf. @@ -4107,6 +4356,7 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( ) { try { auto* ctx = reinterpret_cast(ctx_ptr); + GradientAccumulationFailureGuard accumulation_guard{ctx}; if (ctx->nccl_comm) { c10::cuda::set_device(ctx->cuda_device); cudaSetDevice(ctx->cuda_device); @@ -4221,10 +4471,10 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( double total_loss = 0.0; int64_t num_chunks = (total_adapters + n_max - 1) / n_max; - // Chunking is a memory scheduling detail, not an optimizer step. The - // session clock is kept for backwards-compatible status reporting; - // Adam bias correction below uses each adapter's own clock. - ctx->step_count++; + // Chunking is a memory scheduling detail, not an optimizer step. Both + // the session clock and tenant clocks commit only after Adam launches. + const int64_t next_session_step = ctx->step_count + 1; + bool any_update = false; for (int64_t chunk = 0; chunk < num_chunks; chunk++) { int64_t start = chunk * n_max; @@ -4314,6 +4564,10 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( } else { hidden.backward(hidden_grad); } + // The chunk registry contains intrusive copies of the selected + // adapter tensors. Harvest now; FP32 accumulator contents remain + // shared when the full registry is restored below. + harvest_gradient_accumulators(ctx); // No cudaDeviceSynchronize — let GPU pipeline run asynchronously. // The next chunk's CPU prep (LoRA batch, input expand) will overlap // with the tail of this chunk's GPU backward. @@ -4340,10 +4594,9 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( at::AutoGradMode guard(false); ctx->lora_cache_valid = false; ctx->lora_batch_valid = false; - for (auto& adapter : ctx->adapters) adapter.optimizer_step++; std::map> groups; for (auto& adapter : ctx->adapters) { - groups[adapter.optimizer_step].push_back(&adapter); + groups[adapter.optimizer_step + 1].push_back(&adapter); } for (auto& [logical_step, adapters] : groups) { std::vector h_params, h_grads; @@ -4352,19 +4605,23 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( for (auto* adapter : adapters) { for (auto& [layer_idx, pairs] : adapter->params) { auto& adam_states = adapter->adam_state[layer_idx]; + auto& accumulators = adapter->grad_accum[layer_idx]; for (size_t i = 0; i < pairs.size(); i++) { auto& [a, b] = pairs[i]; auto& [m_a, v_a, m_b, v_b] = adam_states[i]; - if (a.grad().defined() && a.scalar_type() == at::kBFloat16) { + auto& [accum_a, accum_b] = accumulators[i]; + if (a.requires_grad() && accum_a.defined() && + a.scalar_type() == at::kBFloat16) { h_params.push_back(a.data_ptr()); - h_grads.push_back(a.grad().data_ptr()); + h_grads.push_back(accum_a.data_ptr()); h_m.push_back((float*)m_a.data_ptr()); h_v.push_back((float*)v_a.data_ptr()); h_sizes.push_back((int)a.numel()); } - if (b.grad().defined() && b.scalar_type() == at::kBFloat16) { + if (b.requires_grad() && accum_b.defined() && + b.scalar_type() == at::kBFloat16) { h_params.push_back(b.data_ptr()); - h_grads.push_back(b.grad().data_ptr()); + h_grads.push_back(accum_b.data_ptr()); h_m.push_back((float*)m_b.data_ptr()); h_v.push_back((float*)v_b.data_ptr()); h_sizes.push_back((int)b.numel()); @@ -4402,6 +4659,14 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( n_params, (float)ctx->beta1, (float)ctx->beta2, lr_scaled, eps_scaled, one_minus_b1, one_minus_b2, (void*)stream); + auto launch_error = cudaGetLastError(); + TORCH_CHECK(launch_error == cudaSuccess, + "dynamic fused FP32-gradient Adam launch failed: ", + cudaGetErrorString(launch_error)); + for (auto* adapter : adapters) { + adapter->optimizer_step = logical_step; + } + any_update = true; } } @@ -4411,6 +4676,9 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( (long)(chunk + 1), (long)num_chunks, (long)n, loss_val); } + if (any_update) ctx->step_count = next_session_step; + clear_gradient_accumulators(ctx); + accumulation_guard.disarmed = true; return total_loss / total_adapters; } catch (const std::exception& e) { fprintf(stderr, "[train_multi] FAILED: %s\n", e.what()); @@ -4863,6 +5131,7 @@ int64_t qwen36_add_lora( int64_t num_pairs = projection_table.count; std::vector> pairs; std::vector> adam_states; + std::vector> grad_accumulators; for (int64_t k = 0; k < num_pairs; k++) { const auto& projection = projection_table.entries[k]; auto* base = ctx->weight_ptrs[w_offset + projection.weight_index]; @@ -4904,10 +5173,16 @@ int64_t qwen36_add_lora( at::zeros(a.sizes(), opts_f32), at::zeros(a.sizes(), opts_f32), at::zeros(b.sizes(), opts_f32), at::zeros(b.sizes(), opts_f32) }); + grad_accumulators.push_back(active + ? std::array{ + at::zeros(a.sizes(), opts_f32), + at::zeros(b.sizes(), opts_f32)} + : std::array{at::Tensor(), at::Tensor()}); pairs.emplace_back(std::move(a), std::move(b)); } adapter.params[i] = std::move(pairs); adapter.adam_state[i] = std::move(adam_states); + adapter.grad_accum[i] = std::move(grad_accumulators); } int64_t id = adapter.id; ctx->adapters.push_back(std::move(adapter)); diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs index 42f33e1a..c9c48a4f 100644 --- a/crates/rustrain-qwen3-6/src/kernel.rs +++ b/crates/rustrain-qwen3-6/src/kernel.rs @@ -52,6 +52,8 @@ type FnGetLoraCount = unsafe extern "C" fn(*mut c_void) -> i64; type FnGetLoraA = unsafe extern "C" fn(*mut c_void, i64) -> *mut c_void; type FnGetLoraB = unsafe extern "C" fn(*mut c_void, i64) -> *mut c_void; type FnSetLoraTensor = unsafe extern "C" fn(*mut c_void, i64, i32, *mut c_void) -> i32; +type FnGetLoraGradAccumulator = unsafe extern "C" fn(*mut c_void, i64, i32) -> *mut c_void; +type FnAbortGradientAccumulation = unsafe extern "C" fn(*mut c_void) -> i32; type FnGetStepCount = unsafe extern "C" fn(*mut c_void) -> i64; type FnSetStepCount = unsafe extern "C" fn(*mut c_void, i64) -> i32; type FnExportOptimizer = @@ -128,6 +130,8 @@ struct KernelHandles { get_lora_a: FnGetLoraA, get_lora_b: FnGetLoraB, set_lora_tensor: FnSetLoraTensor, + get_lora_grad_accumulator: FnGetLoraGradAccumulator, + abort_gradient_accumulation: FnAbortGradientAccumulation, get_step_count: FnGetStepCount, set_step_count: FnSetStepCount, export_optimizer: FnExportOptimizer, @@ -190,7 +194,7 @@ unsafe fn load_kernels() -> Option { }}; } let abi_version: FnKernelAbiVersion = sym!("qwen36_kernel_abi_version"); - if abi_version() != 10 { + if abi_version() != 11 { return None; } Some(KernelHandles { @@ -204,6 +208,8 @@ unsafe fn load_kernels() -> Option { get_lora_a: sym!("qwen36_get_lora_a"), get_lora_b: sym!("qwen36_get_lora_b"), set_lora_tensor: sym!("qwen36_set_lora_tensor"), + get_lora_grad_accumulator: sym!("qwen36_get_lora_grad_accumulator"), + abort_gradient_accumulation: sym!("qwen36_abort_gradient_accumulation"), get_step_count: sym!("qwen36_get_step_count"), set_step_count: sym!("qwen36_set_step_count"), export_optimizer: sym!("qwen36_export_optimizer_state"), @@ -676,6 +682,29 @@ impl CppTrainingContext { Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) }) } + /// Inspect the native FP32 gradient accumulator for one fixed LoRA slot. + /// The returned tensor is a shallow handle owned by the C++ context. + pub fn get_lora_gradient_accumulator(&self, index: i64, is_b: bool) -> Option { + let kh = get_kernels()?; + let ptr = + unsafe { (kh.get_lora_grad_accumulator)(self.ptr, index, if is_b { 1 } else { 0 }) }; + if ptr.is_null() { + return None; + } + Some(unsafe { Tensor::clone_from_ptr(ptr as *mut _) }) + } + + /// Abort an incomplete micro-batch window without changing parameters, + /// Adam state, or optimizer clocks. Safe to call when no window is active. + pub fn abort_gradient_accumulation(&self) -> Result<()> { + let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; + let status = unsafe { (kh.abort_gradient_accumulation)(self.ptr) }; + if status != 0 { + bail!("C++ abort_gradient_accumulation failed"); + } + Ok(()) + } + pub fn set_lora_tensor(&self, index: i64, is_b: bool, tensor: &Tensor) -> Result<()> { let kh = get_kernels().ok_or_else(|| anyhow::anyhow!("kernels not loaded"))?; let status = unsafe { diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index 141ea426..624221ac 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -28,6 +28,8 @@ extern "C" int32_t qwen36_init_nccl(void*); extern "C" int64_t qwen36_get_lora_count(void*); extern "C" void* qwen36_get_lora_a(void*, int64_t); extern "C" void* qwen36_get_lora_b(void*, int64_t); +extern "C" void* qwen36_get_lora_grad_accumulator(void*, int64_t, int32_t); +extern "C" int32_t qwen36_abort_gradient_accumulation(void*); extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*); extern "C" double qwen36_train_step(void*, void*, void*, void*); extern "C" double qwen36_train_micro_step( @@ -54,7 +56,7 @@ static at::Tensor cuda_rand(std::initializer_list shape) { } int main() { - assert(qwen36_kernel_abi_version() == 10); + assert(qwen36_kernel_abi_version() == 11); const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); const int process_rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0"); const int local_rank = std::atoi(std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); @@ -415,6 +417,12 @@ int main() { auto linear_b_value = at::ones(linear_b->sizes(), linear_b->options()); assert(qwen36_set_lora_tensor(ctx, 0, 1, &linear_b_value) == 0); auto linear_a_before = linear_a->clone(); + auto* linear_a_accum = reinterpret_cast( + qwen36_get_lora_grad_accumulator(ctx, 0, 0)); + assert(linear_a_accum); + assert(linear_a_accum->scalar_type() == at::kFloat); + assert(linear_a_accum->sizes() == linear_a->sizes()); + assert(linear_a_accum->abs().sum().item() == 0.0); assert(qwen36_get_step_count(ctx) == 0); assert(qwen36_set_step_count(ctx, -1) != 0); assert(qwen36_get_step_count(ctx) == 0); @@ -427,12 +435,45 @@ int main() { assert(accum_loss_0 == accum_loss_0); assert(qwen36_get_step_count(ctx) == 0); assert((*linear_a - linear_a_before).abs().sum().item() == 0.0); + assert(linear_a_accum->abs().sum().item() > 0.0); + // The BF16 leaf gradient is consumed at the micro-step boundary; only the + // FP32 accumulator may survive. + assert(!linear_a->grad().defined()); + + // A failing final micro-step aborts the existing window transactionally. + const double failed_accum = qwen36_train_micro_step( + ctx, &input_ids, &target_mask, &attention_mask, NAN, 1); + c10::cuda::device_synchronize(); + assert(failed_accum < 0.0); + assert(qwen36_get_step_count(ctx) == 0); + assert(linear_a_accum->abs().sum().item() == 0.0); + assert((*linear_a - linear_a_before).abs().sum().item() == 0.0); + + // Explicit abort is idempotent and clears a successful non-final micro. + assert(qwen36_train_micro_step( + ctx, &input_ids, &target_mask, &attention_mask, 0.5, 0) > 0.0); + c10::cuda::device_synchronize(); + assert(linear_a_accum->abs().sum().item() > 0.0); + assert(qwen36_abort_gradient_accumulation(ctx) == 0); + c10::cuda::device_synchronize(); + assert(linear_a_accum->abs().sum().item() == 0.0); + assert(qwen36_get_step_count(ctx) == 0); + assert((*linear_a - linear_a_before).abs().sum().item() == 0.0); + + // Two clean micro-batches accumulate into FP32 and commit one Adam step. + const double clean_accum_loss_0 = qwen36_train_micro_step( + ctx, &input_ids, &target_mask, &attention_mask, 0.5, 0); + c10::cuda::device_synchronize(); + assert(clean_accum_loss_0 == clean_accum_loss_0); + assert(linear_a_accum->scalar_type() == at::kFloat); + assert(linear_a_accum->abs().sum().item() > 0.0); const double accum_loss_1 = qwen36_train_micro_step( ctx, &input_ids, &target_mask, &attention_mask, 0.5, 1); c10::cuda::device_synchronize(); assert(accum_loss_1 == accum_loss_1); assert(qwen36_get_step_count(ctx) == 1); assert((*linear_a - linear_a_before).abs().sum().item() > 0.0); + assert(linear_a_accum->abs().sum().item() == 0.0); linear_a_before = linear_a->clone(); const double linear_loss = qwen36_train_step( ctx, &input_ids, &target_mask, &attention_mask); From ed25daa0cd673ba8c6abc1b30de5c4dcbb785711 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 06:07:28 +0800 Subject: [PATCH 013/156] feat: add rank-aware tp lora checkpoints --- crates/rustrain-server/src/checkpoint.rs | 747 ++++++++++++++++++++++- crates/rustrain-server/src/session.rs | 7 +- 2 files changed, 739 insertions(+), 15 deletions(-) diff --git a/crates/rustrain-server/src/checkpoint.rs b/crates/rustrain-server/src/checkpoint.rs index 3e265585..35403ce0 100644 --- a/crates/rustrain-server/src/checkpoint.rs +++ b/crates/rustrain-server/src/checkpoint.rs @@ -1,10 +1,119 @@ //! Checkpoint save/load: adapter (LoRA A/B) + optimizer state (Adam m/v) + step count. -use anyhow::{Context, Result}; +use anyhow::{Context, Result, bail}; use serde::{Deserialize, Serialize}; -use std::path::Path; +use std::collections::BTreeMap; +use std::env; +use std::path::{Path, PathBuf}; use tch::Tensor; +const TP_CHECKPOINT_FORMAT: &str = "rustrain-checkpoint-v3-tp"; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ParallelCheckpointManifest { + pub world_size: usize, + pub tensor_model_parallel_size: usize, + pub pipeline_model_parallel_size: usize, + pub data_parallel_size: usize, + pub expert_model_parallel_size: usize, + pub context_parallel_size: usize, + pub global_rank: usize, + pub tensor_model_parallel_rank: usize, +} + +impl ParallelCheckpointManifest { + pub fn new( + world_size: usize, + global_rank: usize, + tensor_model_parallel_size: usize, + pipeline_model_parallel_size: usize, + data_parallel_size: usize, + expert_model_parallel_size: usize, + context_parallel_size: usize, + ) -> Result { + let sizes = [ + tensor_model_parallel_size, + pipeline_model_parallel_size, + data_parallel_size, + expert_model_parallel_size, + context_parallel_size, + ]; + if sizes.contains(&0) { + bail!("parallel checkpoint sizes must be positive"); + } + let expected_world_size = sizes + .into_iter() + .try_fold(1usize, |product, size| product.checked_mul(size)) + .context("parallel checkpoint world size overflow")?; + if world_size != expected_world_size { + bail!( + "parallel checkpoint topology product is {expected_world_size}, but WORLD_SIZE is {world_size}" + ); + } + if global_rank >= world_size { + bail!("global rank {global_rank} is outside WORLD_SIZE={world_size}"); + } + Ok(Self { + world_size, + tensor_model_parallel_size, + pipeline_model_parallel_size, + data_parallel_size, + expert_model_parallel_size, + context_parallel_size, + global_rank, + tensor_model_parallel_rank: global_rank % tensor_model_parallel_size, + }) + } + + pub fn from_env() -> Result { + let world_size = env_usize(&["WORLD_SIZE"], 1)?; + let global_rank = env_usize(&["RANK"], 0)?; + let tp = env_usize(&["TP_SIZE", "RUSTRAIN_TP_SIZE"], 1)?; + let pp = env_usize(&["PP_SIZE", "RUSTRAIN_PP_SIZE"], 1)?; + let ep = env_usize(&["EP_SIZE", "RUSTRAIN_EP_SIZE"], 1)?; + let cp = env_usize(&["CP_SIZE", "RUSTRAIN_CP_SIZE"], 1)?; + let non_dp = tp + .checked_mul(pp) + .and_then(|size| size.checked_mul(ep)) + .and_then(|size| size.checked_mul(cp)) + .context("model-parallel topology product overflow")?; + let dp = match env_usize_optional(&["DP_SIZE", "RUSTRAIN_DP_SIZE"])? { + Some(dp) => dp, + None if world_size % non_dp == 0 => world_size / non_dp, + None => { + bail!("WORLD_SIZE={world_size} is not divisible by model-parallel product {non_dp}") + } + }; + Self::new(world_size, global_rank, tp, pp, dp, ep, cp) + } + + fn is_tensor_parallel(&self) -> bool { + self.tensor_model_parallel_size > 1 + } + + fn replica_identity(&self) -> String { + format!( + "global-rank-{}-tp-rank-{}", + self.global_rank, self.tensor_model_parallel_rank + ) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct TensorShardManifest { + pub file: String, + pub tensor_name: String, + pub state: String, + #[serde(default)] + pub adapter_id: Option, + pub global_lora_rank: i64, + pub global_shape: Vec, + pub local_shape: Vec, + pub partition_axis: usize, + pub global_offset: Vec, + pub replica_identity: String, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CheckpointManifest { pub format: String, @@ -16,6 +125,10 @@ pub struct CheckpointManifest { pub files: Vec, #[serde(default)] pub dynamic_adapters: Vec, + #[serde(default)] + pub parallel: Option, + #[serde(default)] + pub tensor_shards: Vec, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -93,6 +206,104 @@ pub fn save_checkpoint_with_dynamic( adam_v: &[Tensor], dynamic_adapters: &[DynamicAdapterCheckpoint], ) -> Result<()> { + save_checkpoint_with_dynamic_at( + dir, + step, + loss, + model_path, + lora_rank, + lora_alpha, + lora_a, + lora_b, + adam_m, + adam_v, + dynamic_adapters, + None, + ) +} + +#[allow(clippy::too_many_arguments)] +pub fn save_checkpoint_with_dynamic_for_topology( + dir: &Path, + step: u64, + loss: f64, + model_path: &str, + lora_rank: i64, + lora_alpha: f64, + lora_a: &[Tensor], + lora_b: &[Tensor], + adam_m: &[Tensor], + adam_v: &[Tensor], + dynamic_adapters: &[DynamicAdapterCheckpoint], + parallel: &ParallelCheckpointManifest, +) -> Result<()> { + if !parallel.is_tensor_parallel() { + return save_checkpoint_with_dynamic( + dir, + step, + loss, + model_path, + lora_rank, + lora_alpha, + lora_a, + lora_b, + adam_m, + adam_v, + dynamic_adapters, + ); + } + let rank_dir = rank_checkpoint_dir(dir, parallel.global_rank); + save_checkpoint_with_dynamic_at( + &rank_dir, + step, + loss, + model_path, + lora_rank, + lora_alpha, + lora_a, + lora_b, + adam_m, + adam_v, + dynamic_adapters, + Some(parallel), + ) +} + +#[allow(clippy::too_many_arguments)] +fn save_checkpoint_with_dynamic_at( + dir: &Path, + step: u64, + loss: f64, + model_path: &str, + lora_rank: i64, + lora_alpha: f64, + lora_a: &[Tensor], + lora_b: &[Tensor], + adam_m: &[Tensor], + adam_v: &[Tensor], + dynamic_adapters: &[DynamicAdapterCheckpoint], + parallel: Option<&ParallelCheckpointManifest>, +) -> Result<()> { + validate_tensor_counts( + lora_a, + lora_b, + adam_m, + adam_v, + dynamic_adapters, + parallel.is_some(), + )?; + let tensor_shards = match parallel { + Some(parallel) => build_tensor_shard_manifest( + parallel, + lora_rank, + lora_a, + lora_b, + adam_m, + adam_v, + dynamic_adapters, + )?, + None => Vec::new(), + }; std::fs::create_dir_all(dir) .with_context(|| format!("create checkpoint dir {}", dir.display()))?; @@ -101,16 +312,6 @@ pub fn save_checkpoint_with_dynamic( let mut adapter_tensors = named_tensors(lora_a, lora_b, "a_", "b_"); let mut dynamic_manifests = Vec::with_capacity(dynamic_adapters.len()); for adapter in dynamic_adapters { - if adapter.lora_a.len() != adapter.lora_b.len() - || adapter.adam_m.len() != adapter.adam_v.len() - || adapter.manifest.parameter_count != adapter.lora_a.len() - || adapter.manifest.optimizer_count != adapter.adam_m.len() - { - anyhow::bail!( - "dynamic adapter {} checkpoint count mismatch", - adapter.manifest.id - ); - } let id = adapter.manifest.id; adapter_tensors.extend(named_tensors( &adapter.lora_a, @@ -138,7 +339,9 @@ pub fn save_checkpoint_with_dynamic( // Write manifest let manifest = CheckpointManifest { - format: if dynamic_manifests.is_empty() { + format: if parallel.is_some() { + TP_CHECKPOINT_FORMAT.to_string() + } else if dynamic_manifests.is_empty() { "rustrain-checkpoint-v1".to_string() } else { "rustrain-checkpoint-v2".to_string() @@ -150,6 +353,8 @@ pub fn save_checkpoint_with_dynamic( lora_alpha, files: vec!["adapter.safetensors".into(), "optimizer.safetensors".into()], dynamic_adapters: dynamic_manifests, + parallel: parallel.cloned(), + tensor_shards, }; let manifest_path = dir.join("manifest.json"); std::fs::write(&manifest_path, serde_json::to_string_pretty(&manifest)?) @@ -166,12 +371,53 @@ pub fn save_checkpoint_with_dynamic( /// Load checkpoint from a directory. pub fn load_checkpoint(dir: &Path) -> Result { + load_checkpoint_at(dir, None) +} + +pub fn load_checkpoint_for_topology( + dir: &Path, + parallel: &ParallelCheckpointManifest, +) -> Result { + if !parallel.is_tensor_parallel() { + return load_checkpoint(dir); + } + let rank_dir = rank_checkpoint_dir(dir, parallel.global_rank); + load_checkpoint_at(&rank_dir, Some(parallel)) +} + +fn load_checkpoint_at( + dir: &Path, + expected_parallel: Option<&ParallelCheckpointManifest>, +) -> Result { let manifest_path = dir.join("manifest.json"); let manifest: CheckpointManifest = serde_json::from_str( &std::fs::read_to_string(&manifest_path) .with_context(|| format!("read {}", manifest_path.display()))?, ) .with_context(|| "parse manifest.json")?; + match expected_parallel { + Some(expected) => { + if manifest.format != TP_CHECKPOINT_FORMAT { + bail!( + "tensor-parallel resume requires {TP_CHECKPOINT_FORMAT}, found {}", + manifest.format + ); + } + let saved = manifest + .parallel + .as_ref() + .context("tensor-parallel checkpoint is missing topology metadata")?; + if saved != expected { + bail!( + "tensor-parallel checkpoint topology mismatch: saved={saved:?}, current={expected:?}" + ); + } + } + None if manifest.format == TP_CHECKPOINT_FORMAT => { + bail!("tensor-parallel checkpoint must be loaded with rank topology"); + } + None => {} + } let adapter_path = dir.join("adapter.safetensors"); let adapter_named = read_named_tensors(&adapter_path)?; @@ -205,6 +451,19 @@ pub fn load_checkpoint(dir: &Path) -> Result { }); } + if let Some(parallel) = expected_parallel { + let expected_shards = build_tensor_shard_manifest( + parallel, + manifest.lora_rank, + &lora_a, + &lora_b, + &adam_m, + &adam_v, + &dynamic_adapters, + )?; + validate_saved_shards(&manifest.tensor_shards, &expected_shards)?; + } + tracing::info!( step = manifest.step, loss = manifest.loss, @@ -221,6 +480,286 @@ pub fn load_checkpoint(dir: &Path) -> Result { }) } +#[derive(Clone, Copy)] +enum LoraSide { + A, + B, +} + +fn validate_tensor_counts( + lora_a: &[Tensor], + lora_b: &[Tensor], + adam_m: &[Tensor], + adam_v: &[Tensor], + dynamic_adapters: &[DynamicAdapterCheckpoint], + strict_optimizer_layout: bool, +) -> Result<()> { + if lora_a.len() != lora_b.len() || adam_m.len() != adam_v.len() { + bail!("fixed adapter checkpoint tensor count mismatch"); + } + if strict_optimizer_layout + && !adam_m.is_empty() + && adam_m.len() != lora_a.len().saturating_mul(2) + { + bail!("tensor-parallel fixed optimizer state must contain A/B entries for every LoRA slot"); + } + for adapter in dynamic_adapters { + if adapter.lora_a.len() != adapter.lora_b.len() + || adapter.adam_m.len() != adapter.adam_v.len() + || adapter.manifest.parameter_count != adapter.lora_a.len() + || adapter.manifest.optimizer_count != adapter.adam_m.len() + { + bail!( + "dynamic adapter {} checkpoint count mismatch", + adapter.manifest.id + ); + } + if strict_optimizer_layout + && !adapter.adam_m.is_empty() + && adapter.adam_m.len() != adapter.lora_a.len().saturating_mul(2) + { + bail!( + "tensor-parallel dynamic adapter {} optimizer state must contain A/B entries for every LoRA slot", + adapter.manifest.id + ); + } + } + Ok(()) +} + +fn build_tensor_shard_manifest( + parallel: &ParallelCheckpointManifest, + fixed_lora_rank: i64, + lora_a: &[Tensor], + lora_b: &[Tensor], + adam_m: &[Tensor], + adam_v: &[Tensor], + dynamic_adapters: &[DynamicAdapterCheckpoint], +) -> Result> { + validate_tensor_counts(lora_a, lora_b, adam_m, adam_v, dynamic_adapters, true)?; + let mut shards = Vec::new(); + append_adapter_shards( + &mut shards, + parallel, + None, + fixed_lora_rank, + lora_a, + lora_b, + adam_m, + adam_v, + )?; + for adapter in dynamic_adapters { + append_adapter_shards( + &mut shards, + parallel, + Some(adapter.manifest.id), + adapter.manifest.rank, + &adapter.lora_a, + &adapter.lora_b, + &adapter.adam_m, + &adapter.adam_v, + )?; + } + Ok(shards) +} + +#[allow(clippy::too_many_arguments)] +fn append_adapter_shards( + shards: &mut Vec, + parallel: &ParallelCheckpointManifest, + adapter_id: Option, + global_lora_rank: i64, + lora_a: &[Tensor], + lora_b: &[Tensor], + adam_m: &[Tensor], + adam_v: &[Tensor], +) -> Result<()> { + let adapter_prefix = adapter_id + .map(|id| format!("dynamic_{id}_")) + .unwrap_or_default(); + for (index, tensor) in lora_a.iter().enumerate() { + shards.push(tensor_shard( + parallel, + adapter_id, + global_lora_rank, + "adapter.safetensors", + format!("{adapter_prefix}a_{index}"), + "lora_a", + LoraSide::A, + tensor, + )?); + } + for (index, tensor) in lora_b.iter().enumerate() { + shards.push(tensor_shard( + parallel, + adapter_id, + global_lora_rank, + "adapter.safetensors", + format!("{adapter_prefix}b_{index}"), + "lora_b", + LoraSide::B, + tensor, + )?); + } + for (index, tensor) in adam_m.iter().enumerate() { + shards.push(tensor_shard( + parallel, + adapter_id, + global_lora_rank, + "optimizer.safetensors", + format!("{adapter_prefix}a_{index}"), + "adam_m", + if index % 2 == 0 { + LoraSide::A + } else { + LoraSide::B + }, + tensor, + )?); + } + for (index, tensor) in adam_v.iter().enumerate() { + shards.push(tensor_shard( + parallel, + adapter_id, + global_lora_rank, + "optimizer.safetensors", + format!("{adapter_prefix}b_{index}"), + "adam_v", + if index % 2 == 0 { + LoraSide::A + } else { + LoraSide::B + }, + tensor, + )?); + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +fn tensor_shard( + parallel: &ParallelCheckpointManifest, + adapter_id: Option, + global_lora_rank: i64, + file: &str, + tensor_name: String, + state: &str, + side: LoraSide, + tensor: &Tensor, +) -> Result { + let tp_size = i64::try_from(parallel.tensor_model_parallel_size) + .context("TP size exceeds checkpoint tensor shape range")?; + let tp_rank = i64::try_from(parallel.tensor_model_parallel_rank) + .context("TP rank exceeds checkpoint tensor shape range")?; + if global_lora_rank <= 0 || global_lora_rank % tp_size != 0 { + bail!( + "global LoRA rank {global_lora_rank} must be positive and divisible by TP size {tp_size}" + ); + } + let local_shape = tensor.size(); + if local_shape.len() < 2 { + bail!("checkpoint tensor {file}:{tensor_name} must have at least two dimensions"); + } + let partition_axis = match side { + LoraSide::A => local_shape.len() - 2, + LoraSide::B => local_shape.len() - 1, + }; + let local_lora_rank = global_lora_rank / tp_size; + if local_shape[partition_axis] != local_lora_rank { + bail!( + "checkpoint tensor {file}:{tensor_name} has local rank {} on axis {partition_axis}, expected {local_lora_rank}", + local_shape[partition_axis] + ); + } + let mut global_shape = local_shape.clone(); + global_shape[partition_axis] = global_lora_rank; + let mut global_offset = vec![0; local_shape.len()]; + global_offset[partition_axis] = tp_rank * local_lora_rank; + Ok(TensorShardManifest { + file: file.to_string(), + tensor_name, + state: state.to_string(), + adapter_id, + global_lora_rank, + global_shape, + local_shape, + partition_axis, + global_offset, + replica_identity: parallel.replica_identity(), + }) +} + +fn validate_saved_shards( + saved: &[TensorShardManifest], + expected: &[TensorShardManifest], +) -> Result<()> { + let keyed = |shards: &[TensorShardManifest]| -> Result> { + let mut by_key = BTreeMap::new(); + for shard in shards { + let key = ( + shard.file.clone(), + shard.tensor_name.clone(), + shard.state.clone(), + ); + if by_key.insert(key, shard.clone()).is_some() { + bail!( + "duplicate tensor shard metadata for {}:{} ({})", + shard.file, + shard.tensor_name, + shard.state + ); + } + } + Ok(by_key) + }; + let saved = keyed(saved)?; + let expected = keyed(expected)?; + if saved.len() != expected.len() { + bail!( + "tensor shard metadata count mismatch: saved={}, expected={}", + saved.len(), + expected.len() + ); + } + for (key, expected_shard) in expected { + let saved_shard = saved + .get(&key) + .with_context(|| format!("missing tensor shard metadata for {}:{}", key.0, key.1))?; + if saved_shard != &expected_shard { + bail!( + "tensor shard metadata mismatch for {}:{}: saved={saved_shard:?}, expected={expected_shard:?}", + key.0, + key.1 + ); + } + } + Ok(()) +} + +fn rank_checkpoint_dir(root: &Path, global_rank: usize) -> PathBuf { + root.join(format!("rank-{global_rank:05}")) +} + +fn env_usize(names: &[&str], default: usize) -> Result { + Ok(env_usize_optional(names)?.unwrap_or(default)) +} + +fn env_usize_optional(names: &[&str]) -> Result> { + let Some((name, value)) = names + .iter() + .find_map(|name| env::var(name).ok().map(|value| (*name, value))) + else { + return Ok(None); + }; + let value = value + .parse::() + .with_context(|| format!("{name} must be a non-negative integer"))?; + if value == 0 && name != "RANK" { + bail!("{name} must be positive"); + } + Ok(Some(value)) +} + fn named_tensors( a: &[Tensor], b: &[Tensor], @@ -308,6 +847,8 @@ mod tests { ) .unwrap(); let loaded = load_checkpoint(dir.path()).unwrap(); + assert!(loaded.manifest.parallel.is_none()); + assert!(loaded.manifest.tensor_shards.is_empty()); assert_eq!(loaded.manifest.lora_alpha, 4.0); assert_eq!(loaded.manifest.step, 7); assert_eq!(loaded.lora_a[0].size(), [2, 3]); @@ -361,6 +902,8 @@ mod tests { .unwrap(); let loaded = load_checkpoint(dir.path()).unwrap(); assert_eq!(loaded.manifest.format, "rustrain-checkpoint-v2"); + assert!(loaded.manifest.parallel.is_none()); + assert!(loaded.manifest.tensor_shards.is_empty()); assert_eq!(loaded.dynamic_adapters.len(), 1); let loaded_dynamic = &loaded.dynamic_adapters[0]; assert_eq!(loaded_dynamic.manifest.id, 7); @@ -391,4 +934,182 @@ mod tests { let manifest: DynamicAdapterManifest = serde_json::from_str(json).unwrap(); assert_eq!(manifest.optimizer_step, 0); } + + fn tp_topology(global_rank: usize, tp_size: usize) -> ParallelCheckpointManifest { + ParallelCheckpointManifest::new(tp_size, global_rank, tp_size, 1, 1, 1, 1).unwrap() + } + + fn tp_state(value: f64) -> (Vec, Vec, Vec, Vec) { + let a = Tensor::full([2, 3], value, (tch::Kind::Float, tch::Device::Cpu)); + let b = Tensor::full([4, 2], value + 1.0, (tch::Kind::Float, tch::Device::Cpu)); + let m = vec![a.zeros_like(), b.zeros_like()]; + let v = vec![a.ones_like(), b.ones_like()]; + (vec![a], vec![b], m, v) + } + + fn tp_dynamic_adapter(value: f64) -> DynamicAdapterCheckpoint { + let a = Tensor::full([3, 5], value, (tch::Kind::Float, tch::Device::Cpu)); + let b = Tensor::full([7, 3], value + 1.0, (tch::Kind::Float, tch::Device::Cpu)); + DynamicAdapterCheckpoint { + manifest: DynamicAdapterManifest { + id: 9, + rank: 6, + alpha: 12.0, + optimizer_step: 4, + target_layers: vec![1], + target_modules: vec!["q_proj".into()], + parameter_count: 1, + optimizer_count: 2, + }, + lora_a: vec![a.shallow_clone()], + lora_b: vec![b.shallow_clone()], + adam_m: vec![a.zeros_like(), b.zeros_like()], + adam_v: vec![a.ones_like(), b.ones_like()], + } + } + + #[test] + fn tensor_parallel_ranks_use_distinct_paths_and_resume_same_topology() { + let dir = tempfile::tempdir().unwrap(); + for global_rank in 0..2 { + let topology = tp_topology(global_rank, 2); + let (a, b, m, v) = tp_state(global_rank as f64 + 1.0); + let dynamic = tp_dynamic_adapter(global_rank as f64 + 10.0); + save_checkpoint_with_dynamic_for_topology( + dir.path(), + 7, + 0.25, + "Qwen/test", + 4, + 8.0, + &a, + &b, + &m, + &v, + &[dynamic], + &topology, + ) + .unwrap(); + } + + let rank0_dir = dir.path().join("rank-00000"); + let rank1_dir = dir.path().join("rank-00001"); + assert!(rank0_dir.join("manifest.json").is_file()); + assert!(rank1_dir.join("manifest.json").is_file()); + + let topology = tp_topology(1, 2); + let loaded = load_checkpoint_for_topology(dir.path(), &topology).unwrap(); + assert_eq!(loaded.manifest.format, TP_CHECKPOINT_FORMAT); + assert_eq!(loaded.manifest.parallel.as_ref(), Some(&topology)); + assert_eq!(loaded.manifest.tensor_shards.len(), 12); + assert_eq!(loaded.lora_a[0].double_value(&[0, 0]), 2.0); + assert_eq!(loaded.dynamic_adapters.len(), 1); + assert_eq!( + loaded.dynamic_adapters[0].lora_a[0].double_value(&[0, 0]), + 11.0 + ); + + let a_shard = loaded + .manifest + .tensor_shards + .iter() + .find(|shard| shard.state == "lora_a") + .unwrap(); + assert_eq!(a_shard.global_lora_rank, 4); + assert_eq!(a_shard.local_shape, vec![2, 3]); + assert_eq!(a_shard.global_shape, vec![4, 3]); + assert_eq!(a_shard.partition_axis, 0); + assert_eq!(a_shard.global_offset, vec![2, 0]); + assert_eq!(a_shard.replica_identity, "global-rank-1-tp-rank-1"); + + let b_shard = loaded + .manifest + .tensor_shards + .iter() + .find(|shard| shard.state == "lora_b") + .unwrap(); + assert_eq!(b_shard.global_shape, vec![4, 4]); + assert_eq!(b_shard.partition_axis, 1); + assert_eq!(b_shard.global_offset, vec![0, 2]); + + let dynamic_a_shard = loaded + .manifest + .tensor_shards + .iter() + .find(|shard| shard.state == "lora_a" && shard.adapter_id == Some(9)) + .unwrap(); + assert_eq!(dynamic_a_shard.global_lora_rank, 6); + assert_eq!(dynamic_a_shard.global_shape, vec![6, 5]); + assert_eq!(dynamic_a_shard.global_offset, vec![3, 0]); + } + + #[test] + fn tensor_parallel_resume_rejects_different_tp_size() { + let dir = tempfile::tempdir().unwrap(); + let topology = tp_topology(0, 2); + let (a, b, m, v) = tp_state(1.0); + save_checkpoint_with_dynamic_for_topology( + dir.path(), + 7, + 0.25, + "Qwen/test", + 4, + 8.0, + &a, + &b, + &m, + &v, + &[], + &topology, + ) + .unwrap(); + + let different_tp = tp_topology(0, 4); + let error = load_checkpoint_for_topology(dir.path(), &different_tp) + .err() + .expect("different TP size must fail"); + assert!(error.to_string().contains("topology mismatch")); + } + + #[test] + fn tensor_parallel_resume_rejects_another_ranks_shard() { + let dir = tempfile::tempdir().unwrap(); + let rank0 = tp_topology(0, 2); + let (a, b, m, v) = tp_state(1.0); + save_checkpoint_with_dynamic_for_topology( + dir.path(), + 7, + 0.25, + "Qwen/test", + 4, + 8.0, + &a, + &b, + &m, + &v, + &[], + &rank0, + ) + .unwrap(); + + let rank1_dir = dir.path().join("rank-00001"); + std::fs::create_dir(&rank1_dir).unwrap(); + for file in [ + "manifest.json", + "adapter.safetensors", + "optimizer.safetensors", + ] { + std::fs::copy( + dir.path().join("rank-00000").join(file), + rank1_dir.join(file), + ) + .unwrap(); + } + + let rank1 = tp_topology(1, 2); + let error = load_checkpoint_for_topology(dir.path(), &rank1) + .err() + .expect("loading another rank's shard must fail"); + assert!(error.to_string().contains("topology mismatch")); + } } diff --git a/crates/rustrain-server/src/session.rs b/crates/rustrain-server/src/session.rs index 878b2c13..d188b19e 100644 --- a/crates/rustrain-server/src/session.rs +++ b/crates/rustrain-server/src/session.rs @@ -700,7 +700,8 @@ impl TrainingSession for Qwen36Session { } } - checkpoint::save_checkpoint_with_dynamic( + let parallel = checkpoint::ParallelCheckpointManifest::from_env()?; + checkpoint::save_checkpoint_with_dynamic_for_topology( std::path::Path::new(path), self.step, self.last_loss, @@ -712,13 +713,15 @@ impl TrainingSession for Qwen36Session { &adam_m, &adam_v, &dynamic_adapters, + ¶llel, )?; Ok((self.step, self.last_loss)) } fn load_checkpoint(&mut self, path: &str) -> Result<(u64, f64)> { - let data = checkpoint::load_checkpoint(std::path::Path::new(path))?; + let parallel = checkpoint::ParallelCheckpointManifest::from_env()?; + let data = checkpoint::load_checkpoint_for_topology(std::path::Path::new(path), ¶llel)?; if !data.dynamic_adapters.is_empty() { let model_path = self .model_path From ff382a05e0d0f2a48610116a3242e736d12b3fa4 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 06:18:24 +0800 Subject: [PATCH 014/156] fix: reject unsafe native parallel topologies --- .../kernels/qwen3_6_kernels.cpp | 59 ++++++++++++++++++- .../rustrain-qwen3-6/tests/native_smoke.cpp | 29 +++++++++ 2 files changed, 86 insertions(+), 2 deletions(-) diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index ef8047c0..d6d82437 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -1584,6 +1584,9 @@ struct TrainingContext { cudaStream_t tp_stream = nullptr; int tp_world_size = 1; int tp_rank = 0; + // Set when a legacy NCCL setter supplies an incompatible mixed topology. + // Training entry points reject the context before touching parameters. + bool topology_invalid = false; int cuda_device = 0; // ────────────────────────────────────────────────────────────────────── }; @@ -3779,6 +3782,19 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( if (!tp_size_env) tp_size_env = getenv("RUSTRAIN_TP_SIZE"); ctx->tp_world_size = tp_size_env ? atoi(tp_size_env) : 1; TORCH_CHECK(ctx->tp_world_size > 0, "TP_SIZE must be positive"); + const char* world_size_env = getenv("WORLD_SIZE"); + const int configured_world_size = world_size_env ? atoi(world_size_env) : 1; + TORCH_CHECK(configured_world_size > 0, "WORLD_SIZE must be positive"); + const bool data_parallel_requested = env_enabled("RUSTRAIN_DATA_PARALLEL"); + TORCH_CHECK( + ctx->tp_world_size <= 1 || + (!data_parallel_requested && configured_world_size == ctx->tp_world_size), + "native Qwen LoRA supports TP-only topology when TP_SIZE>1; " + "TP_SIZE=", ctx->tp_world_size, " WORLD_SIZE=", configured_world_size, + " DATA_PARALLEL=", data_parallel_requested ? 1 : 0, + " is an incompatible mixed TP/DP/EP topology"); + ctx->ep_world_size = configured_world_size; + ctx->data_parallel = data_parallel_requested; const char* rank_env = getenv("RANK"); const int global_rank = rank_env ? atoi(rank_env) : 0; ctx->tp_rank = global_rank % ctx->tp_world_size; @@ -3977,6 +3993,8 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( ) { try { auto* ctx = reinterpret_cast(ctx_ptr); + TORCH_CHECK(!ctx->topology_invalid, + "native Qwen context rejected an incompatible TP/DP/EP topology"); GradientAccumulationFailureGuard accumulation_guard{ctx}; TORCH_CHECK(gradient_scale > 0.0 && std::isfinite(gradient_scale), "gradient_scale must be finite and positive"); @@ -4356,6 +4374,8 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( ) { try { auto* ctx = reinterpret_cast(ctx_ptr); + TORCH_CHECK(!ctx->topology_invalid, + "native Qwen context rejected an incompatible TP/DP/EP topology"); GradientAccumulationFailureGuard accumulation_guard{ctx}; if (ctx->nccl_comm) { c10::cuda::set_device(ctx->cuda_device); @@ -4376,6 +4396,10 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( TORCH_CHECK(input_batch == 1 || input_batch == n_total, "multi-LoRA input batch must be 1 or n_total (batch=", input_batch, ", n_total=", n_total, ")"); + TORCH_CHECK( + !(input_batch == n_total && ctx->data_parallel && ctx->ep_world_size > 1), + "dynamic multi-LoRA DP requires shared batch-1/equal tenant masks; " + "per-tenant batch with DP unsupported"); TORCH_CHECK(target_mask.size(0) == input_batch && target_mask.size(1) == input_ids.size(1), "target_mask must match input_ids shape"); @@ -4824,6 +4848,16 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( ctx->tp_world_size = tp_size_str2 ? atoi(tp_size_str2) : 1; ctx->tp_rank = ctx->tp_world_size > 0 ? ctx->ep_rank % ctx->tp_world_size : 0; + if (ctx->tp_world_size <= 0 || + (ctx->tp_world_size > 1 && + (ctx->data_parallel || ctx->ep_world_size != ctx->tp_world_size))) { + ctx->topology_invalid = true; + fprintf(stderr, + "[tp_nccl] reject mixed topology: TP_SIZE=%d WORLD_SIZE=%d DATA_PARALLEL=%d\n", + ctx->tp_world_size, ctx->ep_world_size, ctx->data_parallel ? 1 : 0); + return -1; + } + ctx->topology_invalid = false; ctx->tp_comm = ctx->tp_world_size > 1 ? g_tp_comm : nullptr; ctx->tp_stream = ctx->tp_world_size > 1 ? g_tp_stream : nullptr; // In TP-only mode the parent communicator is reserved for the TP @@ -4846,6 +4880,19 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( if (!rank_str || !world_str) return -1; int rank = atoi(rank_str); int world_size = atoi(world_str); + const char* tp_size_str = getenv("TP_SIZE"); + if (!tp_size_str) tp_size_str = getenv("RUSTRAIN_TP_SIZE"); + const int configured_tp_size = tp_size_str ? atoi(tp_size_str) : 1; + const bool data_parallel_requested = env_enabled("RUSTRAIN_DATA_PARALLEL"); + if (configured_tp_size <= 0 || + (configured_tp_size > 1 && + (data_parallel_requested || world_size != configured_tp_size))) { + ctx->topology_invalid = true; + fprintf(stderr, + "[tp_nccl] reject mixed topology: TP_SIZE=%d WORLD_SIZE=%d DATA_PARALLEL=%d\n", + configured_tp_size, world_size, data_parallel_requested ? 1 : 0); + return -1; + } if (world_size <= 1) return 0; // no EP needed // Set CUDA device and initialize PyTorch CUDA context on this device. @@ -4944,8 +4991,6 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( // Store as process-level singleton g_nccl_comm = comm; g_nccl_stream = nccl_stream; - const char* tp_size_str = getenv("TP_SIZE"); - if (!tp_size_str) tp_size_str = getenv("RUSTRAIN_TP_SIZE"); const int tp_size = tp_size_str ? atoi(tp_size_str) : 1; if (tp_size <= 0 || world_size % tp_size != 0) { fprintf(stderr, "[tp_nccl] invalid TP_SIZE=%d for WORLD_SIZE=%d\n", @@ -5003,6 +5048,16 @@ __attribute__((visibility("default"))) void qwen36_set_nccl_comm( ctx->ep_rank = ep_rank; ctx->ep_world_size = ep_world_size; ctx->data_parallel = env_enabled("RUSTRAIN_DATA_PARALLEL"); + if (ctx->tp_world_size <= 0 || + (ctx->tp_world_size > 1 && + (ctx->data_parallel || ep_world_size != ctx->tp_world_size))) { + ctx->topology_invalid = true; + fprintf(stderr, + "[tp_nccl] reject mixed topology: TP_SIZE=%d WORLD_SIZE=%d DATA_PARALLEL=%d\n", + ctx->tp_world_size, ep_world_size, ctx->data_parallel ? 1 : 0); + return; + } + ctx->topology_invalid = false; int current_device = g_cuda_device; cudaGetDevice(¤t_device); ctx->cuda_device = current_device; diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index 624221ac..6c10cb07 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -25,6 +25,7 @@ extern "C" void* qwen36_create_training_context( const int64_t*, int64_t, const char*); extern "C" int64_t qwen36_kernel_abi_version(); extern "C" int32_t qwen36_init_nccl(void*); +extern "C" void qwen36_set_nccl_comm(void*, void*, void*, int32_t, int32_t); extern "C" int64_t qwen36_get_lora_count(void*); extern "C" void* qwen36_get_lora_a(void*, int64_t); extern "C" void* qwen36_get_lora_b(void*, int64_t); @@ -128,6 +129,22 @@ int main() { config.nccl_stream = nullptr; const int64_t target_layer = 0; + // Native C++ must reject a mixed TP/DP/EP topology even when Rust-side + // validation is bypassed. The smoke process is single-rank, so simulate + // an invalid world before creating the real context. + if (world == 1 && tp_size == 1) { + setenv("TP_SIZE", "2", 1); + setenv("WORLD_SIZE", "4", 1); + void* invalid_topology_ctx = qwen36_create_training_context( + weight_ptrs.data(), static_cast(weight_ptrs.size()), + &embed, &final_norm, &lm_head, &config, 1, + static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank, + &target_layer, 1, "experts_gate_up_proj,experts_down_proj"); + assert(invalid_topology_ctx == nullptr); + setenv("TP_SIZE", "1", 1); + setenv("WORLD_SIZE", "1", 1); + } void* ctx = qwen36_create_training_context( weight_ptrs.data(), static_cast(weight_ptrs.size()), &embed, &final_norm, &lm_head, &config, 1, @@ -313,6 +330,18 @@ int main() { ctx, adapter_one, 0, "shared_gate_proj", 1) != nullptr); assert(qwen36_get_adapter_lora_tensor( ctx, adapter_two, 0, "shared_gate_proj", 1) != nullptr); + // Different tenant rows can have different token counts. The current + // native DP path only has one aggregate count, so reject batch=n_total + // instead of silently applying the wrong cross-rank weighting. + setenv("RUSTRAIN_DATA_PARALLEL", "1", 1); + qwen36_set_nccl_comm(ctx, nullptr, nullptr, 0, 2); + assert(qwen36_train_multi_lora( + ctx, &multi_input_ids, &multi_target_mask, &multi_attention_mask, + 2, rank) < 0.0); + assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 2); + assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 1); + unsetenv("RUSTRAIN_DATA_PARALLEL"); + qwen36_set_nccl_comm(ctx, nullptr, nullptr, 0, 1); qwen36_free_training_context(ctx); // Dense Qwen3.5 variants use the same per-sample activation path for From 2702bf03f17aa76cfa270d8dd0b76bbb58898150 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 06:20:13 +0800 Subject: [PATCH 015/156] docs: record abi11 tp checkpoint evidence --- docs/plans/qwen-lora-megatron-progress.md | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/docs/plans/qwen-lora-megatron-progress.md b/docs/plans/qwen-lora-megatron-progress.md index b69b3765..1d5fde1f 100644 --- a/docs/plans/qwen-lora-megatron-progress.md +++ b/docs/plans/qwen-lora-megatron-progress.md @@ -12,9 +12,9 @@ timestamp: 2026-07-17T00:00:00Z # Current State -Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP smoke, dense replicated-DP smoke, TP-only latent-rank-sharded LoRA smoke, dynamic batch logical-step update, ABI10 fixed and per-tenant optimizer-step restore, selected-tenant isolation, and 5D topology mapping. +Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP smoke, dense replicated-DP smoke, TP-only latent-rank-sharded LoRA smoke, dynamic batch logical-step update, ABI11 FP32 gradient storage/aggregation, per-tenant optimizer-step restore, selected-tenant isolation, same-topology rank-aware checkpointing, and 5D topology mapping. -Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axis TP+DP/EP, PP/CP, FP32 accumulation/abort, rank-sharded checkpoint topology, DeepEP/TE prebuilt integration, and matched Megatron throughput. +Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axis TP+DP/EP, PP/CP, variable-split token A2A, DeepEP/TE prebuilt integration, cross-topology checkpoint resharding, and matched Megatron throughput. # Durable Milestones @@ -22,9 +22,13 @@ Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axi - `afcf091`: restored native C++ Adam step on checkpoint load and added ABI8 smoke assertions; server checkpoint tests passed 2/2. - ABI9 working tree: selected adapter IDs flow through HTTP, IPC, Rust, and C++; each dynamic tenant owns its Adam clock, checkpoint metadata preserves it, and failed selection restores the complete registry. - ABI10 working tree: TP-only mode shards each fixed and dynamic adapter's latent rank, all-reduces only activation-level LoRA deltas on a split TP communicator, and keeps replicated MoE layers off the EP communicator. +- `5aa4ea4`: ABI11 stores fixed and dynamic LoRA leaf gradients in FP32 accumulators between micro-batches; fused Adam consumes those accumulators directly, fixed windows use token-weighted numerator/denominator reduction, and explicit abort/failure cleanup clears the pending window. This is FP32 storage/aggregation; the autograd leaf backward remains BF16-typed. +- `ed25daa`: TP rank-aware checkpoint format v3 writes `rank-xxxxx/` manifests and adapter/optimizer shards with topology, global/local shape, latent partition axis, offset, and replica identity metadata. v1/v2 single-rank loading remains compatible; v3 requires the same topology and rank. +- `ff382a0`: native ABI rejects unsupported TP+DP/EP mixtures before training. Dynamic per-tenant batch input under DP is explicitly rejected because the current reduction contract has one aggregate token count; shared batch-1/equal-mask dynamic DP remains the supported contract. - H20 `123.57.26.97:28004`: ABI8 native smoke passed grouped/fallback parity, GDN, dense/MoE LoRA, dynamic adapters, and step setter validation. - H20 `123.57.26.97:28004`: ABI9 native smoke passed selected-tenant training with a positive selected update, exactly zero unselected update, independent clocks (`2` vs `1`), and registry preservation after an unknown ID. - H20 `123.57.26.97:28004`: ABI10 two-rank TP native smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, and selected-tenant isolation. Rank-local LoRA tensors used distinct rank `4` slices for global rank `8`; losses matched and both shards had positive updates. +- H20 `123.57.26.97:28004`: ABI11 single-rank smoke passed FP32 accumulator dtype, two-micro accumulation, NaN/explicit abort cleanup, successful step commit, topology rejection, and dynamic-DP per-tenant batch rejection. ABI11 two-rank TP smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, selected-tenant isolation, and the new guards (`rank_statuses=0,0`). - Target runtime probe: PyTorch 2.5.1+cu121, ABI0; Transformer Engine, flash-attn, DeepEP, Triton, and DeepSpeed are not importable. # Decisions During Execution @@ -36,6 +40,6 @@ Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axi # Verification -Passed: `cargo test -p rustrain-core --lib` (8), `cargo test -p rustrain-parallel --lib` (14), `cargo test -p rustrain-server --lib` (3), Qwen unit tests (3), Qwen integration (6), remote ABI8 smoke, remote ABI9 selected-tenant native smoke, remote ABI10 single-rank native smoke, and remote ABI10 two-rank TP native smoke. +Passed: `cargo check -p rustrain-qwen3-6 -p rustrain-server -p rustrain-ipc`, `cargo test -p rustrain-server --lib` (6), Qwen unit tests (3), Qwen integration (6), remote ABI8 smoke, remote ABI9 selected-tenant native smoke, remote ABI10 single-rank native smoke, remote ABI10 two-rank TP native smoke, and remote ABI11 single/two-rank native smoke. -Not run: Megatron-style base-model TP, multi-axis TP+DP/EP, PP/CP, FP32 accumulation equivalence, rank-sharded checkpoint resume, and matched Megatron performance benchmark. +Not run: Megatron-style base-model TP, variable-split EP token dispatch, multi-axis TP+DP/EP, PP/CP, cross-topology resharding, a numerical FP32-accumulation oracle against a concatenated batch, and matched Megatron performance benchmark. From 76eace62d9c2876ae360c5c1fbda7c69c4fec9b8 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 06:24:24 +0800 Subject: [PATCH 016/156] fix: reject native pp and cp topologies --- crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp | 11 +++++++++++ crates/rustrain-qwen3-6/tests/native_smoke.cpp | 11 +++++++++++ 2 files changed, 22 insertions(+) diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index d6d82437..24b3b275 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -3786,6 +3786,17 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( const int configured_world_size = world_size_env ? atoi(world_size_env) : 1; TORCH_CHECK(configured_world_size > 0, "WORLD_SIZE must be positive"); const bool data_parallel_requested = env_enabled("RUSTRAIN_DATA_PARALLEL"); + const char* pp_size_env = getenv("PP_SIZE"); + if (!pp_size_env) pp_size_env = getenv("RUSTRAIN_PP_SIZE"); + const char* cp_size_env = getenv("CP_SIZE"); + if (!cp_size_env) cp_size_env = getenv("RUSTRAIN_CP_SIZE"); + const int configured_pp_size = pp_size_env ? atoi(pp_size_env) : 1; + const int configured_cp_size = cp_size_env ? atoi(cp_size_env) : 1; + TORCH_CHECK(configured_pp_size > 0 && configured_cp_size > 0, + "PP_SIZE and CP_SIZE must be positive"); + TORCH_CHECK(configured_pp_size == 1 && configured_cp_size == 1, + "native Qwen LoRA does not implement PP/CP yet; ", + "PP_SIZE=", configured_pp_size, " CP_SIZE=", configured_cp_size); TORCH_CHECK( ctx->tp_world_size <= 1 || (!data_parallel_requested && configured_world_size == ctx->tp_world_size), diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index 6c10cb07..80b93272 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -145,6 +145,17 @@ int main() { setenv("TP_SIZE", "1", 1); setenv("WORLD_SIZE", "1", 1); } + if (world == 1 && tp_size == 1) { + setenv("PP_SIZE", "2", 1); + void* invalid_pp_ctx = qwen36_create_training_context( + weight_ptrs.data(), static_cast(weight_ptrs.size()), + &embed, &final_norm, &lm_head, &config, 1, + static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank, + &target_layer, 1, "experts_gate_up_proj,experts_down_proj"); + assert(invalid_pp_ctx == nullptr); + unsetenv("PP_SIZE"); + } void* ctx = qwen36_create_training_context( weight_ptrs.data(), static_cast(weight_ptrs.size()), &embed, &final_norm, &lm_head, &config, 1, From 56a6ea1037bc1f0cf9eb930ddd88972022971349 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 06:24:44 +0800 Subject: [PATCH 017/156] docs: record native pp cp rejection --- docs/plans/qwen-lora-megatron-progress.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/plans/qwen-lora-megatron-progress.md b/docs/plans/qwen-lora-megatron-progress.md index 1d5fde1f..8b6af249 100644 --- a/docs/plans/qwen-lora-megatron-progress.md +++ b/docs/plans/qwen-lora-megatron-progress.md @@ -25,10 +25,11 @@ Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axi - `5aa4ea4`: ABI11 stores fixed and dynamic LoRA leaf gradients in FP32 accumulators between micro-batches; fused Adam consumes those accumulators directly, fixed windows use token-weighted numerator/denominator reduction, and explicit abort/failure cleanup clears the pending window. This is FP32 storage/aggregation; the autograd leaf backward remains BF16-typed. - `ed25daa`: TP rank-aware checkpoint format v3 writes `rank-xxxxx/` manifests and adapter/optimizer shards with topology, global/local shape, latent partition axis, offset, and replica identity metadata. v1/v2 single-rank loading remains compatible; v3 requires the same topology and rank. - `ff382a0`: native ABI rejects unsupported TP+DP/EP mixtures before training. Dynamic per-tenant batch input under DP is explicitly rejected because the current reduction contract has one aggregate token count; shared batch-1/equal-mask dynamic DP remains the supported contract. +- `76eace6`: direct native context creation also rejects PP/CP sizes greater than one, so unsupported pipeline/context parallelism cannot silently fall back to replicated single-stage execution. - H20 `123.57.26.97:28004`: ABI8 native smoke passed grouped/fallback parity, GDN, dense/MoE LoRA, dynamic adapters, and step setter validation. - H20 `123.57.26.97:28004`: ABI9 native smoke passed selected-tenant training with a positive selected update, exactly zero unselected update, independent clocks (`2` vs `1`), and registry preservation after an unknown ID. - H20 `123.57.26.97:28004`: ABI10 two-rank TP native smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, and selected-tenant isolation. Rank-local LoRA tensors used distinct rank `4` slices for global rank `8`; losses matched and both shards had positive updates. -- H20 `123.57.26.97:28004`: ABI11 single-rank smoke passed FP32 accumulator dtype, two-micro accumulation, NaN/explicit abort cleanup, successful step commit, topology rejection, and dynamic-DP per-tenant batch rejection. ABI11 two-rank TP smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, selected-tenant isolation, and the new guards (`rank_statuses=0,0`). +- H20 `123.57.26.97:28004`: ABI11 single-rank smoke passed FP32 accumulator dtype, two-micro accumulation, NaN/explicit abort cleanup, successful step commit, TP/DP/EP and PP/CP topology rejection, and dynamic-DP per-tenant batch rejection. ABI11 two-rank TP smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, selected-tenant isolation, and the new guards (`rank_statuses=0,0`). - Target runtime probe: PyTorch 2.5.1+cu121, ABI0; Transformer Engine, flash-attn, DeepEP, Triton, and DeepSpeed are not importable. # Decisions During Execution From b7050ea47e82d33791f00728589e87fe620e4a4f Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 07:03:04 +0800 Subject: [PATCH 018/156] fix: make dynamic lora dp and ep numerically correct --- .../kernels/qwen3_6_kernels.cpp | 214 +++++++--- .../tests/native_ep_smoke.cpp | 377 +++++++++++++++--- .../rustrain-qwen3-6/tests/native_smoke.cpp | 366 ++++++++++++++++- 3 files changed, 833 insertions(+), 124 deletions(-) diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 24b3b275..715a1699 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -1511,6 +1511,7 @@ struct TrainingContext { std::vector adapters; int64_t next_adapter_id = 0; + int64_t multi_lora_invocation = 0; // LoRA cache: pre-concatenated A/B per (layer, module) pair // Invalidated when adapters change or after Adam update @@ -1810,18 +1811,84 @@ static void reduce_lora_accumulator( accumulator.copy_(reduced); } +static void reduce_lora_accumulator_weighted( + TrainingContext* ctx, at::Tensor& accumulator, + const at::Tensor& local_weight, const at::Tensor& global_weight, + bool allreduce +) { + if (!accumulator.defined()) return; + TORCH_CHECK(accumulator.scalar_type() == at::kFloat, + "LoRA DP gradient accumulator must be FP32"); + TORCH_CHECK(local_weight.numel() == 1 && global_weight.numel() == 1, + "per-adapter LoRA token weights must be scalar"); + auto weighted = accumulator.contiguous() * local_weight; + at::Tensor reduced; + if (allreduce) { + TORCH_CHECK(ctx->nccl_comm, + "LoRA gradient all-reduce has no communicator"); + reduced = at::empty_like(weighted); + int dev = weighted.device().index(); + cudaSetDevice(dev); + auto stream = c10::cuda::getCurrentCUDAStream(dev).stream(); + auto err = ncclAllReduce( + weighted.data_ptr(), reduced.data_ptr(), weighted.numel(), + nccl_dtype_for(weighted), ncclSum, ctx->nccl_comm, stream); + TORCH_CHECK(err == ncclSuccess, + "NCCL weighted LoRA gradient all-reduce failed: ", + ncclGetErrorString(err)); + } else { + reduced = weighted; + } + reduced = reduced / global_weight.clamp_min(1.0); + at::NoGradGuard guard; + accumulator.copy_(reduced); +} + // Every rank evaluates the complete loss. Average replicated LoRA gradients // across DP ranks with token-count weighting so their Adam update matches a -// single global batch. EP ranks already receive the complete routed activation -// in forward and keep replicated gradients local; routed expert adapters remain -// local because their parameter tensors are sharded. +// single global batch. Pure DP keeps the complete routed-expert tensors +// replicated, so those accumulators reduce too; EP ranks receive the complete +// routed activation in forward and keep their sharded expert gradients local. static void synchronize_lora_gradients( TrainingContext* ctx, const at::Tensor& target_mask, - double accumulated_token_weight = 0.0 + double accumulated_token_weight = 0.0, + const at::Tensor* per_adapter_token_counts = nullptr ) { const bool allreduce = ctx->nccl_comm && ctx->data_parallel; + const bool per_adapter_weighting = per_adapter_token_counts && + per_adapter_token_counts->defined(); + at::Tensor local_adapter_weights; + at::Tensor global_adapter_weights; double scale = 1.0; - if (accumulated_token_weight > 0.0) { + if (per_adapter_weighting) { + TORCH_CHECK(per_adapter_token_counts->dim() == 1 && + per_adapter_token_counts->size(0) == + static_cast(ctx->adapters.size()), + "dynamic LoRA token-count vector must match adapter registry"); + local_adapter_weights = per_adapter_token_counts->to(at::kFloat) + .contiguous(); + TORCH_CHECK(at::isfinite(local_adapter_weights).all().item(), + "dynamic LoRA token counts must be finite"); + TORCH_CHECK((local_adapter_weights >= 0).all().item(), + "dynamic LoRA token counts must be non-negative"); + if (allreduce) { + global_adapter_weights = at::empty_like(local_adapter_weights); + auto stream = c10::cuda::getCurrentCUDAStream( + local_adapter_weights.device().index()).stream(); + auto err = ncclAllReduce( + local_adapter_weights.data_ptr(), + global_adapter_weights.data_ptr(), + local_adapter_weights.numel(), ncclFloat, ncclSum, + ctx->nccl_comm, stream); + TORCH_CHECK(err == ncclSuccess, + "NCCL per-adapter token-count all-reduce failed: ", + ncclGetErrorString(err)); + } else { + global_adapter_weights = local_adapter_weights; + } + TORCH_CHECK((global_adapter_weights > 0).all().item(), + "every dynamic LoRA adapter must have at least one global target token"); + } else if (accumulated_token_weight > 0.0) { double global_weight = accumulated_token_weight; if (allreduce) { auto local = at::full({1}, accumulated_token_weight, @@ -1857,7 +1924,9 @@ static void synchronize_lora_gradients( const double global_tokens = global_mask.item(); scale = local_tokens / std::max(global_tokens, 1.0); } - for (auto& adapter : ctx->adapters) { + for (size_t adapter_index = 0; + adapter_index < ctx->adapters.size(); ++adapter_index) { + auto& adapter = ctx->adapters[adapter_index]; for (auto& [layer_idx, pairs] : adapter.params) { auto table = lora_projection_table(ctx->layer_configs[layer_idx]); for (int64_t pair = 0; pair < (int64_t)pairs.size(); ++pair) { @@ -1865,17 +1934,29 @@ static void synchronize_lora_gradients( TORCH_CHECK(accum_it != adapter.grad_accum.end() && pair < (int64_t)accum_it->second.size(), "dynamic LoRA gradient accumulator layout mismatch"); - // Dynamic routed-expert tensors are sharded exactly like the - // base experts. In a fixed accumulation window they still - // need local token normalization, but must never be reduced - // across EP ranks; the legacy per-row path leaves them - // untouched to preserve its independent-sample contract. + if (per_adapter_weighting) { + auto local_weight = local_adapter_weights.index( + {static_cast(adapter_index)}); + auto global_weight = global_adapter_weights.index( + {static_cast(adapter_index)}); + reduce_lora_accumulator_weighted( + ctx, accum_it->second[pair][0], local_weight, + global_weight, allreduce); + reduce_lora_accumulator_weighted( + ctx, accum_it->second[pair][1], local_weight, + global_weight, allreduce); + continue; + } + // Routed-expert tensors are sharded only in EP. Pure DP has + // replicated experts and therefore uses the same reduction as + // shared projections; the legacy accumulation path preserves + // local normalization when no DP communicator is present. if (table.entries[pair].grouped_expert) { if (accumulated_token_weight > 0.0) { reduce_lora_accumulator( - ctx, accum_it->second[pair][0], scale, false); + ctx, accum_it->second[pair][0], scale, allreduce); reduce_lora_accumulator( - ctx, accum_it->second[pair][1], scale, false); + ctx, accum_it->second[pair][1], scale, allreduce); } continue; } @@ -1886,19 +1967,19 @@ static void synchronize_lora_gradients( } } } + if (per_adapter_weighting) return; for (int64_t layer = 0; layer < ctx->num_layers; ++layer) { auto table = lora_projection_table(ctx->layer_configs[layer]); int64_t offset = ctx->lora_layer_offset[layer]; for (int64_t pair = 0; pair < table.count; ++pair) { - // Routed expert LoRA is sharded with the base expert weights. Its - // local gradients belong only to this EP rank and must not be - // summed with a different expert shard on another rank. + // Routed expert LoRA is local-only for EP, while pure DP owns the + // complete replicated expert tensor and must all-reduce it. if (table.entries[pair].grouped_expert) { if (accumulated_token_weight > 0.0) { reduce_lora_accumulator( - ctx, ctx->grad_accum_a[offset + pair], scale, false); + ctx, ctx->grad_accum_a[offset + pair], scale, allreduce); reduce_lora_accumulator( - ctx, ctx->grad_accum_b[offset + pair], scale, false); + ctx, ctx->grad_accum_b[offset + pair], scale, allreduce); } continue; } @@ -4104,8 +4185,10 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( double step_f = (double)next_step; double bias_correction1 = 1.0 - std::pow(ctx->beta1, step_f); double bias_correction2 = 1.0 - std::pow(ctx->beta2, step_f); - float lr_scaled = (float)(ctx->lr / bias_correction1); - float eps_scaled = (float)(ctx->eps / std::sqrt(bias_correction2)); + double sqrt_bias_correction2 = std::sqrt(bias_correction2); + float lr_scaled = (float)( + ctx->lr * sqrt_bias_correction2 / bias_correction1); + float eps_scaled = (float)(ctx->eps * sqrt_bias_correction2); float one_minus_b1 = (float)(1.0 - ctx->beta1); float one_minus_b2 = (float)(1.0 - ctx->beta2); @@ -4407,13 +4490,16 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( TORCH_CHECK(input_batch == 1 || input_batch == n_total, "multi-LoRA input batch must be 1 or n_total (batch=", input_batch, ", n_total=", n_total, ")"); - TORCH_CHECK( - !(input_batch == n_total && ctx->data_parallel && ctx->ep_world_size > 1), - "dynamic multi-LoRA DP requires shared batch-1/equal tenant masks; " - "per-tenant batch with DP unsupported"); TORCH_CHECK(target_mask.size(0) == input_batch && target_mask.size(1) == input_ids.size(1), "target_mask must match input_ids shape"); + auto input_row_token_counts = target_mask + .narrow(1, 1, target_mask.size(1) - 1) + .to(at::kFloat).sum(1); + auto adapter_token_counts = input_batch == 1 + ? input_row_token_counts.repeat({total_adapters}) + : input_row_token_counts; + const int64_t multi_lora_invocation = ++ctx->multi_lora_invocation; // Keep the caller's mask intact. Each chunk receives either the // corresponding rows or a repeated batch-1 mask; this also prevents @@ -4469,10 +4555,11 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( // Use file-based barrier: rank 0 computes n_max, writes to file, others read. size_t free_mem, total_mem; cudaMemGetInfo(&free_mem, &total_mem); - int64_t n_max; + int64_t n_max = 0; if (ctx->nccl_comm && ctx->ep_world_size > 1) { const std::string sync_path = nccl_sync_dir() + "/nmax_sync_" + - std::to_string(ctx->context_sequence) + ".txt"; + std::to_string(ctx->context_sequence) + "_" + + std::to_string(multi_lora_invocation) + ".txt"; if (ctx->ep_rank == 0) { n_max = compute_n_max( (int64_t)free_mem, lora_rank, @@ -4481,15 +4568,32 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( ); n_max = std::min(n_max, total_adapters); if (n_max < 1) n_max = 1; - FILE* f = fopen(sync_path.c_str(), "w"); - fprintf(f, "%ld\n", (long)n_max); + const std::string temporary_path = sync_path + ".tmp." + + std::to_string(static_cast(getpid())); + FILE* f = fopen(temporary_path.c_str(), "w"); + TORCH_CHECK(f, "failed to create n_max rendezvous file: ", + temporary_path); + TORCH_CHECK(fprintf(f, "%ld\n", (long)n_max) > 0, + "failed to write n_max rendezvous file: ", temporary_path); fclose(f); + TORCH_CHECK(rename(temporary_path.c_str(), sync_path.c_str()) == 0, + "failed to publish n_max rendezvous file: ", sync_path); } else { + bool loaded = false; for (int i = 0; i < 600; i++) { FILE* f = fopen(sync_path.c_str(), "r"); - if (f) { fscanf(f, "%ld", (long*)&n_max); fclose(f); break; } + if (f) { + const int parsed = fscanf(f, "%ld", (long*)&n_max); + fclose(f); + TORCH_CHECK(parsed == 1, + "invalid n_max rendezvous file: ", sync_path); + loaded = true; + break; + } usleep(10000); } + TORCH_CHECK(loaded, + "timed out waiting for n_max rendezvous file: ", sync_path); } } else { n_max = compute_n_max( @@ -4620,7 +4724,8 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( if (chunk == num_chunks - 1) { // DP gradient synchronization and Adam belong to the logical // multi-tenant step, never to an activation-memory chunk. - synchronize_lora_gradients(ctx, target_mask); + synchronize_lora_gradients( + ctx, target_mask, 0.0, &adapter_token_counts); // Adam step. Group tenants by their own logical clock so // newly-added or resumed tenants do not inherit another @@ -4666,8 +4771,16 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( } if (h_params.empty()) continue; const double step_f = (double)logical_step; - const float lr_scaled = (float)(ctx->lr / (1.0 - std::pow(ctx->beta1, step_f))); - const float eps_scaled = (float)(ctx->eps / std::sqrt(1.0 - std::pow(ctx->beta2, step_f))); + const double bias_correction1 = + 1.0 - std::pow(ctx->beta1, step_f); + const double bias_correction2 = + 1.0 - std::pow(ctx->beta2, step_f); + const double sqrt_bias_correction2 = + std::sqrt(bias_correction2); + const float lr_scaled = (float)( + ctx->lr * sqrt_bias_correction2 / bias_correction1); + const float eps_scaled = (float)( + ctx->eps * sqrt_bias_correction2); const float one_minus_b1 = (float)(1.0 - ctx->beta1); const float one_minus_b2 = (float)(1.0 - ctx->beta2); int n_params = (int)h_params.size(); @@ -4874,13 +4987,17 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( // In TP-only mode the parent communicator is reserved for the TP // split; EP layer collectives must remain disabled on replicated MoE. if (ctx->tp_world_size <= 1) { + void* layer_comm = ctx->data_parallel + ? nullptr : (void*)g_nccl_comm; + void* layer_stream = ctx->data_parallel + ? nullptr : (void*)g_nccl_stream; for (auto& lc : ctx->layer_configs) { - lc.nccl_comm = (void*)g_nccl_comm; - lc.nccl_stream = (void*)g_nccl_stream; + lc.nccl_comm = layer_comm; + lc.nccl_stream = layer_stream; } for (auto& lc : ctx->mtp_layer_configs) { - lc.nccl_comm = (void*)g_nccl_comm; - lc.nccl_stream = (void*)g_nccl_stream; + lc.nccl_comm = layer_comm; + lc.nccl_stream = layer_stream; } } return 0; @@ -5035,13 +5152,16 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( // Propagate to layer configs if (tp_size <= 1) { + void* layer_comm = ctx->data_parallel ? nullptr : (void*)comm; + void* layer_stream = ctx->data_parallel + ? nullptr : (void*)nccl_stream; for (auto& lc : ctx->layer_configs) { - lc.nccl_comm = (void*)comm; - lc.nccl_stream = (void*)nccl_stream; + lc.nccl_comm = layer_comm; + lc.nccl_stream = layer_stream; } for (auto& lc : ctx->mtp_layer_configs) { - lc.nccl_comm = (void*)comm; - lc.nccl_stream = (void*)nccl_stream; + lc.nccl_comm = layer_comm; + lc.nccl_stream = layer_stream; } } @@ -5072,14 +5192,18 @@ __attribute__((visibility("default"))) void qwen36_set_nccl_comm( int current_device = g_cuda_device; cudaGetDevice(¤t_device); ctx->cuda_device = current_device; - // Propagate NCCL handles to all layer configs so moe_forward can access them + // Only EP owns routed-output collectives. In pure replicated DP the world + // communicator belongs exclusively to LoRA gradient synchronization; + // exposing it to moe_forward would mix activations from unrelated samples. + void* layer_comm = ctx->data_parallel ? nullptr : comm_ptr; + void* layer_stream = ctx->data_parallel ? nullptr : stream_ptr; for (auto& lc : ctx->layer_configs) { - lc.nccl_comm = comm_ptr; - lc.nccl_stream = stream_ptr; + lc.nccl_comm = layer_comm; + lc.nccl_stream = layer_stream; } for (auto& lc : ctx->mtp_layer_configs) { - lc.nccl_comm = comm_ptr; - lc.nccl_stream = stream_ptr; + lc.nccl_comm = layer_comm; + lc.nccl_stream = layer_stream; } } diff --git a/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp index 29799d57..1f2c33b5 100644 --- a/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp @@ -1,7 +1,9 @@ #include #include +#include #include +#include #include #include #include @@ -22,8 +24,11 @@ extern "C" void* qwen36_create_training_context( double, double, double, double, double, int64_t, double, int64_t, const int64_t*, int64_t, const char*); extern "C" int32_t qwen36_init_nccl(void*); +extern "C" int64_t qwen36_get_lora_count(void*); +extern "C" void* qwen36_get_lora_a(void*, int64_t); extern "C" void* qwen36_get_lora_b(void*, int64_t); extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*); +extern "C" int64_t qwen36_export_optimizer_state(void*, void**, void**, int64_t); extern "C" double qwen36_train_step(void*, void*, void*, void*); extern "C" void qwen36_free_training_context(void*); @@ -31,39 +36,123 @@ static at::Tensor cuda_rand(std::initializer_list shape) { return at::randn(shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)); } +static std::vector tensor_ptrs(std::vector& tensors) { + std::vector ptrs; + ptrs.reserve(tensors.size()); + for (auto& tensor : tensors) ptrs.push_back(&tensor); + return ptrs; +} + +static double max_abs_diff(const at::Tensor& actual, const at::Tensor& expected) { + assert(actual.sizes() == expected.sizes()); + return (actual.to(at::kFloat) - expected.to(at::kFloat)) + .abs().max().item(); +} + +static double update_norm(const at::Tensor& after, const at::Tensor& before) { + return (after.to(at::kFloat) - before.to(at::kFloat)) + .abs().sum().item(); +} + +static double first_adam_step_diff( + const at::Tensor& actual, + const at::Tensor& before, + const at::Tensor& first_m, + const at::Tensor& first_v +) { + constexpr double lr = 1e-3; + constexpr double beta1 = 0.9; + constexpr double beta2 = 0.999; + constexpr double eps = 1e-8; + auto m_hat = first_m.to(at::kFloat) / (1.0 - beta1); + auto v_hat = first_v.to(at::kFloat) / (1.0 - beta2); + auto expected = (before.to(at::kFloat) - + lr * m_hat / (v_hat.sqrt() + eps)).to(actual.scalar_type()); + return max_abs_diff(actual, expected); +} + +struct ExpertLora { + at::Tensor* gate_up_a; + at::Tensor* gate_up_b; + at::Tensor* down_a; + at::Tensor* down_b; +}; + +static ExpertLora get_expert_lora(void* ctx) { + assert(qwen36_get_lora_count(ctx) == 9); + ExpertLora lora{ + reinterpret_cast(qwen36_get_lora_a(ctx, 7)), + reinterpret_cast(qwen36_get_lora_b(ctx, 7)), + reinterpret_cast(qwen36_get_lora_a(ctx, 8)), + reinterpret_cast(qwen36_get_lora_b(ctx, 8)), + }; + assert(lora.gate_up_a && lora.gate_up_b && lora.down_a && lora.down_b); + return lora; +} + +static void set_expert_lora( + void* ctx, + at::Tensor& gate_up_a, + at::Tensor& gate_up_b, + at::Tensor& down_a, + at::Tensor& down_b +) { + assert(qwen36_set_lora_tensor(ctx, 7, 0, &gate_up_a) == 0); + assert(qwen36_set_lora_tensor(ctx, 7, 1, &gate_up_b) == 0); + assert(qwen36_set_lora_tensor(ctx, 8, 0, &down_a) == 0); + assert(qwen36_set_lora_tensor(ctx, 8, 1, &down_b) == 0); +} + int main() { const int rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0"); const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); - const int local_rank = std::atoi(std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); + const int local_rank = std::atoi( + std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); assert(world == 2 && rank >= 0 && rank < world); + assert(!std::getenv("TP_SIZE") || std::atoi(std::getenv("TP_SIZE")) == 1); c10::cuda::CUDAGuard guard(local_rank); - // Replicated weights must be identical across EP ranks. - at::manual_seed(100); + // Every process deterministically creates the same global model. The EP + // context receives a distinct contiguous expert slice from that model. + at::manual_seed(100); constexpr int64_t hidden = 16; constexpr int64_t vocab = 8; constexpr int64_t experts = 2; constexpr int64_t head_dim = 8; constexpr int64_t intermediate = 8; - constexpr int64_t rank_lora = 4; - std::vector weights; - weights.push_back(cuda_rand({hidden})); - weights.push_back(cuda_rand({hidden})); - weights.push_back(cuda_rand({2 * head_dim, hidden})); - weights.push_back(cuda_rand({head_dim})); - weights.push_back(cuda_rand({head_dim, hidden})); - weights.push_back(cuda_rand({head_dim})); - weights.push_back(cuda_rand({head_dim, hidden})); - weights.push_back(cuda_rand({hidden, head_dim})); - weights.push_back(cuda_rand({experts, hidden})); - weights.push_back(cuda_rand({1, hidden})); - weights.push_back(cuda_rand({intermediate, hidden})); - weights.push_back(cuda_rand({intermediate, hidden})); - weights.push_back(cuda_rand({hidden, intermediate})); - // Each process owns exactly one expert row. - weights.push_back(cuda_rand({1, 2 * intermediate, hidden})); - weights.push_back(cuda_rand({1, hidden, intermediate})); - for (auto& weight : weights) weight.set_requires_grad(false); + constexpr int64_t lora_rank = 4; + + std::vector global_weights; + global_weights.push_back(cuda_rand({hidden})); + global_weights.push_back(cuda_rand({hidden})); + global_weights.push_back(cuda_rand({2 * head_dim, hidden})); + global_weights.push_back(cuda_rand({head_dim})); + global_weights.push_back(cuda_rand({head_dim, hidden})); + global_weights.push_back(cuda_rand({head_dim})); + global_weights.push_back(cuda_rand({head_dim, hidden})); + global_weights.push_back(cuda_rand({hidden, head_dim})); + global_weights.push_back(cuda_rand({experts, hidden})); + global_weights.push_back(cuda_rand({1, hidden})); + global_weights.push_back(cuda_rand({intermediate, hidden})); + global_weights.push_back(cuda_rand({intermediate, hidden})); + global_weights.push_back(cuda_rand({hidden, intermediate})); + global_weights.push_back(cuda_rand({experts, 2 * intermediate, hidden})); + global_weights.push_back(cuda_rand({experts, hidden, intermediate})); + for (auto& weight : global_weights) weight.set_requires_grad(false); + + assert(max_abs_diff( + global_weights[13].narrow(0, 0, 1), + global_weights[13].narrow(0, 1, 1)) > 0.0); + assert(max_abs_diff( + global_weights[14].narrow(0, 0, 1), + global_weights[14].narrow(0, 1, 1)) > 0.0); + + std::vector distributed_weights = global_weights; + distributed_weights[13] = global_weights[13].narrow(0, rank, 1).contiguous(); + distributed_weights[14] = global_weights[14].narrow(0, rank, 1).contiguous(); + auto distributed_weight_ptrs = tensor_ptrs(distributed_weights); + auto reference_weight_ptrs = tensor_ptrs(global_weights); + auto embed = cuda_rand({vocab, hidden}); auto final_norm = cuda_rand({hidden}); auto lm_head = cuda_rand({vocab, hidden}); @@ -71,53 +160,213 @@ int main() { final_norm.set_requires_grad(false); lm_head.set_requires_grad(false); - std::vector weight_ptrs; - for (auto& weight : weights) weight_ptrs.push_back(&weight); - LayerConfig config{}; - config.layer_type = 0; - config.num_heads = 1; - config.num_kv_heads = 1; - config.head_dim = head_dim; - config.partial_rotary_factor = 1.0; - config.rope_theta = 10000.0; - config.rms_eps = 1e-5; - config.num_experts = experts; - // Route every token to both experts so both local shards exercise their - // expert LoRA optimizer path in this two-rank smoke. - config.top_k = 2; - config.moe_intermediate = intermediate; - config.expert_start = rank; - config.expert_count = 1; - config.norm_topk_prob = 1; + LayerConfig distributed_config{}; + distributed_config.layer_type = 0; + distributed_config.num_heads = 1; + distributed_config.num_kv_heads = 1; + distributed_config.head_dim = head_dim; + distributed_config.partial_rotary_factor = 1.0; + distributed_config.rope_theta = 10000.0; + distributed_config.rms_eps = 1e-5; + distributed_config.num_experts = experts; + // With two experts, top_k=2 guarantees every token exercises both EP + // shards and both corresponding expert LoRA optimizer paths. + distributed_config.top_k = 2; + distributed_config.moe_intermediate = intermediate; + distributed_config.expert_start = rank; + distributed_config.expert_count = 1; + distributed_config.norm_topk_prob = 1; + distributed_config.nccl_comm = nullptr; + distributed_config.nccl_stream = nullptr; + + LayerConfig reference_config = distributed_config; + reference_config.expert_start = 0; + reference_config.expert_count = experts; const int64_t target_layer = 0; - void* ctx = qwen36_create_training_context( - weight_ptrs.data(), static_cast(weight_ptrs.size()), - &embed, &final_norm, &lm_head, &config, 1, + const char* targets = "experts_gate_up_proj,experts_down_proj"; + void* distributed_ctx = qwen36_create_training_context( + distributed_weight_ptrs.data(), + static_cast(distributed_weight_ptrs.size()), + &embed, &final_norm, &lm_head, &distributed_config, 1, + static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + &target_layer, 1, targets); + assert(distributed_ctx); + + // The reference owns all experts and deliberately never initializes NCCL. + // Its copied LayerConfig therefore retains null communication handles even + // though WORLD_SIZE remains two for the distributed process. + void* reference_ctx = qwen36_create_training_context( + reference_weight_ptrs.data(), + static_cast(reference_weight_ptrs.size()), + &embed, &final_norm, &lm_head, &reference_config, 1, static_cast(at::kBFloat16), - 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, rank_lora, - &target_layer, 1, "experts_gate_up_proj,experts_down_proj"); - assert(ctx); - assert(qwen36_init_nccl(ctx) == 0); - auto* lora_b = reinterpret_cast(qwen36_get_lora_b(ctx, 7)); - assert(lora_b); - auto lora_b_value = at::ones(lora_b->sizes(), lora_b->options()); - assert(qwen36_set_lora_tensor(ctx, 7, 1, &lora_b_value) == 0); - auto before = lora_b->clone(); - - auto input_ids = at::tensor({1, 2}, - at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 2}); - auto target_mask = at::ones({1, 2}, + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + &target_layer, 1, targets); + assert(reference_ctx); + + auto opts = at::TensorOptions().device(at::kCUDA).dtype(at::kFloat); + auto global_gate_up_a = + ((at::arange(experts * lora_rank * hidden, opts) + 1.0) * 1e-4) + .reshape({experts, lora_rank, hidden}).to(at::kBFloat16); + auto global_gate_up_b = + ((at::arange(experts * 2 * intermediate * lora_rank, opts) + 3.0) * 5e-5) + .reshape({experts, 2 * intermediate, lora_rank}).to(at::kBFloat16); + auto global_down_a = + ((at::arange(experts * lora_rank * intermediate, opts) + 5.0) * 8e-5) + .reshape({experts, lora_rank, intermediate}).to(at::kBFloat16); + auto global_down_b = + ((at::arange(experts * hidden * lora_rank, opts) + 7.0) * 6e-5) + .reshape({experts, hidden, lora_rank}).to(at::kBFloat16); + + assert(max_abs_diff( + global_gate_up_a.narrow(0, 0, 1), + global_gate_up_a.narrow(0, 1, 1)) > 0.0); + assert(max_abs_diff( + global_down_b.narrow(0, 0, 1), + global_down_b.narrow(0, 1, 1)) > 0.0); + + auto local_gate_up_a = global_gate_up_a.narrow(0, rank, 1).contiguous(); + auto local_gate_up_b = global_gate_up_b.narrow(0, rank, 1).contiguous(); + auto local_down_a = global_down_a.narrow(0, rank, 1).contiguous(); + auto local_down_b = global_down_b.narrow(0, rank, 1).contiguous(); + set_expert_lora(distributed_ctx, + local_gate_up_a, local_gate_up_b, local_down_a, local_down_b); + set_expert_lora(reference_ctx, + global_gate_up_a, global_gate_up_b, global_down_a, global_down_b); + + auto distributed_lora = get_expert_lora(distributed_ctx); + auto reference_lora = get_expert_lora(reference_ctx); + assert(distributed_lora.gate_up_a->sizes() == + at::IntArrayRef({1, lora_rank, hidden})); + assert(reference_lora.gate_up_a->sizes() == + at::IntArrayRef({experts, lora_rank, hidden})); + assert(distributed_lora.gate_up_b->sizes() == + at::IntArrayRef({1, 2 * intermediate, lora_rank})); + assert(distributed_lora.down_a->sizes() == + at::IntArrayRef({1, lora_rank, intermediate})); + assert(distributed_lora.down_b->sizes() == + at::IntArrayRef({1, hidden, lora_rank})); + + auto gate_up_a_before = distributed_lora.gate_up_a->clone(); + auto gate_up_b_before = distributed_lora.gate_up_b->clone(); + auto down_a_before = distributed_lora.down_a->clone(); + auto down_b_before = distributed_lora.down_b->clone(); + + // Initialize NCCL only after the reference context is fully constructed. + // qwen36_init_nccl mutates only the distributed context's copied configs. + assert(qwen36_init_nccl(distributed_ctx) == 0); + + auto input_ids = at::tensor({1, 2, 3}, + at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 3}); + auto target_mask = at::ones({1, 3}, at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); - auto attention_mask = at::ones({1, 2}, + auto attention_mask = at::ones({1, 3}, at::TensorOptions().device(at::kCUDA).dtype(at::kBool)); - const double loss = qwen36_train_step(ctx, &input_ids, &target_mask, &attention_mask); + + const double distributed_loss = qwen36_train_step( + distributed_ctx, &input_ids, &target_mask, &attention_mask); + const double reference_loss = qwen36_train_step( + reference_ctx, &input_ids, &target_mask, &attention_mask); c10::cuda::device_synchronize(); - const double update = (*lora_b - before).abs().sum().item(); - std::printf("native_qwen36_ep_smoke rank=%d world=%d loss=%0.8f lora_b_update=%0.8e\n", - rank, world, loss, update); - assert(loss == loss && loss > 0.0); - assert(update > 0.0); - qwen36_free_training_context(ctx); + assert(distributed_loss > 0.0 && std::isfinite(distributed_loss)); + assert(reference_loss > 0.0 && std::isfinite(reference_loss)); + + const auto reference_slice = [rank](const at::Tensor& tensor) { + return tensor.narrow(0, rank, 1); + }; + const double gate_up_a_diff = max_abs_diff( + *distributed_lora.gate_up_a, reference_slice(*reference_lora.gate_up_a)); + const double gate_up_b_diff = max_abs_diff( + *distributed_lora.gate_up_b, reference_slice(*reference_lora.gate_up_b)); + const double down_a_diff = max_abs_diff( + *distributed_lora.down_a, reference_slice(*reference_lora.down_a)); + const double down_b_diff = max_abs_diff( + *distributed_lora.down_b, reference_slice(*reference_lora.down_b)); + const double loss_diff = std::abs(distributed_loss - reference_loss); + + const double gate_up_a_update = update_norm( + *distributed_lora.gate_up_a, gate_up_a_before); + const double gate_up_b_update = update_norm( + *distributed_lora.gate_up_b, gate_up_b_before); + const double down_a_update = update_norm( + *distributed_lora.down_a, down_a_before); + const double down_b_update = update_norm( + *distributed_lora.down_b, down_b_before); + + constexpr int64_t optimizer_slots = 18; + std::vector distributed_m(optimizer_slots), distributed_v(optimizer_slots); + std::vector reference_m(optimizer_slots), reference_v(optimizer_slots); + assert(qwen36_export_optimizer_state(distributed_ctx, + distributed_m.data(), distributed_v.data(), optimizer_slots) == optimizer_slots); + assert(qwen36_export_optimizer_state(reference_ctx, + reference_m.data(), reference_v.data(), optimizer_slots) == optimizer_slots); + + double optimizer_m_diff = 0.0; + double optimizer_v_diff = 0.0; + for (int64_t state_idx = 14; state_idx < optimizer_slots; ++state_idx) { + auto* local_m = reinterpret_cast(distributed_m[state_idx]); + auto* local_v = reinterpret_cast(distributed_v[state_idx]); + auto* full_m = reinterpret_cast(reference_m[state_idx]); + auto* full_v = reinterpret_cast(reference_v[state_idx]); + assert(local_m && local_v && full_m && full_v); + optimizer_m_diff = std::max( + optimizer_m_diff, max_abs_diff(*local_m, reference_slice(*full_m))); + optimizer_v_diff = std::max( + optimizer_v_diff, max_abs_diff(*local_v, reference_slice(*full_v))); + } + + const double gate_up_a_adam_diff = first_adam_step_diff( + *distributed_lora.gate_up_a, gate_up_a_before, + *reinterpret_cast(distributed_m[14]), + *reinterpret_cast(distributed_v[14])); + const double gate_up_b_adam_diff = first_adam_step_diff( + *distributed_lora.gate_up_b, gate_up_b_before, + *reinterpret_cast(distributed_m[15]), + *reinterpret_cast(distributed_v[15])); + const double down_a_adam_diff = first_adam_step_diff( + *distributed_lora.down_a, down_a_before, + *reinterpret_cast(distributed_m[16]), + *reinterpret_cast(distributed_v[16])); + const double down_b_adam_diff = first_adam_step_diff( + *distributed_lora.down_b, down_b_before, + *reinterpret_cast(distributed_m[17]), + *reinterpret_cast(distributed_v[17])); + + std::printf( + "native_qwen36_ep_parity rank=%d world=%d top_k=2 " + "distributed_loss=%0.8f reference_loss=%0.8f loss_diff=%0.8e " + "gate_up_a_diff=%0.8e gate_up_b_diff=%0.8e " + "down_a_diff=%0.8e down_b_diff=%0.8e " + "adam_m_diff=%0.8e adam_v_diff=%0.8e " + "adam_step_diffs=[%0.8e,%0.8e,%0.8e,%0.8e] " + "updates=[%0.8e,%0.8e,%0.8e,%0.8e]\n", + rank, world, distributed_loss, reference_loss, loss_diff, + gate_up_a_diff, gate_up_b_diff, down_a_diff, down_b_diff, + optimizer_m_diff, optimizer_v_diff, + gate_up_a_adam_diff, gate_up_b_adam_diff, + down_a_adam_diff, down_b_adam_diff, + gate_up_a_update, gate_up_b_update, down_a_update, down_b_update); + std::fflush(stdout); + + assert(gate_up_a_update > 0.0); + assert(gate_up_b_update > 0.0); + assert(down_a_update > 0.0); + assert(down_b_update > 0.0); + assert(loss_diff <= 2e-2); + assert(gate_up_a_diff <= 1e-5); + assert(gate_up_b_diff <= 1e-5); + assert(down_a_diff <= 1e-5); + assert(down_b_diff <= 1e-5); + assert(optimizer_m_diff <= 1e-5); + assert(optimizer_v_diff <= 1e-6); + assert(gate_up_a_adam_diff <= 1e-5); + assert(gate_up_b_adam_diff <= 1e-5); + assert(down_a_adam_diff <= 1e-5); + assert(down_b_adam_diff <= 1e-5); + + qwen36_free_training_context(reference_ctx); + qwen36_free_training_context(distributed_ctx); return 0; } diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index 80b93272..4a363df1 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -2,10 +2,15 @@ #include #include +#include #include #include #include #include +#include +#include +#include +#include #include struct LayerConfig { @@ -37,6 +42,8 @@ extern "C" double qwen36_train_micro_step( void*, void*, void*, void*, double, int32_t); extern "C" int64_t qwen36_get_step_count(void*); extern "C" int32_t qwen36_set_step_count(void*, int64_t); +extern "C" int64_t qwen36_export_optimizer_state( + void*, void**, void**, int64_t); extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t); extern "C" int32_t qwen36_set_adapter_step_count(void*, int64_t, int64_t); extern "C" double qwen36_eval_step(void*, void*, void*, void*); @@ -50,12 +57,314 @@ extern "C" void* qwen36_get_adapter_lora_tensor( void*, int64_t, int64_t, const char*, int32_t); extern "C" int32_t qwen36_set_adapter_lora_tensor( void*, int64_t, int64_t, const char*, int32_t, void*); +extern "C" void* qwen36_get_adapter_optimizer_tensor( + void*, int64_t, int64_t, const char*, int32_t, int32_t); extern "C" void qwen36_free_training_context(void*); static at::Tensor cuda_rand(std::initializer_list shape) { return at::randn(shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)); } +static std::string dp_smoke_sync_dir() { + const char* run_id = std::getenv("RUSTRAIN_NCCL_RUN_ID"); + std::string sanitized = run_id && run_id[0] ? run_id : "native-dp-smoke"; + for (char& ch : sanitized) { + if (!((ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || + (ch >= '0' && ch <= '9') || ch == '-' || ch == '_')) { + ch = '_'; + } + } + mkdir("/tmp/rustrain-nccl", 0777); + std::string path = "/tmp/rustrain-nccl/" + sanitized; + mkdir(path.c_str(), 0777); + return path; +} + +static void write_tensor_file(const std::string& path, const at::Tensor& tensor) { + auto cpu = tensor.to(at::kCPU).to(at::kFloat).contiguous(); + const std::string temporary_path = path + ".tmp." + + std::to_string(static_cast(getpid())); + FILE* file = std::fopen(temporary_path.c_str(), "wb"); + assert(file); + const int64_t rows = cpu.size(0); + const int64_t cols = cpu.size(1); + assert(std::fwrite(&rows, sizeof(rows), 1, file) == 1); + assert(std::fwrite(&cols, sizeof(cols), 1, file) == 1); + assert(std::fwrite(cpu.data_ptr(), sizeof(float), cpu.numel(), file) == + static_cast(cpu.numel())); + std::fclose(file); + assert(std::rename(temporary_path.c_str(), path.c_str()) == 0); +} + +static at::Tensor read_tensor_file(const std::string& path) { + FILE* file = nullptr; + for (int attempt = 0; attempt < 6000 && !file; ++attempt) { + file = std::fopen(path.c_str(), "rb"); + if (!file) usleep(10000); + } + assert(file); + int64_t rows = 0; + int64_t cols = 0; + assert(std::fread(&rows, sizeof(rows), 1, file) == 1); + assert(std::fread(&cols, sizeof(cols), 1, file) == 1); + auto result = at::empty({rows, cols}, at::TensorOptions().dtype(at::kFloat)); + assert(std::fread(result.data_ptr(), sizeof(float), result.numel(), file) == + static_cast(result.numel())); + std::fclose(file); + return result; +} + +struct DpAdapterInitial { + at::Tensor shared_a; + at::Tensor shared_b; + at::Tensor expert_a; + at::Tensor expert_b; +}; + +static DpAdapterInitial capture_dp_adapter(void* ctx, int64_t adapter_id) { + auto* shared_a = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_id, 0, "shared_gate_proj", 0)); + auto* shared_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_id, 0, "shared_gate_proj", 1)); + auto* expert_a = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_id, 0, "experts_gate_up_proj", 0)); + auto* expert_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_id, 0, "experts_gate_up_proj", 1)); + assert(shared_a && shared_b && expert_a && expert_b); + return { + shared_a->clone(), shared_b->clone(), + expert_a->clone(), expert_b->clone() + }; +} + +static void restore_dp_adapter( + void* ctx, int64_t adapter_id, const DpAdapterInitial& initial +) { + auto shared_a = initial.shared_a; + auto shared_b = initial.shared_b; + auto expert_a = initial.expert_a; + auto expert_b = initial.expert_b; + assert(qwen36_set_adapter_lora_tensor( + ctx, adapter_id, 0, "shared_gate_proj", 0, &shared_a) == 0); + assert(qwen36_set_adapter_lora_tensor( + ctx, adapter_id, 0, "shared_gate_proj", 1, &shared_b) == 0); + assert(qwen36_set_adapter_lora_tensor( + ctx, adapter_id, 0, "experts_gate_up_proj", 0, &expert_a) == 0); + assert(qwen36_set_adapter_lora_tensor( + ctx, adapter_id, 0, "experts_gate_up_proj", 1, &expert_b) == 0); +} + +static at::Tensor dp_adapter_b_delta( + void* ctx, int64_t adapter_id, const DpAdapterInitial& initial +) { + auto* shared_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_id, 0, "shared_gate_proj", 1)); + auto* expert_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + ctx, adapter_id, 0, "experts_gate_up_proj", 1)); + assert(shared_b && expert_b); + return at::cat({ + (*shared_b - initial.shared_b).to(at::kFloat).reshape({-1}), + (*expert_b - initial.expert_b).to(at::kFloat).reshape({-1}) + }); +} + +static at::Tensor dp_adapter_b_optimizer( + void* ctx, int64_t adapter_id, bool is_v +) { + auto* shared_state = reinterpret_cast( + qwen36_get_adapter_optimizer_tensor( + ctx, adapter_id, 0, "shared_gate_proj", 1, is_v ? 1 : 0)); + auto* expert_state = reinterpret_cast( + qwen36_get_adapter_optimizer_tensor( + ctx, adapter_id, 0, "experts_gate_up_proj", 1, is_v ? 1 : 0)); + assert(shared_state && expert_state); + assert(shared_state->scalar_type() == at::kFloat); + assert(expert_state->scalar_type() == at::kFloat); + return at::cat({ + shared_state->reshape({-1}), expert_state->reshape({-1})}); +} + +static at::Tensor dp_adapter_initial_b(const DpAdapterInitial& initial) { + return at::cat({ + initial.shared_b.to(at::kFloat).reshape({-1}), + initial.expert_b.to(at::kFloat).reshape({-1}) + }); +} + +static int run_dynamic_dp_smoke( + std::vector& weight_ptrs, + at::Tensor& embed, at::Tensor& final_norm, at::Tensor& lm_head, + LayerConfig& config, int process_rank, int64_t lora_rank, + int64_t vocab, int64_t intermediate, int64_t experts +) { + constexpr double learning_rate = 1e-3; + constexpr double adam_eps = 1e-8; + constexpr double beta1 = 0.9; + constexpr double beta2 = 0.999; + const int64_t target_layer = 0; + const char* targets = "shared_gate_proj,experts_gate_up_proj"; + auto create_context = [&]() { + return qwen36_create_training_context( + weight_ptrs.data(), static_cast(weight_ptrs.size()), + &embed, &final_norm, &lm_head, &config, 1, + static_cast(at::kBFloat16), + 1.0, learning_rate, beta1, beta2, adam_eps, + vocab, 1e-5, lora_rank, &target_layer, 1, targets); + }; + auto add_adapters = [&](void* ctx) { + std::vector ids; + ids.push_back(qwen36_add_lora( + ctx, lora_rank, 1.0, &target_layer, 1, targets)); + ids.push_back(qwen36_add_lora( + ctx, lora_rank, 1.0, &target_layer, 1, targets)); + assert(ids[0] > 0 && ids[1] > ids[0]); + return ids; + }; + + auto input_ids = process_rank == 0 + ? at::tensor({1, 2, 3, 4, 4, 3, 2, 1}, + at::TensorOptions().device(at::kCUDA).dtype(at::kLong)) + : at::tensor({2, 4, 1, 3, 3, 1, 4, 2}, + at::TensorOptions().device(at::kCUDA).dtype(at::kLong)); + input_ids = input_ids.reshape({2, 4}); + auto target_mask = process_rank == 0 + ? at::tensor({1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0}, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)) + : at::tensor({1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0}, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); + target_mask = target_mask.reshape({2, 4}); + auto attention_mask = at::ones({2, 4}, + at::TensorOptions().device(at::kCUDA).dtype(at::kBool)); + + // First obtain each rank's independently-normalized local first moment. + // At the first step m=(1-beta1)*g, so token weighting remains linear and + // can be checked without inferring gradients from BF16 parameter deltas. + void* local_ctx = create_context(); + assert(local_ctx); + auto local_ids = add_adapters(local_ctx); + std::vector initial; + initial.push_back(capture_dp_adapter(local_ctx, local_ids[0])); + initial.push_back(capture_dp_adapter(local_ctx, local_ids[1])); + const double local_loss = qwen36_train_multi_lora( + local_ctx, &input_ids, &target_mask, &attention_mask, 2, lora_rank); + assert(local_loss == local_loss && local_loss > 0.0); + c10::cuda::device_synchronize(); + auto local_gradients = at::stack({ + dp_adapter_b_optimizer(local_ctx, local_ids[0], false), + dp_adapter_b_optimizer(local_ctx, local_ids[1], false) + }).to(at::kCPU); + qwen36_free_training_context(local_ctx); + + const std::string sync_dir = dp_smoke_sync_dir(); + const std::string local_path = sync_dir + "/dp-local-" + + std::to_string(process_rank) + ".bin"; + const std::string peer_path = sync_dir + "/dp-local-" + + std::to_string(1 - process_rank) + ".bin"; + std::remove(local_path.c_str()); + write_tensor_file(local_path, local_gradients); + auto peer_gradients = read_tensor_file(peer_path); + assert(peer_gradients.sizes() == local_gradients.sizes()); + + void* distributed_ctx = create_context(); + assert(distributed_ctx); + assert(qwen36_init_nccl(distributed_ctx) == 0); + auto distributed_ids = add_adapters(distributed_ctx); + restore_dp_adapter(distributed_ctx, distributed_ids[0], initial[0]); + restore_dp_adapter(distributed_ctx, distributed_ids[1], initial[1]); + const double distributed_loss = qwen36_train_multi_lora( + distributed_ctx, &input_ids, &target_mask, &attention_mask, + 2, lora_rank); + assert(distributed_loss == distributed_loss && distributed_loss > 0.0); + c10::cuda::device_synchronize(); + auto distributed_deltas = at::stack({ + dp_adapter_b_delta( + distributed_ctx, distributed_ids[0], initial[0]), + dp_adapter_b_delta( + distributed_ctx, distributed_ids[1], initial[1]) + }).to(at::kCPU); + auto distributed_gradients = at::stack({ + dp_adapter_b_optimizer(distributed_ctx, distributed_ids[0], false), + dp_adapter_b_optimizer(distributed_ctx, distributed_ids[1], false) + }).to(at::kCPU); + auto distributed_v = at::stack({ + dp_adapter_b_optimizer(distributed_ctx, distributed_ids[0], true), + dp_adapter_b_optimizer(distributed_ctx, distributed_ids[1], true) + }).to(at::kCPU); + auto initial_b = at::stack({ + dp_adapter_initial_b(initial[0]), dp_adapter_initial_b(initial[1]) + }).to(at::kCPU); + + const std::string distributed_path = sync_dir + "/dp-global-" + + std::to_string(process_rank) + ".bin"; + const std::string peer_distributed_path = sync_dir + "/dp-global-" + + std::to_string(1 - process_rank) + ".bin"; + std::remove(distributed_path.c_str()); + write_tensor_file(distributed_path, distributed_gradients); + auto peer_distributed_gradients = read_tensor_file(peer_distributed_path); + assert(at::allclose( + distributed_gradients, peer_distributed_gradients, 0.0, 0.0)); + + auto gradient0 = process_rank == 0 ? local_gradients : peer_gradients; + auto gradient1 = process_rank == 0 ? peer_gradients : local_gradients; + auto counts0 = at::tensor( + {1.0, 3.0}, at::TensorOptions().dtype(at::kFloat)).reshape({2, 1}); + auto counts1 = at::tensor( + {3.0, 1.0}, at::TensorOptions().dtype(at::kFloat)).reshape({2, 1}); + auto expected_gradient = + (gradient0 * counts0 + gradient1 * counts1) / (counts0 + counts1); + auto aggregate_count_gradient = (gradient0 + gradient1) * 0.5; + + auto relative_error = [&](const at::Tensor& actual, + const at::Tensor& expected) { + return (actual - expected).abs().sum().item() / + std::max(expected.abs().sum().item(), 1e-12); + }; + const int64_t shared_numel = intermediate * lora_rank; + const int64_t expert_numel = experts * 2 * intermediate * lora_rank; + assert(distributed_gradients.size(1) == shared_numel + expert_numel); + const double all_error = relative_error( + distributed_gradients, expected_gradient); + const double grouped_error = relative_error( + distributed_gradients.narrow(1, shared_numel, expert_numel), + expected_gradient.narrow(1, shared_numel, expert_numel)); + const double old_formula_gap = + (expected_gradient - aggregate_count_gradient).abs().sum().item(); + auto expected_v = distributed_gradients.square() * + ((1.0 - beta2) / ((1.0 - beta1) * (1.0 - beta1))); + const double v_error = relative_error(distributed_v, expected_v); + auto expected_parameter = ( + initial_b - learning_rate * + (distributed_gradients / (1.0 - beta1)) / + ((distributed_v / (1.0 - beta2)).sqrt() + adam_eps)) + .to(at::kBFloat16).to(at::kFloat); + auto expected_parameter_delta = expected_parameter - initial_b; + const double adam_delta_max_error = + (distributed_deltas - expected_parameter_delta) + .abs().max().item(); + const double grouped_update = distributed_deltas + .narrow(1, shared_numel, expert_numel).abs().sum().item(); + std::printf( + "native_qwen36_dynamic_dp_weighting rank=%d loss=%0.8f " + "relative_error=%0.6e grouped_error=%0.6e old_formula_gap=%0.6e " + "v_error=%0.6e adam_delta_max_error=%0.6e grouped_update=%0.6e\n", + process_rank, distributed_loss, all_error, grouped_error, + old_formula_gap, v_error, adam_delta_max_error, grouped_update); + assert(old_formula_gap > 1e-8); + assert(grouped_update > 0.0); + assert(all_error < 2e-5); + assert(grouped_error < 2e-5); + assert(v_error < 2e-5); + assert(adam_delta_max_error < 2e-5); + qwen36_free_training_context(distributed_ctx); + return 0; +} + int main() { assert(qwen36_kernel_abi_version() == 11); const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); @@ -128,6 +437,14 @@ int main() { config.nccl_comm = nullptr; config.nccl_stream = nullptr; + const bool data_parallel = std::getenv("RUSTRAIN_DATA_PARALLEL") && + std::strcmp(std::getenv("RUSTRAIN_DATA_PARALLEL"), "0") != 0; + if (world == 2 && tp_size == 1 && data_parallel) { + return run_dynamic_dp_smoke( + weight_ptrs, embed, final_norm, lm_head, config, + process_rank, rank, vocab, intermediate, experts); + } + const int64_t target_layer = 0; // Native C++ must reject a mixed TP/DP/EP topology even when Rust-side // validation is bypassed. The smoke process is single-rank, so simulate @@ -221,9 +538,27 @@ int main() { assert(loss > 0.0); const double update_norm = (*expert_a - expert_a_before).abs().sum().item(); const double b_update_norm = (*expert_b - expert_b_before).abs().sum().item(); + std::vector fixed_m(2 * count); + std::vector fixed_v(2 * count); + assert(qwen36_export_optimizer_state( + ctx, fixed_m.data(), fixed_v.data(), 2 * count) == 2 * count); + auto* expert_b_m = reinterpret_cast(fixed_m[2 * 7 + 1]); + auto* expert_b_v = reinterpret_cast(fixed_v[2 * 7 + 1]); + assert(expert_b_m && expert_b_v); + auto expected_expert_b = ( + expert_b_before.to(at::kFloat) - 1e-3 * + (expert_b_m->to(at::kFloat) / (1.0 - 0.9)) / + ((expert_b_v->to(at::kFloat) / (1.0 - 0.999)).sqrt() + 1e-8)) + .to(at::kBFloat16); + const double fixed_adam_max_error = + (*expert_b - expected_expert_b).abs().max().item(); std::printf("native_qwen36_moe_lora_smoke expert_a_update=%0.8e expert_b_update=%0.8e\n", update_norm, b_update_norm); assert(update_norm > 0.0 || b_update_norm > 0.0); + std::printf( + "native_qwen36_fixed_adam_oracle max_error=%0.8e\n", + fixed_adam_max_error); + assert(fixed_adam_max_error == 0.0); // Dynamic multi-adapter batches must apply shared-expert MLP LoRA per // sample and preserve the parameter updates after chunk registry restore. @@ -245,7 +580,8 @@ int main() { qwen36_get_adapter_lora_tensor( ctx, adapter_one, 0, "shared_gate_proj", 1)); assert(dynamic_b && dynamic_b->sizes() == at::IntArrayRef({intermediate, local_lora_rank})); - auto dynamic_b_value = at::ones(dynamic_b->sizes(), dynamic_b->options()); + auto dynamic_b_value = at::full( + dynamic_b->sizes(), 0.01, dynamic_b->options()); assert(qwen36_set_adapter_lora_tensor( ctx, adapter_one, 0, "shared_gate_proj", 1, &dynamic_b_value) == 0); auto dynamic_b_before = dynamic_b->clone(); @@ -253,7 +589,8 @@ int main() { qwen36_get_adapter_lora_tensor( ctx, adapter_two, 0, "shared_gate_proj", 1)); assert(dynamic_b_two && dynamic_b_two->sizes() == dynamic_b->sizes()); - auto dynamic_b_two_value = at::full(dynamic_b_two->sizes(), -1.0, dynamic_b_two->options()); + auto dynamic_b_two_value = at::full( + dynamic_b_two->sizes(), -0.01, dynamic_b_two->options()); assert(qwen36_set_adapter_lora_tensor( ctx, adapter_two, 0, "shared_gate_proj", 1, &dynamic_b_two_value) == 0); auto dynamic_b_two_before = dynamic_b_two->clone(); @@ -262,8 +599,8 @@ int main() { ctx, adapter_one, 0, "experts_gate_up_proj", 1)); assert(dynamic_expert_b && dynamic_expert_b->sizes() == at::IntArrayRef({experts, 2 * intermediate, local_lora_rank})); - auto dynamic_expert_b_value = at::ones( - dynamic_expert_b->sizes(), dynamic_expert_b->options()); + auto dynamic_expert_b_value = at::full( + dynamic_expert_b->sizes(), 0.01, dynamic_expert_b->options()); assert(qwen36_set_adapter_lora_tensor( ctx, adapter_one, 0, "experts_gate_up_proj", 1, &dynamic_expert_b_value) == 0); @@ -341,18 +678,16 @@ int main() { ctx, adapter_one, 0, "shared_gate_proj", 1) != nullptr); assert(qwen36_get_adapter_lora_tensor( ctx, adapter_two, 0, "shared_gate_proj", 1) != nullptr); - // Different tenant rows can have different token counts. The current - // native DP path only has one aggregate count, so reject batch=n_total - // instead of silently applying the wrong cross-rank weighting. - setenv("RUSTRAIN_DATA_PARALLEL", "1", 1); - qwen36_set_nccl_comm(ctx, nullptr, nullptr, 0, 2); + // A tenant may be empty on one DP rank, but it must have at least one + // target token globally. In this single-rank check the second tenant is + // globally empty, so reject the step without advancing either clock. + auto invalid_multi_target_mask = multi_target_mask.clone(); + invalid_multi_target_mask.select(0, 1).zero_(); assert(qwen36_train_multi_lora( - ctx, &multi_input_ids, &multi_target_mask, &multi_attention_mask, + ctx, &multi_input_ids, &invalid_multi_target_mask, &multi_attention_mask, 2, rank) < 0.0); assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 2); assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 1); - unsetenv("RUSTRAIN_DATA_PARALLEL"); - qwen36_set_nccl_comm(ctx, nullptr, nullptr, 0, 1); qwen36_free_training_context(ctx); // Dense Qwen3.5 variants use the same per-sample activation path for @@ -386,7 +721,8 @@ int main() { auto* dense_b = reinterpret_cast( qwen36_get_adapter_lora_tensor(ctx, dense_one, 0, "gate_proj", 1)); assert(dense_b && dense_b->sizes() == at::IntArrayRef({intermediate, local_lora_rank})); - auto dense_b_value = at::ones(dense_b->sizes(), dense_b->options()); + auto dense_b_value = at::full( + dense_b->sizes(), 0.01, dense_b->options()); assert(qwen36_set_adapter_lora_tensor( ctx, dense_one, 0, "gate_proj", 1, &dense_b_value) == 0); auto dense_b_before = dense_b->clone(); @@ -536,8 +872,8 @@ int main() { ctx, linear_adapter_one, 0, "in_proj_qkv", 1)); assert(dynamic_linear_b && dynamic_linear_b->sizes() == at::IntArrayRef({linear_qkv, local_lora_rank})); - auto dynamic_linear_b_value = at::ones( - dynamic_linear_b->sizes(), dynamic_linear_b->options()); + auto dynamic_linear_b_value = at::full( + dynamic_linear_b->sizes(), 0.01, dynamic_linear_b->options()); assert(qwen36_set_adapter_lora_tensor( ctx, linear_adapter_one, 0, "in_proj_qkv", 1, &dynamic_linear_b_value) == 0); From d6f897cfeeb52b7a23ce6be23c99937c605f2bb2 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 07:05:57 +0800 Subject: [PATCH 019/156] docs: record dp ep and optimizer verification --- docs/plans/qwen-lora-megatron-progress.md | 11 +++++++---- docs/plans/qwen-lora-megatron-spec.md | 2 +- docs/qwen35-qwen36-megatron-audit.md | 12 ++++++------ 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/docs/plans/qwen-lora-megatron-progress.md b/docs/plans/qwen-lora-megatron-progress.md index 8b6af249..36ab22fb 100644 --- a/docs/plans/qwen-lora-megatron-progress.md +++ b/docs/plans/qwen-lora-megatron-progress.md @@ -12,7 +12,7 @@ timestamp: 2026-07-17T00:00:00Z # Current State -Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP smoke, dense replicated-DP smoke, TP-only latent-rank-sharded LoRA smoke, dynamic batch logical-step update, ABI11 FP32 gradient storage/aggregation, per-tenant optimizer-step restore, selected-tenant isolation, same-topology rank-aware checkpointing, and 5D topology mapping. +Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP parity against a full-expert reference, dense replicated-DP smoke with per-tenant token weighting, TP-only latent-rank-sharded LoRA smoke, dynamic batch logical-step update, ABI11 FP32 gradient storage/aggregation, standard Adam bias correction, per-tenant optimizer-step restore, selected-tenant isolation, same-topology rank-aware checkpointing, and 5D topology mapping. Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axis TP+DP/EP, PP/CP, variable-split token A2A, DeepEP/TE prebuilt integration, cross-topology checkpoint resharding, and matched Megatron throughput. @@ -24,12 +24,15 @@ Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axi - ABI10 working tree: TP-only mode shards each fixed and dynamic adapter's latent rank, all-reduces only activation-level LoRA deltas on a split TP communicator, and keeps replicated MoE layers off the EP communicator. - `5aa4ea4`: ABI11 stores fixed and dynamic LoRA leaf gradients in FP32 accumulators between micro-batches; fused Adam consumes those accumulators directly, fixed windows use token-weighted numerator/denominator reduction, and explicit abort/failure cleanup clears the pending window. This is FP32 storage/aggregation; the autograd leaf backward remains BF16-typed. - `ed25daa`: TP rank-aware checkpoint format v3 writes `rank-xxxxx/` manifests and adapter/optimizer shards with topology, global/local shape, latent partition axis, offset, and replica identity metadata. v1/v2 single-rank loading remains compatible; v3 requires the same topology and rank. -- `ff382a0`: native ABI rejects unsupported TP+DP/EP mixtures before training. Dynamic per-tenant batch input under DP is explicitly rejected because the current reduction contract has one aggregate token count; shared batch-1/equal-mask dynamic DP remains the supported contract. +- `ff382a0`: native ABI rejects unsupported TP+DP/EP mixtures before training. Dynamic per-tenant DP batches now use an all-reduced per-adapter token-count vector and weighted FP32 LoRA gradient reduction, including grouped expert LoRA. - `76eace6`: direct native context creation also rejects PP/CP sizes greater than one, so unsupported pipeline/context parallelism cannot silently fall back to replicated single-stage execution. +- `b7050ea`: fixed standard Adam bias correction, separated pure-DP gradient NCCL from MoE activation collectives, made dynamic-MoE DP token weighting numerical, scoped n_max rendezvous by invocation with atomic publication, and added full-expert EP parity plus Adam oracles. - H20 `123.57.26.97:28004`: ABI8 native smoke passed grouped/fallback parity, GDN, dense/MoE LoRA, dynamic adapters, and step setter validation. - H20 `123.57.26.97:28004`: ABI9 native smoke passed selected-tenant training with a positive selected update, exactly zero unselected update, independent clocks (`2` vs `1`), and registry preservation after an unknown ID. - H20 `123.57.26.97:28004`: ABI10 two-rank TP native smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, and selected-tenant isolation. Rank-local LoRA tensors used distinct rank `4` slices for global rank `8`; losses matched and both shards had positive updates. -- H20 `123.57.26.97:28004`: ABI11 single-rank smoke passed FP32 accumulator dtype, two-micro accumulation, NaN/explicit abort cleanup, successful step commit, TP/DP/EP and PP/CP topology rejection, and dynamic-DP per-tenant batch rejection. ABI11 two-rank TP smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, selected-tenant isolation, and the new guards (`rank_statuses=0,0`). +- H20 `123.57.26.97:28004`: ABI11 single-rank smoke passed FP32 accumulator dtype, two-micro accumulation, NaN/explicit abort cleanup, successful step commit, standard Adam parameter oracle, TP/DP/EP and PP/CP topology rejection, and dynamic multi-LoRA. ABI11 two-rank TP smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, selected-tenant isolation, invocation-scoped n_max rendezvous, and the new guards (`rank_statuses=0,0`). +- H20 `123.57.26.97:28004`: ABI11 DP2 smoke passed with per-tenant masks `[1,3]` and `[3,1]`: weighted m relative error `2.43e-8`, grouped-expert error `2.27e-8`, v error `7.33e-8`, BF16 Adam delta error `0`, and nonzero gap `7.96e-3` versus the old equal-count formula (`rank_statuses=0,0`). +- H20 `123.57.26.97:28004`: ABI11 EP2 full-expert parity smoke passed with distinct rank-local base and LoRA expert slices; loss, A/B updates, m/v, and standard Adam first-step oracle all matched the rank-local full-expert reference (`rank_statuses=0,0`). - Target runtime probe: PyTorch 2.5.1+cu121, ABI0; Transformer Engine, flash-attn, DeepEP, Triton, and DeepSpeed are not importable. # Decisions During Execution @@ -41,6 +44,6 @@ Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axi # Verification -Passed: `cargo check -p rustrain-qwen3-6 -p rustrain-server -p rustrain-ipc`, `cargo test -p rustrain-server --lib` (6), Qwen unit tests (3), Qwen integration (6), remote ABI8 smoke, remote ABI9 selected-tenant native smoke, remote ABI10 single-rank native smoke, remote ABI10 two-rank TP native smoke, and remote ABI11 single/two-rank native smoke. +Passed: `cargo check -p rustrain-qwen3-6 -p rustrain-server -p rustrain-ipc` (with the repository PyTorch 2.12.1 host venv), `cargo test -p rustrain-server --lib` (6), Qwen unit tests (3), Qwen integration (6), remote ABI8 smoke, remote ABI9 selected-tenant native smoke, remote ABI10 single-rank native smoke, remote ABI10 two-rank TP native smoke, and remote ABI11 single/TP2/DP2/EP2 native smoke with numerical Adam and parity oracles. Not run: Megatron-style base-model TP, variable-split EP token dispatch, multi-axis TP+DP/EP, PP/CP, cross-topology resharding, a numerical FP32-accumulation oracle against a concatenated batch, and matched Megatron performance benchmark. diff --git a/docs/plans/qwen-lora-megatron-spec.md b/docs/plans/qwen-lora-megatron-spec.md index b961932c..96588f95 100644 --- a/docs/plans/qwen-lora-megatron-spec.md +++ b/docs/plans/qwen-lora-megatron-spec.md @@ -8,7 +8,7 @@ timestamp: 2026-07-17T00:00:00Z # Problem -The native Qwen3.5/3.6 path has correct and tested single-rank, EP, and dense-DP slices, but it is not yet a Megatron-LM-level LoRA runtime. In particular, Qwen native TP/PP/CP are not implemented, dynamic tenants share optimizer state semantics, and the MoE path lacks fused asynchronous token dispatch. +The native Qwen3.5/3.6 path has correct and tested single-rank, EP, and dense-DP slices, but it is not yet a Megatron-LM-level LoRA runtime. In particular, Qwen native TP/PP/CP are not implemented, and the MoE path lacks fused asynchronous token dispatch. # Target Outcome diff --git a/docs/qwen35-qwen36-megatron-audit.md b/docs/qwen35-qwen36-megatron-audit.md index 9515b925..92ca95ee 100644 --- a/docs/qwen35-qwen36-megatron-audit.md +++ b/docs/qwen35-qwen36-megatron-audit.md @@ -5,7 +5,7 @@ ## 结论 - 模型语义：Qwen3.5 dense、Qwen3.6 dense/MoE 的 native forward/backward 路径已经覆盖 hybrid full attention、GDN、MoE、MTP 和 LoRA 目标模块；已有配置解析、集成测试及 H20 native smoke 证据。 -- 已实现并可验证的分布式子集：MoE expert parallel，以及 replicated LoRA 的 data parallel；梯度累积和 dynamic multi-LoRA 已有 logical-step 边界。 +- 已实现并可验证的分布式子集：MoE expert parallel，以及 replicated LoRA 的 data parallel；梯度累积和 dynamic multi-LoRA 已有 logical-step 边界。DP 动态租户按 adapter token count 加权，纯 DP 不把 world communicator 传入 MoE activation reduction。 - 性能：MoE grouped dispatch 相对逐 expert matmul 的已有 microbenchmark 为约 3.70x（E=32, N=4096, H=2048, I=768，结果误差为 0）；这不是端到端训练吞吐或 Megatron 对比。 - 尚未实现：Qwen native 路径的 tensor parallel、pipeline parallel、context parallel，以及 TP/PP/CP 与 EP/DP 的组合。当前训练上下文仍由单个进程持有完整 dense 权重和完整层栈。 - 因此当前实现不能宣称“Megatron-LM 级别”。它是一个计算集中在 C++ 的 LoRA/EP/DP 子集，离 Megatron 的完整并行和通信重叠仍有实质差距。 @@ -19,8 +19,8 @@ | Qwen3.6 MoE | 已实现 | grouped dispatch、EP smoke；完整模型仍需目标 GPU/权重运行 | | MTP | 已实现 | C++ hidden gradient 检查和集成测试；可通过环境变量关闭 | | fixed LoRA | 已实现 | attention/GDN/MLP/shared/routed expert 目标模块 | -| dynamic multi-LoRA | 已实现子集 | 请求按 adapter 分组，单个 logical step 统一 backward/Adam；adapter 仍共享 context optimizer step | -| microbatch accumulation | 已实现子集 | non-final microbatch 只 backward，final microbatch 才 optimizer；梯度仍累加在 BF16 leaf 上 | +| dynamic multi-LoRA | 已实现子集 | 请求按 adapter 分组，单个 logical step 统一 backward/Adam；每个 adapter 独立 optimizer clock 和 m/v，DP 按 token count 加权 | +| microbatch accumulation | 已实现子集 | non-final microbatch 只 backward，final microbatch 才 optimizer；FP32 accumulator 存储/聚合，autograd leaf backward 仍为 BF16 | | replicated data parallel | 已实现 | logical-step 边界同步 replicated LoRA；EP expert 参数不走该 reduction | | expert parallel | 已实现子集 | 路由输出 all-reduce 和本地 expert 权重；没有 DeepEP 式 fused A2A/dispatch overlap | | tensor parallel | 未实现于 Qwen native | 不切分 attention/MLP/LM-head 权重，也没有 Qwen TP communicator | @@ -34,11 +34,11 @@ Megatron 的 TP 会按 head、hidden/intermediate 和 vocab 维度切分权重，并在线性层边界执行必要的 reduce-scatter/all-reduce；PP 会把层分到不同 stage 并使用 1F1B 等调度；CP 会在序列和 attention state 上做跨 rank 通信。当前 Qwen native `TrainingContext` 仍加载完整模型并在一个 C++ forward 中执行全部层，因此仅增加 `tensor_model_parallel_size` 等配置不能得到正确的 TP/PP/CP。 -当前 DP/EP 也不是完整 Megatron 语义：DP 只同步 replicated LoRA 梯度，expert 参数留在 EP rank；EP 使用已有 all-reduce，但没有 fused token dispatch/combine、异步 A2A 和通信计算重叠。 +当前 DP/EP 也不是完整 Megatron 语义：DP 同步 replicated LoRA 梯度并按租户 token count 归一化，expert 参数留在 EP rank；EP 使用 routed-output all-reduce，但没有 fused token dispatch/combine、异步 A2A 和通信计算重叠。 ### 优化器与恢复 -固定 LoRA 的 Adam 状态可导出/导入，且 native context 的 logical step 需要与 checkpoint step 对齐。dynamic adapter 的请求频率不同，但目前仍共享 context-level step；尚无每租户独立 optimizer step、FP32 gradient accumulator 或 accumulation window abort/zero API。这些差距会影响长时间多租户训练的数值一致性和故障恢复。 +固定 LoRA 的 Adam 状态可导出/导入，且 native context 的 logical step 需要与 checkpoint step 对齐。dynamic adapter 的请求频率不同，每租户拥有独立 optimizer step、m/v 与 FP32 accumulator；仍缺少跨 optimizer group 的事务性回滚。 ### 性能工程 @@ -46,7 +46,7 @@ Megatron 的 TP 会按 head、hidden/intermediate 和 vocab 维度切分权重 ## 验证边界 -已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试、H20 ABI0 native smoke，以及已有 ABI1 环境中的单卡、EP 和 DP smoke。没有完成 Qwen3.5/3.6 完整大模型的长时间训练、跨节点通信、TP/PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖的模型/LoRA/EP/DP 子集，而不是所有并行配置。 +已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试，以及 H20 ABI11 的单卡、TP2、EP2 和 DP2 native smoke。DP2 的 weighted m/grouped/v/Adam BF16 delta oracle 分别达到 `2.43e-8`、`2.27e-8`、`7.33e-8` 和 `0`；EP2 与 full-expert reference 的 loss、LoRA、Adam state 和标准 Adam 首步 oracle 在两 rank 均为零差异。没有完成 Qwen3.5/3.6 完整大模型的长时间训练、跨节点通信、TP/PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖的模型/LoRA/EP/DP 子集，而不是所有并行配置。 ## 继续达到 Megatron 级别所需的最小工作包 From bfa3d8ec9fbe2317bf08aea8ee626101b12784f1 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 07:23:25 +0800 Subject: [PATCH 020/156] feat: prototype variable split ep token a2a --- .../kernels/qwen3_6_kernels.cpp | 392 +++++++++++++++++- .../tests/native_ep_smoke.cpp | 26 +- 2 files changed, 405 insertions(+), 13 deletions(-) diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 715a1699..edacdd9d 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -845,6 +845,239 @@ static at::Tensor linear_attention_batched(TrainingContext* ctx, const at::Tenso // MoE // ────────────────────────────────────────────────────────────────────── +static ncclDataType_t qwen36_nccl_dtype(at::ScalarType type) { + switch (type) { + case at::kFloat: return ncclFloat32; + case at::kHalf: return ncclFloat16; + case at::kBFloat16: return ncclBfloat16; + case at::kInt: return ncclInt; + case at::kLong: return ncclInt64; + default: + TORCH_CHECK(false, "unsupported NCCL dtype: ", type); + } +} + +static std::vector qwen36_a2a_counts( + const at::Tensor& local_counts, ncclComm_t comm, cudaStream_t stream +) { + const int world = static_cast(local_counts.numel()); + auto all_counts = at::empty({world, world}, local_counts.options()); + auto err = ncclAllGather( + local_counts.data_ptr(), all_counts.data_ptr(), world, + ncclInt, comm, stream); + TORCH_CHECK(err == ncclSuccess, "EP A2A count all-gather failed: ", + ncclGetErrorString(err)); + auto host = all_counts.to(at::TensorOptions().device(at::kCPU)); + std::vector recv_counts(world); + auto ptr = host.data_ptr(); + int rank = 0; + err = ncclCommUserRank(comm, &rank); + TORCH_CHECK(err == ncclSuccess, "ncclCommUserRank failed: ", + ncclGetErrorString(err)); + for (int src = 0; src < world; ++src) { + recv_counts[src] = ptr[src * world + rank]; + TORCH_CHECK(recv_counts[src] >= 0, "negative EP A2A receive count"); + } + return recv_counts; +} + +// Variable-split token dispatch. Metadata is deliberately non-differentiable; +// only the hidden activation participates in the custom backward exchange. +struct Qwen36A2ADispatchFunction : public torch::autograd::Function { + static std::vector forward( + torch::autograd::AutogradContext* ctx, + at::Tensor input, at::Tensor expert_indices, at::Tensor token_indices, + int64_t expert_count, int64_t comm_ptr + ) { + auto comm = reinterpret_cast(comm_ptr); + TORCH_CHECK(comm, "EP A2A requires an NCCL communicator"); + int world = 1, rank = 0; + TORCH_CHECK(ncclCommCount(comm, &world) == ncclSuccess, + "ncclCommCount failed"); + TORCH_CHECK(ncclCommUserRank(comm, &rank) == ncclSuccess, + "ncclCommUserRank failed"); + TORCH_CHECK(expert_count > 0 && expert_indices.scalar_type() == at::kLong, + "invalid EP A2A expert metadata"); + auto stream = c10::cuda::getCurrentCUDAStream(input.device().index()).stream(); + const int64_t hidden = input.size(1); + std::vector indices(world); + std::vector send_counts(world, 0); + for (int dst = 0; dst < world; ++dst) { + auto mask = (expert_indices >= dst * expert_count) & + (expert_indices < (dst + 1) * expert_count); + indices[dst] = at::nonzero(mask).reshape({-1}); + send_counts[dst] = static_cast(indices[dst].numel()); + } + auto count_opts = at::TensorOptions().device(input.device()).dtype(at::kInt); + auto local_counts = at::empty({world}, count_opts); + auto host_counts = at::empty({world}, at::TensorOptions().device(at::kCPU).dtype(at::kInt)); + std::memcpy(host_counts.data_ptr(), send_counts.data(), sizeof(int32_t) * world); + local_counts.copy_(host_counts); + auto recv_counts = qwen36_a2a_counts(local_counts, comm, stream); + std::vector send_offsets(world + 1, 0), recv_offsets(world + 1, 0); + for (int i = 0; i < world; ++i) { + send_offsets[i + 1] = send_offsets[i] + send_counts[i]; + recv_offsets[i + 1] = recv_offsets[i] + recv_counts[i]; + } + auto send_index = at::cat(indices, 0); + auto send_hidden = input.index_select(0, send_index).contiguous(); + auto send_token = token_indices.index_select(0, send_index).contiguous(); + auto send_expert = expert_indices.index_select(0, send_index).contiguous(); + auto recv_hidden = at::empty({recv_offsets.back(), hidden}, input.options()); + auto recv_token = at::empty({recv_offsets.back()}, token_indices.options()); + auto recv_expert = at::empty({recv_offsets.back()}, expert_indices.options()); + auto send_ptr = [&](at::Tensor& t, int64_t row, int64_t width) -> const void* { + return static_cast(t.data_ptr()) + + row * width * t.element_size(); + }; + auto recv_ptr = [&](at::Tensor& t, int64_t row, int64_t width) -> void* { + return static_cast(t.data_ptr()) + + row * width * t.element_size(); + }; + TORCH_CHECK(ncclGroupStart() == ncclSuccess, "ncclGroupStart failed"); + for (int peer = 0; peer < world; ++peer) { + const int64_t rows = send_counts[peer]; + if (rows) { + auto err = ncclSend(send_ptr(send_hidden, send_offsets[peer], hidden), + rows * hidden, qwen36_nccl_dtype(input.scalar_type()), peer, comm, stream); + TORCH_CHECK(err == ncclSuccess, "A2A hidden send failed: ", ncclGetErrorString(err)); + err = ncclSend(send_ptr(send_token, send_offsets[peer], 1), rows, + ncclInt64, peer, comm, stream); + TORCH_CHECK(err == ncclSuccess, "A2A token send failed: ", ncclGetErrorString(err)); + err = ncclSend(send_ptr(send_expert, send_offsets[peer], 1), rows, + ncclInt64, peer, comm, stream); + TORCH_CHECK(err == ncclSuccess, "A2A expert send failed: ", ncclGetErrorString(err)); + } + } + for (int peer = 0; peer < world; ++peer) { + const int64_t rows = recv_counts[peer]; + if (rows) { + auto err = ncclRecv(recv_ptr(recv_hidden, recv_offsets[peer], hidden), + rows * hidden, qwen36_nccl_dtype(input.scalar_type()), peer, comm, stream); + TORCH_CHECK(err == ncclSuccess, "A2A hidden recv failed: ", ncclGetErrorString(err)); + err = ncclRecv(recv_ptr(recv_token, recv_offsets[peer], 1), rows, + ncclInt64, peer, comm, stream); + TORCH_CHECK(err == ncclSuccess, "A2A token recv failed: ", ncclGetErrorString(err)); + err = ncclRecv(recv_ptr(recv_expert, recv_offsets[peer], 1), rows, + ncclInt64, peer, comm, stream); + TORCH_CHECK(err == ncclSuccess, "A2A expert recv failed: ", ncclGetErrorString(err)); + } + } + TORCH_CHECK(ncclGroupEnd() == ncclSuccess, "A2A dispatch group failed"); + auto recv_local = recv_expert - rank * expert_count; + auto send_counts_tensor = at::empty({world}, count_opts); + auto recv_counts_tensor = at::empty({world}, count_opts); + std::memcpy(host_counts.data_ptr(), send_counts.data(), sizeof(int32_t) * world); + send_counts_tensor.copy_(host_counts); + std::memcpy(host_counts.data_ptr(), recv_counts.data(), sizeof(int32_t) * world); + recv_counts_tensor.copy_(host_counts); + ctx->save_for_backward({input, send_index, send_counts_tensor, recv_counts_tensor}); + ctx->saved_data["comm"] = comm_ptr; + ctx->saved_data["expert_count"] = expert_count; + return {recv_hidden, recv_token, recv_local, send_index, send_counts_tensor, recv_counts_tensor}; + } + + static std::vector backward( + torch::autograd::AutogradContext* ctx, std::vector grad_output + ) { + auto saved = ctx->get_saved_variables(); + auto input = saved[0]; + auto send_index = saved[1]; + auto send_counts = saved[2].to(at::TensorOptions().device(at::kCPU)); + auto recv_counts = saved[3].to(at::TensorOptions().device(at::kCPU)); + const int world = send_counts.numel(); + auto comm = reinterpret_cast(ctx->saved_data["comm"].toInt()); + auto stream = c10::cuda::getCurrentCUDAStream(input.device().index()).stream(); + std::vector sc(world), rc(world); + std::memcpy(sc.data(), send_counts.data_ptr(), sizeof(int32_t) * world); + std::memcpy(rc.data(), recv_counts.data_ptr(), sizeof(int32_t) * world); + std::vector so(world + 1, 0), ro(world + 1, 0); + for (int i = 0; i < world; ++i) { so[i + 1] = so[i] + sc[i]; ro[i + 1] = ro[i] + rc[i]; } + auto grad_input = at::zeros_like(input); + auto returned = at::empty({so.back(), input.size(1)}, input.options()); + const size_t elem_bytes = input.element_size(); + TORCH_CHECK(ncclGroupStart() == ncclSuccess, "ncclGroupStart failed"); + for (int peer = 0; peer < world; ++peer) if (sc[peer]) { + auto err = ncclRecv(static_cast(returned.data_ptr()) + so[peer] * input.size(1) * elem_bytes, + sc[peer] * input.size(1), qwen36_nccl_dtype(input.scalar_type()), peer, comm, stream); + TORCH_CHECK(err == ncclSuccess, "A2A backward recv failed: ", ncclGetErrorString(err)); + } + for (int peer = 0; peer < world; ++peer) if (rc[peer]) { + auto err = ncclSend(static_cast(grad_output[0].data_ptr()) + ro[peer] * input.size(1) * elem_bytes, + rc[peer] * input.size(1), qwen36_nccl_dtype(input.scalar_type()), peer, comm, stream); + TORCH_CHECK(err == ncclSuccess, "A2A backward send failed: ", ncclGetErrorString(err)); + } + TORCH_CHECK(ncclGroupEnd() == ncclSuccess, "A2A dispatch backward group failed"); + grad_input.index_add_(0, send_index, returned); + return {grad_input, at::Tensor(), at::Tensor(), at::Tensor(), at::Tensor()}; + } +}; + +struct Qwen36A2ACombineFunction : public torch::autograd::Function { + static at::Tensor forward(torch::autograd::AutogradContext* ctx, + at::Tensor local_output, at::Tensor send_counts, + at::Tensor recv_counts, int64_t comm_ptr) { + auto comm = reinterpret_cast(comm_ptr); + const int world = send_counts.numel(); + auto sc_cpu = send_counts.to(at::TensorOptions().device(at::kCPU)); + auto rc_cpu = recv_counts.to(at::TensorOptions().device(at::kCPU)); + std::vector sc(world), rc(world); + std::memcpy(sc.data(), sc_cpu.data_ptr(), sizeof(int32_t) * world); + std::memcpy(rc.data(), rc_cpu.data_ptr(), sizeof(int32_t) * world); + std::vector so(world + 1, 0), ro(world + 1, 0); + for (int i = 0; i < world; ++i) { so[i + 1] = so[i] + sc[i]; ro[i + 1] = ro[i] + rc[i]; } + auto stream = c10::cuda::getCurrentCUDAStream(local_output.device().index()).stream(); + auto returned = at::empty({so.back(), local_output.size(1)}, local_output.options()); + const size_t elem_bytes = local_output.element_size(); + TORCH_CHECK(ncclGroupStart() == ncclSuccess, "ncclGroupStart failed"); + for (int peer = 0; peer < world; ++peer) if (rc[peer]) { + auto err = ncclSend(static_cast(local_output.data_ptr()) + ro[peer] * local_output.size(1) * elem_bytes, + rc[peer] * local_output.size(1), qwen36_nccl_dtype(local_output.scalar_type()), peer, comm, stream); + TORCH_CHECK(err == ncclSuccess, "A2A combine send failed: ", ncclGetErrorString(err)); + } + for (int peer = 0; peer < world; ++peer) if (sc[peer]) { + auto err = ncclRecv(static_cast(returned.data_ptr()) + so[peer] * local_output.size(1) * elem_bytes, + sc[peer] * local_output.size(1), qwen36_nccl_dtype(local_output.scalar_type()), peer, comm, stream); + TORCH_CHECK(err == ncclSuccess, "A2A combine recv failed: ", ncclGetErrorString(err)); + } + TORCH_CHECK(ncclGroupEnd() == ncclSuccess, "A2A combine group failed"); + ctx->save_for_backward({send_counts, recv_counts}); + ctx->saved_data["comm"] = comm_ptr; + return returned; + } + + static std::vector backward(torch::autograd::AutogradContext* ctx, + std::vector grad_output) { + auto saved = ctx->get_saved_variables(); + auto sc_cpu = saved[0].to(at::TensorOptions().device(at::kCPU)); + auto rc_cpu = saved[1].to(at::TensorOptions().device(at::kCPU)); + const int world = sc_cpu.numel(); + std::vector sc(world), rc(world); + std::memcpy(sc.data(), sc_cpu.data_ptr(), sizeof(int32_t) * world); + std::memcpy(rc.data(), rc_cpu.data_ptr(), sizeof(int32_t) * world); + std::vector so(world + 1, 0), ro(world + 1, 0); + for (int i = 0; i < world; ++i) { so[i + 1] = so[i] + sc[i]; ro[i + 1] = ro[i] + rc[i]; } + auto comm = reinterpret_cast(ctx->saved_data["comm"].toInt()); + auto stream = c10::cuda::getCurrentCUDAStream(grad_output[0].device().index()).stream(); + auto packed = grad_output[0].contiguous(); + auto grad_local = at::empty({ro.back(), grad_output[0].size(1)}, grad_output[0].options()); + const size_t elem_bytes = grad_output[0].element_size(); + TORCH_CHECK(ncclGroupStart() == ncclSuccess, "ncclGroupStart failed"); + for (int peer = 0; peer < world; ++peer) if (sc[peer]) { + auto err = ncclSend(static_cast(packed.data_ptr()) + so[peer] * grad_output[0].size(1) * elem_bytes, + sc[peer] * grad_output[0].size(1), qwen36_nccl_dtype(grad_output[0].scalar_type()), peer, comm, stream); + TORCH_CHECK(err == ncclSuccess, "A2A combine backward send failed: ", ncclGetErrorString(err)); + } + for (int peer = 0; peer < world; ++peer) if (rc[peer]) { + auto err = ncclRecv(static_cast(grad_local.data_ptr()) + ro[peer] * grad_output[0].size(1) * elem_bytes, + rc[peer] * grad_output[0].size(1), qwen36_nccl_dtype(grad_output[0].scalar_type()), peer, comm, stream); + TORCH_CHECK(err == ncclSuccess, "A2A combine backward recv failed: ", ncclGetErrorString(err)); + } + TORCH_CHECK(ncclGroupEnd() == ncclSuccess, "A2A combine backward group failed"); + return {grad_local, at::Tensor(), at::Tensor(), at::Tensor()}; + } +}; + struct RoutedExpertLora { const at::Tensor* gate_up_a = nullptr; // [local_experts, rank, hidden] const at::Tensor* gate_up_b = nullptr; // [local_experts, 2*intermediate, rank] @@ -927,6 +1160,115 @@ static at::Tensor dynamic_expert_lora_delta( return tp_allreduce_lora_delta(ctx, delta * scaling); } +static at::Tensor moe_routed_a2a( + TrainingContext* training_ctx, ncclComm_t comm, + const at::Tensor& flat, const at::Tensor& topk_weights, + const at::Tensor& topk_indices, + const at::Tensor& experts_gate_up, const at::Tensor& experts_down, + const RoutedExpertLora& expert_lora, + int64_t top_k, int64_t intermediate, int64_t expert_count, + int64_t batch, int64_t seq, + const LoraBatchEntry* expert_gate_up_lora, + const LoraBatchEntry* expert_down_lora +) { + const int64_t hidden_dim = flat.size(1); + auto routed_output = at::zeros_like(flat); + auto token_indices = at::arange( + flat.size(0), at::TensorOptions().device(flat.device()).dtype(at::kLong)); + for (int64_t kk = 0; kk < top_k; ++kk) { + auto expert_indices = topk_indices.select(-1, kk).contiguous(); + auto dispatched = Qwen36A2ADispatchFunction::apply( + flat, expert_indices, token_indices, expert_count, + static_cast(reinterpret_cast(comm))); + auto received = dispatched[0]; + auto received_tokens = dispatched[1]; + auto received_experts = dispatched[2]; + auto send_index = dispatched[3]; + auto send_counts = dispatched[4]; + auto recv_counts = dispatched[5]; + auto local_output = at::zeros_like(received); + for (int64_t e_local = 0; e_local < expert_count; ++e_local) { + auto rows = at::nonzero(received_experts == e_local).reshape({-1}); + if (rows.numel() == 0) continue; + auto selected = received.index_select(0, rows); + auto selected_tokens = received_tokens.index_select(0, rows); + auto selected_experts = received_experts.index_select(0, rows); + auto gu = at::matmul(selected, experts_gate_up.select(0, e_local).t()); + if (expert_lora.gate_up_a && expert_lora.gate_up_b) { + auto a = expert_lora.gate_up_a->select(0, e_local); + auto b = expert_lora.gate_up_b->select(0, e_local); + auto delta = at::matmul(at::matmul(selected, a.t()), b.t()) * + expert_lora.scaling; + gu = gu + tp_allreduce_lora_delta(training_ctx, delta); + } + if (expert_gate_up_lora) { + gu = gu + dynamic_expert_lora_delta( + training_ctx, selected, selected_tokens, selected_experts, + batch, seq, expert_gate_up_lora); + } + auto activated = fused_swiglu_op( + gu.narrow(-1, 0, intermediate), + gu.narrow(-1, intermediate, intermediate), 0.0); + auto expert_out = at::matmul( + activated, experts_down.select(0, e_local).t()); + if (expert_lora.down_a && expert_lora.down_b) { + auto a = expert_lora.down_a->select(0, e_local); + auto b = expert_lora.down_b->select(0, e_local); + auto delta = at::matmul(at::matmul(activated, a.t()), b.t()) * + expert_lora.scaling; + expert_out = expert_out + + tp_allreduce_lora_delta(training_ctx, delta); + } + if (expert_down_lora) { + expert_out = expert_out + dynamic_expert_lora_delta( + training_ctx, activated, selected_tokens, selected_experts, + batch, seq, expert_down_lora); + } + local_output = local_output.index_add(0, rows, expert_out); + } + + // Preserve a zero-valued dependency on every local expert LoRA when a + // rank receives no tokens. This keeps optimizer collective order equal + // across ranks without changing the output. + if (!local_output.requires_grad()) { + // Even without routed-expert LoRA, an empty destination must keep + // the dispatch activation edge alive. Its zero gradient is sent + // back to the source in the inverse dispatch backward. + at::Tensor anchor = received.sum().to(local_output.scalar_type()); + auto include = [&](const at::Tensor* tensor) { + if (!tensor || !tensor->defined() || !tensor->requires_grad()) return; + auto value = tensor->sum().to(local_output.scalar_type()); + anchor = anchor + value; + }; + include(expert_lora.gate_up_a); + include(expert_lora.gate_up_b); + include(expert_lora.down_a); + include(expert_lora.down_b); + if (expert_gate_up_lora) { + include(&expert_gate_up_lora->a_stack); + include(&expert_gate_up_lora->b_stack); + } + if (expert_down_lora) { + include(&expert_down_lora->a_stack); + include(&expert_down_lora->b_stack); + } + local_output = local_output + anchor * 0.0; + } + + // Expert results return in the source's packed dispatch order. Apply + // routing weights on the source rank, then restore source token order; + // this keeps router/hidden gradients on the correct source graph. + auto returned = Qwen36A2ACombineFunction::apply( + local_output, send_counts, recv_counts, + static_cast(reinterpret_cast(comm))); + auto source_weights = topk_weights.select(-1, kk) + .index_select(0, send_index).unsqueeze(-1); + routed_output = routed_output.index_add( + 0, send_index, returned * source_weights); + } + return routed_output; +} + static at::Tensor moe_forward( TrainingContext* training_ctx, void* nccl_comm_v, void* nccl_stream_v, @@ -959,6 +1301,32 @@ static at::Tensor moe_forward( topk_weights = topk_weights.to(compute_type); auto routed_output = at::zeros(flat.sizes(), flat.options()); + bool use_a2a = false; + if (nccl_comm_v && env_enabled("QWEN36_EP_A2A") && + !expert_gate_up_lora && !expert_down_lora) { + auto comm = reinterpret_cast(nccl_comm_v); + int world = 1, rank = 0; + auto err = ncclCommCount(comm, &world); + TORCH_CHECK(err == ncclSuccess, "ncclCommCount failed: ", + ncclGetErrorString(err)); + err = ncclCommUserRank(comm, &rank); + TORCH_CHECK(err == ncclSuccess, "ncclCommUserRank failed: ", + ncclGetErrorString(err)); + TORCH_CHECK(world * expert_count == num_experts, + "EP A2A requires equal contiguous expert partitions: world=", world, + " local_experts=", expert_count, " global_experts=", num_experts); + TORCH_CHECK(expert_start == rank * expert_count, + "EP A2A requires rank-contiguous expert ownership: rank=", rank, + " expert_start=", expert_start, " local_experts=", expert_count); + use_a2a = world > 1; + if (use_a2a) { + routed_output = moe_routed_a2a( + training_ctx, comm, flat, topk_weights, topk_indices, + experts_gate_up, experts_down, expert_lora, + top_k, intermediate, expert_count, batch, seq, + expert_gate_up_lora, expert_down_lora); + } + } // Debug: dump MoE routing and weight stats if (getenv("QWEN36_DUMP_MOE")) { @@ -978,7 +1346,7 @@ static at::Tensor moe_forward( // Autograd note: routing indices/weights are computed in no-grad forward (detached). // Only matmul inputs/outputs participate in autograd. sort/index_select/index_add // gradients are handled by PyTorch automatically. - for (int64_t kk = 0; kk < top_k; kk++) { + if (!use_a2a) for (int64_t kk = 0; kk < top_k; kk++) { auto expert_indices = topk_indices.select(-1, kk); // [N*seq] auto expert_weights = topk_weights.select(-1, kk); // [N*seq] @@ -1147,7 +1515,7 @@ static at::Tensor moe_forward( // zero-valued dependency on routed-expert LoRA tensors so autograd still // produces defined zero gradients and every rank reaches the same NCCL // collectives. This changes only the graph, not the routed output values. - if (!routed_output.requires_grad()) { + if (!use_a2a && !routed_output.requires_grad()) { at::Tensor graph_anchor; auto include_anchor = [&](const at::Tensor* tensor) { if (!tensor || !tensor->defined() || !tensor->requires_grad()) return; @@ -1174,7 +1542,7 @@ static at::Tensor moe_forward( } // EP all-reduce via NcclAllReduceFunction — custom autograd Function. - if (nccl_comm_v) { + if (nccl_comm_v && !use_a2a) { auto nccl_comm = reinterpret_cast(nccl_comm_v); routed_output = NcclAllReduceFunction::apply( routed_output, (int64_t)nccl_comm, @@ -1968,6 +2336,18 @@ static void synchronize_lora_gradients( } } if (per_adapter_weighting) return; + // The gated A2A prototype currently runs on the repository's replicated + // EP batch layout. Every source rank therefore sends an identical token + // batch to the owning expert. Average only the sharded expert parameter + // gradients here; scaling the combine backward would also under-scale the + // activation gradient returned to each source graph. Unequal source-token + // shards will require global-token weighting before this gate can become + // the default dispatcher. + const double replicated_a2a_expert_scale = + env_enabled("QWEN36_EP_A2A") && !ctx->data_parallel && + ctx->ep_world_size > 1 + ? 1.0 / static_cast(ctx->ep_world_size) + : 1.0; for (int64_t layer = 0; layer < ctx->num_layers; ++layer) { auto table = lora_projection_table(ctx->layer_configs[layer]); int64_t offset = ctx->lora_layer_offset[layer]; @@ -1977,9 +2357,11 @@ static void synchronize_lora_gradients( if (table.entries[pair].grouped_expert) { if (accumulated_token_weight > 0.0) { reduce_lora_accumulator( - ctx, ctx->grad_accum_a[offset + pair], scale, allreduce); + ctx, ctx->grad_accum_a[offset + pair], + scale * replicated_a2a_expert_scale, allreduce); reduce_lora_accumulator( - ctx, ctx->grad_accum_b[offset + pair], scale, allreduce); + ctx, ctx->grad_accum_b[offset + pair], + scale * replicated_a2a_expert_scale, allreduce); } continue; } diff --git a/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp index 1f2c33b5..d9eae509 100644 --- a/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include struct LayerConfig { @@ -108,6 +109,15 @@ int main() { const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); const int local_rank = std::atoi( std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); + const int a2a = std::getenv("QWEN36_EP_A2A") && + std::strcmp(std::getenv("QWEN36_EP_A2A"), "0") != 0; + // Variable-split A2A receives replicated source rows in rank order. The + // resulting BF16 GEMM/accumulation order is not bit-identical to the + // single-rank full-expert reference, while the Adam update oracle remains + // exact. Keep the legacy parity threshold strict and bound A2A drift. + const double param_tol = a2a ? 2e-4 : 1e-5; + const double m_tol = a2a ? 5e-3 : 1e-5; + const double v_tol = a2a ? 1e-5 : 1e-6; assert(world == 2 && rank >= 0 && rank < world); assert(!std::getenv("TP_SIZE") || std::atoi(std::getenv("TP_SIZE")) == 1); c10::cuda::CUDAGuard guard(local_rank); @@ -335,14 +345,14 @@ int main() { *reinterpret_cast(distributed_v[17])); std::printf( - "native_qwen36_ep_parity rank=%d world=%d top_k=2 " + "native_qwen36_ep_parity rank=%d world=%d top_k=2 a2a=%d " "distributed_loss=%0.8f reference_loss=%0.8f loss_diff=%0.8e " "gate_up_a_diff=%0.8e gate_up_b_diff=%0.8e " "down_a_diff=%0.8e down_b_diff=%0.8e " "adam_m_diff=%0.8e adam_v_diff=%0.8e " "adam_step_diffs=[%0.8e,%0.8e,%0.8e,%0.8e] " "updates=[%0.8e,%0.8e,%0.8e,%0.8e]\n", - rank, world, distributed_loss, reference_loss, loss_diff, + rank, world, a2a, distributed_loss, reference_loss, loss_diff, gate_up_a_diff, gate_up_b_diff, down_a_diff, down_b_diff, optimizer_m_diff, optimizer_v_diff, gate_up_a_adam_diff, gate_up_b_adam_diff, @@ -355,12 +365,12 @@ int main() { assert(down_a_update > 0.0); assert(down_b_update > 0.0); assert(loss_diff <= 2e-2); - assert(gate_up_a_diff <= 1e-5); - assert(gate_up_b_diff <= 1e-5); - assert(down_a_diff <= 1e-5); - assert(down_b_diff <= 1e-5); - assert(optimizer_m_diff <= 1e-5); - assert(optimizer_v_diff <= 1e-6); + assert(gate_up_a_diff <= param_tol); + assert(gate_up_b_diff <= param_tol); + assert(down_a_diff <= param_tol); + assert(down_b_diff <= param_tol); + assert(optimizer_m_diff <= m_tol); + assert(optimizer_v_diff <= v_tol); assert(gate_up_a_adam_diff <= 1e-5); assert(gate_up_b_adam_diff <= 1e-5); assert(down_a_adam_diff <= 1e-5); From 1a721839f1b7bdf46b3a709aa6b0d18c9f86a3ad Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 07:37:02 +0800 Subject: [PATCH 021/156] feat: add sharded fixed lora ep a2a --- .../kernels/qwen3_6_kernels.cpp | 61 +++++++++------ crates/rustrain-qwen3-6/src/session.rs | 15 +++- .../tests/native_ep_smoke.cpp | 75 ++++++++++++++----- docs/plans/qwen-lora-megatron-progress.md | 9 ++- docs/qwen35-qwen36-megatron-audit.md | 14 ++-- 5 files changed, 126 insertions(+), 48 deletions(-) diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index edacdd9d..5ece667e 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -1301,6 +1301,11 @@ static at::Tensor moe_forward( topk_weights = topk_weights.to(compute_type); auto routed_output = at::zeros(flat.sizes(), flat.options()); + if (env_enabled("QWEN36_EP_A2A_SHARDED") && expert_count < num_experts) { + TORCH_CHECK(nccl_comm_v && env_enabled("QWEN36_EP_A2A"), + "QWEN36_EP_A2A_SHARDED=1 requires an initialized EP communicator " + "and QWEN36_EP_A2A=1"); + } bool use_a2a = false; if (nccl_comm_v && env_enabled("QWEN36_EP_A2A") && !expert_gate_up_lora && !expert_down_lora) { @@ -2212,19 +2217,28 @@ static void reduce_lora_accumulator_weighted( accumulator.copy_(reduced); } -// Every rank evaluates the complete loss. Average replicated LoRA gradients -// across DP ranks with token-count weighting so their Adam update matches a -// single global batch. Pure DP keeps the complete routed-expert tensors -// replicated, so those accumulators reduce too; EP ranks receive the complete -// routed activation in forward and keep their sharded expert gradients local. +// Fixed-LoRA gradients are accumulated as token-weighted numerators. Replicated +// DP sums all replicated parameters and divides by the global token count; +// legacy EP keeps its replicated batch/local expert semantics. Sharded A2A +// sums non-expert parameters across source ranks, while expert parameters have +// already received all source numerators through the inverse A2A and therefore +// only divide by the global token count. static void synchronize_lora_gradients( TrainingContext* ctx, const at::Tensor& target_mask, double accumulated_token_weight = 0.0, const at::Tensor* per_adapter_token_counts = nullptr ) { - const bool allreduce = ctx->nccl_comm && ctx->data_parallel; + const bool sharded_a2a = ctx->nccl_comm && !ctx->data_parallel && + env_enabled("QWEN36_EP_A2A_SHARDED"); + TORCH_CHECK(!sharded_a2a || env_enabled("QWEN36_EP_A2A"), + "QWEN36_EP_A2A_SHARDED=1 requires QWEN36_EP_A2A=1"); + const bool dp_allreduce = ctx->nccl_comm && ctx->data_parallel; + const bool normalization_allreduce = dp_allreduce || sharded_a2a; const bool per_adapter_weighting = per_adapter_token_counts && per_adapter_token_counts->defined(); + TORCH_CHECK(!sharded_a2a || !per_adapter_weighting, + "dynamic multi-LoRA is not supported with sharded EP A2A; " + "source tenant metadata is not implemented"); at::Tensor local_adapter_weights; at::Tensor global_adapter_weights; double scale = 1.0; @@ -2239,7 +2253,7 @@ static void synchronize_lora_gradients( "dynamic LoRA token counts must be finite"); TORCH_CHECK((local_adapter_weights >= 0).all().item(), "dynamic LoRA token counts must be non-negative"); - if (allreduce) { + if (dp_allreduce) { global_adapter_weights = at::empty_like(local_adapter_weights); auto stream = c10::cuda::getCurrentCUDAStream( local_adapter_weights.device().index()).stream(); @@ -2258,7 +2272,7 @@ static void synchronize_lora_gradients( "every dynamic LoRA adapter must have at least one global target token"); } else if (accumulated_token_weight > 0.0) { double global_weight = accumulated_token_weight; - if (allreduce) { + if (normalization_allreduce) { auto local = at::full({1}, accumulated_token_weight, at::TensorOptions().dtype(at::kFloat).device(target_mask.device())); auto global = at::empty_like(local); @@ -2277,7 +2291,7 @@ static void synchronize_lora_gradients( // Dynamic multi-LoRA currently contributes one independently-normalized // row per tenant. Preserve that contract while weighting replicated DP // ranks by the selected batch's token count. - if (!allreduce) return; + if (!dp_allreduce) return; auto shifted_mask = target_mask.narrow(1, 1, target_mask.size(1) - 1) .to(at::kFloat).sum().reshape({1}); auto global_mask = at::empty_like(shifted_mask); @@ -2309,10 +2323,10 @@ static void synchronize_lora_gradients( {static_cast(adapter_index)}); reduce_lora_accumulator_weighted( ctx, accum_it->second[pair][0], local_weight, - global_weight, allreduce); + global_weight, dp_allreduce); reduce_lora_accumulator_weighted( ctx, accum_it->second[pair][1], local_weight, - global_weight, allreduce); + global_weight, dp_allreduce); continue; } // Routed-expert tensors are sharded only in EP. Pure DP has @@ -2322,16 +2336,16 @@ static void synchronize_lora_gradients( if (table.entries[pair].grouped_expert) { if (accumulated_token_weight > 0.0) { reduce_lora_accumulator( - ctx, accum_it->second[pair][0], scale, allreduce); + ctx, accum_it->second[pair][0], scale, dp_allreduce); reduce_lora_accumulator( - ctx, accum_it->second[pair][1], scale, allreduce); + ctx, accum_it->second[pair][1], scale, dp_allreduce); } continue; } reduce_lora_accumulator( - ctx, accum_it->second[pair][0], scale, allreduce); + ctx, accum_it->second[pair][0], scale, dp_allreduce); reduce_lora_accumulator( - ctx, accum_it->second[pair][1], scale, allreduce); + ctx, accum_it->second[pair][1], scale, dp_allreduce); } } } @@ -2344,8 +2358,8 @@ static void synchronize_lora_gradients( // shards will require global-token weighting before this gate can become // the default dispatcher. const double replicated_a2a_expert_scale = - env_enabled("QWEN36_EP_A2A") && !ctx->data_parallel && - ctx->ep_world_size > 1 + ctx->nccl_comm && env_enabled("QWEN36_EP_A2A") && !sharded_a2a && + !ctx->data_parallel && ctx->ep_world_size > 1 ? 1.0 / static_cast(ctx->ep_world_size) : 1.0; for (int64_t layer = 0; layer < ctx->num_layers; ++layer) { @@ -2358,17 +2372,19 @@ static void synchronize_lora_gradients( if (accumulated_token_weight > 0.0) { reduce_lora_accumulator( ctx, ctx->grad_accum_a[offset + pair], - scale * replicated_a2a_expert_scale, allreduce); + scale * replicated_a2a_expert_scale, dp_allreduce); reduce_lora_accumulator( ctx, ctx->grad_accum_b[offset + pair], - scale * replicated_a2a_expert_scale, allreduce); + scale * replicated_a2a_expert_scale, dp_allreduce); } continue; } reduce_lora_accumulator( - ctx, ctx->grad_accum_a[offset + pair], scale, allreduce); + ctx, ctx->grad_accum_a[offset + pair], scale, + dp_allreduce || sharded_a2a); reduce_lora_accumulator( - ctx, ctx->grad_accum_b[offset + pair], scale, allreduce); + ctx, ctx->grad_accum_b[offset + pair], scale, + dp_allreduce || sharded_a2a); } } } @@ -4852,6 +4868,9 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( auto* ctx = reinterpret_cast(ctx_ptr); TORCH_CHECK(!ctx->topology_invalid, "native Qwen context rejected an incompatible TP/DP/EP topology"); + TORCH_CHECK(!env_enabled("QWEN36_EP_A2A_SHARDED"), + "dynamic multi-LoRA is not supported with sharded EP A2A; " + "source tenant metadata is not implemented"); GradientAccumulationFailureGuard accumulation_guard{ctx}; if (ctx->nccl_comm) { c10::cuda::set_device(ctx->cuda_device); diff --git a/crates/rustrain-qwen3-6/src/session.rs b/crates/rustrain-qwen3-6/src/session.rs index 04513ee5..83f1599e 100644 --- a/crates/rustrain-qwen3-6/src/session.rs +++ b/crates/rustrain-qwen3-6/src/session.rs @@ -249,6 +249,19 @@ fn train_impl( let shard_ref = ep_shard.as_ref(); let is_ep = shard_ref.is_some(); + let env_enabled = |name: &str| { + std::env::var(name) + .map(|value| !value.is_empty() && value != "0") + .unwrap_or(false) + }; + let ep_a2a = env_enabled("QWEN36_EP_A2A"); + let ep_a2a_sharded = env_enabled("QWEN36_EP_A2A_SHARDED"); + if ep_a2a_sharded && !is_ep { + bail!("QWEN36_EP_A2A_SHARDED=1 requires expert-parallel training"); + } + if ep_a2a_sharded && !ep_a2a { + bail!("QWEN36_EP_A2A_SHARDED=1 requires QWEN36_EP_A2A=1"); + } // Non-EP Qwen sessions may run replicated-weight LoRA data parallelism. // The launcher supplies the standard torchrun environment; EP keeps its // explicit shard metadata as the source of truth. @@ -472,7 +485,7 @@ fn train_impl( let mut loss_value = 0.0; for accumulation_index in 0..gradient_accumulation_steps { let micro_step = step * gradient_accumulation_steps + accumulation_index; - let data_start = if is_data_parallel { + let data_start = if is_data_parallel || ep_a2a_sharded { (micro_step * batch_size * world_size + rank * batch_size) % data.len() } else { (micro_step * batch_size) % data.len() diff --git a/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp index d9eae509..c5af117f 100644 --- a/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp @@ -111,12 +111,16 @@ int main() { std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); const int a2a = std::getenv("QWEN36_EP_A2A") && std::strcmp(std::getenv("QWEN36_EP_A2A"), "0") != 0; - // Variable-split A2A receives replicated source rows in rank order. The - // resulting BF16 GEMM/accumulation order is not bit-identical to the - // single-rank full-expert reference, while the Adam update oracle remains - // exact. Keep the legacy parity threshold strict and bound A2A drift. - const double param_tol = a2a ? 2e-4 : 1e-5; - const double m_tol = a2a ? 5e-3 : 1e-5; + const int sharded = a2a && std::getenv("QWEN36_EP_A2A_SHARDED") && + std::strcmp(std::getenv("QWEN36_EP_A2A_SHARDED"), "0") != 0; + // Replicated A2A receives duplicate source rows in rank order, so its BF16 + // accumulation is not bit-identical to the full-expert reference. Sharded + // A2A uses distinct rows; its optimizer state matches closely, while a + // near-boundary update can land in the adjacent BF16 parameter bin. Keep + // the legacy threshold strict, bound each BF16 case separately, and retain + // the exact standard-Adam oracle below for all modes. + const double param_tol = sharded ? 2e-3 : (a2a ? 2e-4 : 1e-5); + const double m_tol = sharded ? 5e-4 : (a2a ? 5e-3 : 1e-5); const double v_tol = a2a ? 1e-5 : 1e-6; assert(world == 2 && rank >= 0 && rank < world); assert(!std::getenv("TP_SIZE") || std::atoi(std::getenv("TP_SIZE")) == 1); @@ -268,17 +272,47 @@ int main() { // qwen36_init_nccl mutates only the distributed context's copied configs. assert(qwen36_init_nccl(distributed_ctx) == 0); - auto input_ids = at::tensor({1, 2, 3}, - at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 3}); - auto target_mask = at::ones({1, 3}, - at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); - auto attention_mask = at::ones({1, 3}, - at::TensorOptions().device(at::kCUDA).dtype(at::kBool)); + at::Tensor input_ids; + at::Tensor target_mask; + at::Tensor attention_mask; + at::Tensor reference_input_ids; + at::Tensor reference_target_mask; + at::Tensor reference_attention_mask; + auto long_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kLong); + auto float_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kFloat); + auto bool_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kBool); + if (sharded) { + // Deliberately unequal supervised-token counts: rank 0 contributes one + // response token, rank 1 contributes three. The reference evaluates + // the deterministic global batch on a no-NCCL full-expert context. + input_ids = rank == 0 + ? at::tensor({1, 2, 3, 4}, long_opts).reshape({1, 4}) + : at::tensor({4, 5, 6, 7}, long_opts).reshape({1, 4}); + target_mask = rank == 0 + ? at::tensor({0, 1, 0, 0}, float_opts).reshape({1, 4}) + : at::tensor({0, 1, 1, 1}, float_opts).reshape({1, 4}); + attention_mask = at::ones({1, 4}, bool_opts); + reference_input_ids = at::cat({ + at::tensor({1, 2, 3, 4}, long_opts).reshape({1, 4}), + at::tensor({4, 5, 6, 7}, long_opts).reshape({1, 4})}, 0); + reference_target_mask = at::cat({ + at::tensor({0, 1, 0, 0}, float_opts).reshape({1, 4}), + at::tensor({0, 1, 1, 1}, float_opts).reshape({1, 4})}, 0); + reference_attention_mask = at::ones({2, 4}, bool_opts); + } else { + input_ids = at::tensor({1, 2, 3}, long_opts).reshape({1, 3}); + target_mask = at::ones({1, 3}, float_opts); + attention_mask = at::ones({1, 3}, bool_opts); + reference_input_ids = input_ids; + reference_target_mask = target_mask; + reference_attention_mask = attention_mask; + } const double distributed_loss = qwen36_train_step( distributed_ctx, &input_ids, &target_mask, &attention_mask); const double reference_loss = qwen36_train_step( - reference_ctx, &input_ids, &target_mask, &attention_mask); + reference_ctx, &reference_input_ids, &reference_target_mask, + &reference_attention_mask); c10::cuda::device_synchronize(); assert(distributed_loss > 0.0 && std::isfinite(distributed_loss)); assert(reference_loss > 0.0 && std::isfinite(reference_loss)); @@ -294,7 +328,11 @@ int main() { *distributed_lora.down_a, reference_slice(*reference_lora.down_a)); const double down_b_diff = max_abs_diff( *distributed_lora.down_b, reference_slice(*reference_lora.down_b)); - const double loss_diff = std::abs(distributed_loss - reference_loss); + // qwen36_train_step returns each rank's local mean in sharded mode; the + // global weighted scalar is intentionally compared through parameter/state + // parity below rather than against the full-batch reference scalar. + const double loss_diff = sharded ? -1.0 : + std::abs(distributed_loss - reference_loss); const double gate_up_a_update = update_norm( *distributed_lora.gate_up_a, gate_up_a_before); @@ -345,14 +383,17 @@ int main() { *reinterpret_cast(distributed_v[17])); std::printf( - "native_qwen36_ep_parity rank=%d world=%d top_k=2 a2a=%d " + "native_qwen36_ep_parity rank=%d world=%d top_k=2 a2a=%d sharded=%d " + "loss_compare=%s local_tokens=%d " "distributed_loss=%0.8f reference_loss=%0.8f loss_diff=%0.8e " "gate_up_a_diff=%0.8e gate_up_b_diff=%0.8e " "down_a_diff=%0.8e down_b_diff=%0.8e " "adam_m_diff=%0.8e adam_v_diff=%0.8e " "adam_step_diffs=[%0.8e,%0.8e,%0.8e,%0.8e] " "updates=[%0.8e,%0.8e,%0.8e,%0.8e]\n", - rank, world, a2a, distributed_loss, reference_loss, loss_diff, + rank, world, a2a, sharded, sharded ? "skipped" : "direct", + sharded ? (rank == 0 ? 1 : 3) : 2, + distributed_loss, reference_loss, loss_diff, gate_up_a_diff, gate_up_b_diff, down_a_diff, down_b_diff, optimizer_m_diff, optimizer_v_diff, gate_up_a_adam_diff, gate_up_b_adam_diff, @@ -364,7 +405,7 @@ int main() { assert(gate_up_b_update > 0.0); assert(down_a_update > 0.0); assert(down_b_update > 0.0); - assert(loss_diff <= 2e-2); + assert(sharded || loss_diff <= 2e-2); assert(gate_up_a_diff <= param_tol); assert(gate_up_b_diff <= param_tol); assert(down_a_diff <= param_tol); diff --git a/docs/plans/qwen-lora-megatron-progress.md b/docs/plans/qwen-lora-megatron-progress.md index 36ab22fb..839917c3 100644 --- a/docs/plans/qwen-lora-megatron-progress.md +++ b/docs/plans/qwen-lora-megatron-progress.md @@ -12,9 +12,9 @@ timestamp: 2026-07-17T00:00:00Z # Current State -Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP parity against a full-expert reference, dense replicated-DP smoke with per-tenant token weighting, TP-only latent-rank-sharded LoRA smoke, dynamic batch logical-step update, ABI11 FP32 gradient storage/aggregation, standard Adam bias correction, per-tenant optimizer-step restore, selected-tenant isolation, same-topology rank-aware checkpointing, and 5D topology mapping. +Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP parity against a full-expert reference, variable-split EP A2A with fixed-LoRA data sharding, dense replicated-DP smoke with per-tenant token weighting, TP-only latent-rank-sharded LoRA smoke, dynamic batch logical-step update, ABI11 FP32 gradient storage/aggregation, standard Adam bias correction, per-tenant optimizer-step restore, selected-tenant isolation, same-topology rank-aware checkpointing, and 5D topology mapping. -Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axis TP+DP/EP, PP/CP, variable-split token A2A, DeepEP/TE prebuilt integration, cross-topology checkpoint resharding, and matched Megatron throughput. +Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axis TP+DP/EP, PP/CP, DeepEP/TE prebuilt integration, dynamic multi-LoRA source metadata under sharded A2A, cross-topology checkpoint resharding, and matched Megatron throughput. # Durable Milestones @@ -27,12 +27,15 @@ Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axi - `ff382a0`: native ABI rejects unsupported TP+DP/EP mixtures before training. Dynamic per-tenant DP batches now use an all-reduced per-adapter token-count vector and weighted FP32 LoRA gradient reduction, including grouped expert LoRA. - `76eace6`: direct native context creation also rejects PP/CP sizes greater than one, so unsupported pipeline/context parallelism cannot silently fall back to replicated single-stage execution. - `b7050ea`: fixed standard Adam bias correction, separated pure-DP gradient NCCL from MoE activation collectives, made dynamic-MoE DP token weighting numerical, scoped n_max rendezvous by invocation with atomic publication, and added full-expert EP parity plus Adam oracles. +- `bfa3d8e`: added a gated variable-split NCCL dispatch/inverse-combine prototype with differentiable forward/backward collectives; default legacy EP remains unchanged. +- Working tree after `bfa3d8e`: `QWEN36_EP_A2A_SHARDED=1` disjoins EP input rows, all-reduces global fixed-LoRA token numerators, keeps grouped expert LoRA local after A2A, and explicitly rejects dynamic multi-LoRA until source tenant metadata exists. - H20 `123.57.26.97:28004`: ABI8 native smoke passed grouped/fallback parity, GDN, dense/MoE LoRA, dynamic adapters, and step setter validation. - H20 `123.57.26.97:28004`: ABI9 native smoke passed selected-tenant training with a positive selected update, exactly zero unselected update, independent clocks (`2` vs `1`), and registry preservation after an unknown ID. - H20 `123.57.26.97:28004`: ABI10 two-rank TP native smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, and selected-tenant isolation. Rank-local LoRA tensors used distinct rank `4` slices for global rank `8`; losses matched and both shards had positive updates. - H20 `123.57.26.97:28004`: ABI11 single-rank smoke passed FP32 accumulator dtype, two-micro accumulation, NaN/explicit abort cleanup, successful step commit, standard Adam parameter oracle, TP/DP/EP and PP/CP topology rejection, and dynamic multi-LoRA. ABI11 two-rank TP smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, selected-tenant isolation, invocation-scoped n_max rendezvous, and the new guards (`rank_statuses=0,0`). - H20 `123.57.26.97:28004`: ABI11 DP2 smoke passed with per-tenant masks `[1,3]` and `[3,1]`: weighted m relative error `2.43e-8`, grouped-expert error `2.27e-8`, v error `7.33e-8`, BF16 Adam delta error `0`, and nonzero gap `7.96e-3` versus the old equal-count formula (`rank_statuses=0,0`). - H20 `123.57.26.97:28004`: ABI11 EP2 full-expert parity smoke passed with distinct rank-local base and LoRA expert slices; loss, A/B updates, m/v, and standard Adam first-step oracle all matched the rank-local full-expert reference (`rank_statuses=0,0`). +- H20 `123.57.26.97:28004`: gated A2A EP2 smoke passed in legacy, replicated-source, and sharded-source modes (`rank_statuses=0,0`). Sharded mode used token counts `[1,3]`; weighted global loss differed from the full-batch reference by `9.6e-7`, expert m/v maxima were `1.22e-5` / `3.92e-9`, and every standard Adam oracle was zero. One distributed BF16 parameter landed in the adjacent rounding bin (`9.61e-4`), bounded separately by the smoke. - Target runtime probe: PyTorch 2.5.1+cu121, ABI0; Transformer Engine, flash-attn, DeepEP, Triton, and DeepSpeed are not importable. # Decisions During Execution @@ -46,4 +49,4 @@ Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axi Passed: `cargo check -p rustrain-qwen3-6 -p rustrain-server -p rustrain-ipc` (with the repository PyTorch 2.12.1 host venv), `cargo test -p rustrain-server --lib` (6), Qwen unit tests (3), Qwen integration (6), remote ABI8 smoke, remote ABI9 selected-tenant native smoke, remote ABI10 single-rank native smoke, remote ABI10 two-rank TP native smoke, and remote ABI11 single/TP2/DP2/EP2 native smoke with numerical Adam and parity oracles. -Not run: Megatron-style base-model TP, variable-split EP token dispatch, multi-axis TP+DP/EP, PP/CP, cross-topology resharding, a numerical FP32-accumulation oracle against a concatenated batch, and matched Megatron performance benchmark. +Not run: Megatron-style base-model TP, multi-axis TP+DP/EP, PP/CP, dynamic multi-LoRA source metadata under sharded A2A, cross-topology resharding, and matched Megatron performance benchmark. The target host lacks importable Megatron/Transformer Engine/DeepEP/flash-attn prebuilt packages, so no dependency installation or JIT workaround was used. diff --git a/docs/qwen35-qwen36-megatron-audit.md b/docs/qwen35-qwen36-megatron-audit.md index 92ca95ee..a7e52237 100644 --- a/docs/qwen35-qwen36-megatron-audit.md +++ b/docs/qwen35-qwen36-megatron-audit.md @@ -7,7 +7,7 @@ - 模型语义：Qwen3.5 dense、Qwen3.6 dense/MoE 的 native forward/backward 路径已经覆盖 hybrid full attention、GDN、MoE、MTP 和 LoRA 目标模块；已有配置解析、集成测试及 H20 native smoke 证据。 - 已实现并可验证的分布式子集：MoE expert parallel，以及 replicated LoRA 的 data parallel；梯度累积和 dynamic multi-LoRA 已有 logical-step 边界。DP 动态租户按 adapter token count 加权，纯 DP 不把 world communicator 传入 MoE activation reduction。 - 性能：MoE grouped dispatch 相对逐 expert matmul 的已有 microbenchmark 为约 3.70x（E=32, N=4096, H=2048, I=768，结果误差为 0）；这不是端到端训练吞吐或 Megatron 对比。 -- 尚未实现：Qwen native 路径的 tensor parallel、pipeline parallel、context parallel，以及 TP/PP/CP 与 EP/DP 的组合。当前训练上下文仍由单个进程持有完整 dense 权重和完整层栈。 +- 已实现 LoRA latent-rank 的 TP-only 子集，但尚未实现 frozen base-weight tensor parallel、pipeline parallel、context parallel，以及 TP/PP/CP 与 EP/DP 的组合。当前训练上下文仍由单个进程持有完整 dense 权重和完整层栈。 - 因此当前实现不能宣称“Megatron-LM 级别”。它是一个计算集中在 C++ 的 LoRA/EP/DP 子集，离 Megatron 的完整并行和通信重叠仍有实质差距。 ## 当前能力矩阵 @@ -22,11 +22,11 @@ | dynamic multi-LoRA | 已实现子集 | 请求按 adapter 分组，单个 logical step 统一 backward/Adam；每个 adapter 独立 optimizer clock 和 m/v，DP 按 token count 加权 | | microbatch accumulation | 已实现子集 | non-final microbatch 只 backward，final microbatch 才 optimizer；FP32 accumulator 存储/聚合，autograd leaf backward 仍为 BF16 | | replicated data parallel | 已实现 | logical-step 边界同步 replicated LoRA；EP expert 参数不走该 reduction | -| expert parallel | 已实现子集 | 路由输出 all-reduce 和本地 expert 权重；没有 DeepEP 式 fused A2A/dispatch overlap | -| tensor parallel | 未实现于 Qwen native | 不切分 attention/MLP/LM-head 权重，也没有 Qwen TP communicator | +| expert parallel | 已实现子集 | 默认 routed-output all-reduce；gated variable-split A2A 已验证 fixed-LoRA data sharding，但 dynamic tenant metadata、GPU-only split planning 和 overlap 未实现 | +| tensor parallel | LoRA-only 子集 | latent rank 分片和独立 TP communicator 已验证；attention/MLP/LM-head base 权重仍不切分 | | pipeline parallel | 未实现于 Qwen native | 没有 stage 切分、microbatch scheduler 或 activation send/recv | | context parallel | 未实现于 Qwen native | 没有 ring attention、跨 rank KV/索引合并 | -| TP/PP/CP 组合 checkpoint | 未实现 | 当前 checkpoint 不是 Megatron rank-sharded topology | +| distributed checkpoint | 已实现子集 | same-topology TP rank-sharded v3 已验证；跨 topology reshard 和 PP/CP 未实现 | ## 与 Megatron-LM 的关键差距 @@ -34,7 +34,7 @@ Megatron 的 TP 会按 head、hidden/intermediate 和 vocab 维度切分权重，并在线性层边界执行必要的 reduce-scatter/all-reduce；PP 会把层分到不同 stage 并使用 1F1B 等调度；CP 会在序列和 attention state 上做跨 rank 通信。当前 Qwen native `TrainingContext` 仍加载完整模型并在一个 C++ forward 中执行全部层，因此仅增加 `tensor_model_parallel_size` 等配置不能得到正确的 TP/PP/CP。 -当前 DP/EP 也不是完整 Megatron 语义：DP 同步 replicated LoRA 梯度并按租户 token count 归一化，expert 参数留在 EP rank；EP 使用 routed-output all-reduce，但没有 fused token dispatch/combine、异步 A2A 和通信计算重叠。 +当前 DP/EP 也不是完整 Megatron 语义：DP 同步 replicated LoRA 梯度并按租户 token count 归一化，expert 参数留在 EP rank；EP 默认使用 routed-output all-reduce，实验 gate 已加入 variable-split dispatch/inverse combine 和 fixed-LoRA data sharding。该 gate 仍逐 top-k 执行 host-visible count sync，没有 fused permutation、dynamic tenant source metadata、异步 overlap 或 DeepEP backend。 ### 优化器与恢复 @@ -44,9 +44,11 @@ Megatron 的 TP 会按 head、hidden/intermediate 和 vocab 维度切分权重 当前粗粒度 C++ FFI、grouped MoE 和 activation checkpoint/offload 是有效优化，但尚无 Megatron/Transformer Engine 级别的端到端数据：没有完整模型在同一 GPU、序列长度、microbatch、精度和通信配置下的 tokens/s、显存、扩展效率对照，也没有 FP8/FP4 参数与 fused attention/DeepEP 的 Qwen 路径。 +目标 H20 的 ABI1 环境有 PyTorch 2.12.1、Triton 和 NumPy，但没有 Megatron、Transformer Engine、DeepEP、flash-attn、Apex 或缓存的兼容 prebuilt wheel。本地 Megatron 的 Qwen3.5 35B-A3B 入口强制 TE/flash-attn，且是 full-parameter SFT，不提供 trainable LoRA wrapper；其 `moe_perf` 也固定 TE grouped MLP 和 H100 条件。因此当前不能诚实地产出 matched Megatron-LoRA benchmark，且本工作没有通过 JIT 或自构建依赖绕过该限制。 + ## 验证边界 -已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试，以及 H20 ABI11 的单卡、TP2、EP2 和 DP2 native smoke。DP2 的 weighted m/grouped/v/Adam BF16 delta oracle 分别达到 `2.43e-8`、`2.27e-8`、`7.33e-8` 和 `0`；EP2 与 full-expert reference 的 loss、LoRA、Adam state 和标准 Adam 首步 oracle 在两 rank 均为零差异。没有完成 Qwen3.5/3.6 完整大模型的长时间训练、跨节点通信、TP/PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖的模型/LoRA/EP/DP 子集，而不是所有并行配置。 +已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试，以及 H20 ABI11 的单卡、TP2、EP2 和 DP2 native smoke。DP2 的 weighted m/grouped/v/Adam BF16 delta oracle 分别达到 `2.43e-8`、`2.27e-8`、`7.33e-8` 和 `0`；legacy EP2 与 full-expert reference 的 loss、LoRA、Adam state 和标准 Adam 首步 oracle 在两 rank 均为零差异。Replicated A2A 与 sharded A2A 也均通过两 rank full-expert reference；sharded token counts `[1,3]` 的加权 loss 与 global reference 相差约 `9.6e-7`，m/v 最大差 `1.22e-5` / `3.92e-9`，Adam oracle 差为 `0`。没有完成 Qwen3.5/3.6 完整大模型的长时间训练、跨节点通信、PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖且有直接 oracle 的子集，而不是所有并行配置。 ## 继续达到 Megatron 级别所需的最小工作包 From b51e9c109eb5f6de5245fbc533e695930546bb5f Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 07:45:41 +0800 Subject: [PATCH 022/156] docs: add native ep benchmark evidence --- .../tests/native_ep_bench.cpp | 211 ++++++++++++++++++ docs/plans/qwen-lora-megatron-progress.md | 1 + docs/qwen35-qwen36-megatron-audit.md | 28 +++ 3 files changed, 240 insertions(+) create mode 100644 crates/rustrain-qwen3-6/tests/native_ep_bench.cpp diff --git a/crates/rustrain-qwen3-6/tests/native_ep_bench.cpp b/crates/rustrain-qwen3-6/tests/native_ep_bench.cpp new file mode 100644 index 00000000..bbaf70d1 --- /dev/null +++ b/crates/rustrain-qwen3-6/tests/native_ep_bench.cpp @@ -0,0 +1,211 @@ +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +struct LayerConfig { + int64_t layer_type, num_heads, num_kv_heads, head_dim; + int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel; + double partial_rotary_factor, rope_theta, rms_eps; + int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count; + int64_t intermediate_size; + int32_t norm_topk_prob; + void* nccl_comm; + void* nccl_stream; +}; + +extern "C" void* qwen36_create_training_context( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*); +extern "C" int32_t qwen36_init_nccl(void*); +extern "C" double qwen36_train_step(void*, void*, void*, void*); +extern "C" void qwen36_free_training_context(void*); + +static int env_int(const char* name, int fallback) { + const char* value = std::getenv(name); + return value ? std::max(1, std::atoi(value)) : fallback; +} + +static at::Tensor cuda_rand(std::initializer_list shape) { + return at::randn(shape, + at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)); +} + +static std::vector tensor_ptrs(std::vector& tensors) { + std::vector result; + result.reserve(tensors.size()); + for (auto& tensor : tensors) result.push_back(&tensor); + return result; +} + +static double percentile(std::vector values, double q) { + std::sort(values.begin(), values.end()); + const double pos = q * static_cast(values.size() - 1); + const size_t lo = static_cast(pos); + const size_t hi = std::min(lo + 1, values.size() - 1); + const double frac = pos - static_cast(lo); + return values[lo] * (1.0 - frac) + values[hi] * frac; +} + +int main() { + const int rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0"); + const int world = std::atoi( + std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); + const int local_rank = std::atoi( + std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); + const bool a2a = std::getenv("QWEN36_EP_A2A") && + std::atoi(std::getenv("QWEN36_EP_A2A")) != 0; + const bool sharded = a2a && std::getenv("QWEN36_EP_A2A_SHARDED") && + std::atoi(std::getenv("QWEN36_EP_A2A_SHARDED")) != 0; + const int seq = env_int("BENCH_SEQ", 128); + const int hidden = env_int("BENCH_HIDDEN", 256); + const int experts = env_int("BENCH_EXPERTS", 8); + const int intermediate = env_int("BENCH_INTERMEDIATE", 256); + const int warmup = env_int("BENCH_WARMUP", 2); + const int iters = env_int("BENCH_ITERS", 10); + assert(world >= 2 && rank >= 0 && rank < world); + assert(hidden % 8 == 0 && experts % world == 0); + assert(!sharded || a2a); + c10::cuda::CUDAGuard guard(local_rank); + at::manual_seed(1234); + + const int local_experts = experts / world; + const int64_t expert_start = rank * local_experts; + const int vocab = std::max(1024, hidden * 4); + const int head_dim = hidden; + + // Build one deterministic global model, then narrow only expert tensors. + std::vector global_weights; + global_weights.push_back(cuda_rand({hidden})); + global_weights.push_back(cuda_rand({hidden})); + global_weights.push_back(cuda_rand({2 * head_dim, hidden})); + global_weights.push_back(cuda_rand({head_dim})); + global_weights.push_back(cuda_rand({head_dim, hidden})); + global_weights.push_back(cuda_rand({head_dim})); + global_weights.push_back(cuda_rand({head_dim, hidden})); + global_weights.push_back(cuda_rand({hidden, head_dim})); + global_weights.push_back(cuda_rand({experts, hidden})); + global_weights.push_back(cuda_rand({1, hidden})); + global_weights.push_back(cuda_rand({intermediate, hidden})); + global_weights.push_back(cuda_rand({intermediate, hidden})); + global_weights.push_back(cuda_rand({hidden, intermediate})); + global_weights.push_back(cuda_rand({experts, 2 * intermediate, hidden})); + global_weights.push_back(cuda_rand({experts, hidden, intermediate})); + for (auto& weight : global_weights) weight.set_requires_grad(false); + + std::vector local_weights = global_weights; + local_weights[13] = global_weights[13] + .narrow(0, expert_start, local_experts).contiguous(); + local_weights[14] = global_weights[14] + .narrow(0, expert_start, local_experts).contiguous(); + auto weights = tensor_ptrs(local_weights); + auto embed = cuda_rand({vocab, hidden}); + auto final_norm = cuda_rand({hidden}); + auto lm_head = cuda_rand({vocab, hidden}); + embed.set_requires_grad(false); + final_norm.set_requires_grad(false); + lm_head.set_requires_grad(false); + + LayerConfig config{}; + config.layer_type = 0; + config.num_heads = 1; + config.num_kv_heads = 1; + config.head_dim = head_dim; + config.partial_rotary_factor = 1.0; + config.rope_theta = 10000.0; + config.rms_eps = 1e-5; + config.num_experts = experts; + config.top_k = 2; + config.moe_intermediate = intermediate; + config.expert_start = expert_start; + config.expert_count = local_experts; + config.norm_topk_prob = 1; + + const int64_t target_layer = 0; + const char* targets = "experts_gate_up_proj,experts_down_proj"; + void* ctx = qwen36_create_training_context( + weights.data(), static_cast(weights.size()), + &embed, &final_norm, &lm_head, &config, 1, + static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, 8, + &target_layer, 1, targets); + assert(ctx); + assert(qwen36_init_nccl(ctx) == 0); + + std::vector ids(seq); + for (int i = 0; i < seq; ++i) { + const int offset = sharded ? rank * seq : 0; + ids[i] = (offset + i + 1) % vocab; + } + auto input_ids = at::from_blob(ids.data(), {1, seq}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)).clone().to(at::kCUDA); + auto target_mask = at::ones({1, seq}, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); + auto attention_mask = at::ones({1, seq}, + at::TensorOptions().device(at::kCUDA).dtype(at::kBool)); + + for (int i = 0; i < warmup; ++i) { + (void)qwen36_train_step(ctx, &input_ids, &target_mask, &attention_mask); + } + cudaDeviceSynchronize(); + + cudaEvent_t start = nullptr, stop = nullptr; + cudaEventCreate(&start); + cudaEventCreate(&stop); + std::vector times; + times.reserve(iters); + double last_loss = 0.0; + for (int i = 0; i < iters; ++i) { + cudaEventRecord(start, c10::cuda::getCurrentCUDAStream().stream()); + last_loss = qwen36_train_step( + ctx, &input_ids, &target_mask, &attention_mask); + cudaEventRecord(stop, c10::cuda::getCurrentCUDAStream().stream()); + cudaEventSynchronize(stop); + float elapsed_ms = 0.0f; + cudaEventElapsedTime(&elapsed_ms, start, stop); + times.push_back(static_cast(elapsed_ms)); + } + cudaEventDestroy(start); + cudaEventDestroy(stop); + + const double mean = std::accumulate(times.begin(), times.end(), 0.0) / + static_cast(times.size()); + double variance = 0.0; + for (const double value : times) variance += (value - mean) * (value - mean); + variance /= static_cast(times.size()); + size_t free_bytes = 0, total_bytes = 0; + cudaMemGetInfo(&free_bytes, &total_bytes); + const double local_tokens = static_cast(std::max(seq - 1, 1)); + const double processed_tokens = local_tokens * world; + // Legacy EP replicates the input batch on every rank; only sharded A2A + // represents distinct global samples in this synthetic harness. + const double unique_tokens = sharded ? processed_tokens : local_tokens; + const double median_seconds = percentile(times, 0.5) / 1000.0; + const double processed_tokens_per_sec = processed_tokens / median_seconds; + const double unique_tokens_per_sec = unique_tokens / median_seconds; + std::printf( + "native_qwen36_ep_bench rank=%d world=%d a2a=%d sharded=%d " + "seq=%d hidden=%d experts=%d intermediate=%d warmup=%d iters=%d " + "local_tokens=%.0f processed_tokens=%.0f unique_tokens=%.0f last_loss=%.8f " + "step_ms_mean=%.4f step_ms_median=%.4f step_ms_std=%.4f " + "processed_tokens_per_sec=%.2f unique_tokens_per_sec=%.2f " + "free_mem_gib=%.3f\n", + rank, world, a2a, sharded, seq, hidden, experts, intermediate, + warmup, iters, local_tokens, processed_tokens, unique_tokens, last_loss, + mean, percentile(times, 0.5), std::sqrt(variance), + processed_tokens_per_sec, unique_tokens_per_sec, + static_cast(free_bytes) / (1024.0 * 1024.0 * 1024.0)); + std::fflush(stdout); + qwen36_free_training_context(ctx); + return 0; +} diff --git a/docs/plans/qwen-lora-megatron-progress.md b/docs/plans/qwen-lora-megatron-progress.md index 839917c3..0acc20b3 100644 --- a/docs/plans/qwen-lora-megatron-progress.md +++ b/docs/plans/qwen-lora-megatron-progress.md @@ -37,6 +37,7 @@ Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axi - H20 `123.57.26.97:28004`: ABI11 EP2 full-expert parity smoke passed with distinct rank-local base and LoRA expert slices; loss, A/B updates, m/v, and standard Adam first-step oracle all matched the rank-local full-expert reference (`rank_statuses=0,0`). - H20 `123.57.26.97:28004`: gated A2A EP2 smoke passed in legacy, replicated-source, and sharded-source modes (`rank_statuses=0,0`). Sharded mode used token counts `[1,3]`; weighted global loss differed from the full-batch reference by `9.6e-7`, expert m/v maxima were `1.22e-5` / `3.92e-9`, and every standard Adam oracle was zero. One distributed BF16 parameter landed in the adjacent rounding bin (`9.61e-4`), bounded separately by the smoke. - Target runtime probe: PyTorch 2.5.1+cu121, ABI0; Transformer Engine, flash-attn, DeepEP, Triton, and DeepSpeed are not importable. +- H20 `123.57.26.97:28004`: `native_ep_bench.cpp` fresh ABI0 benchmark (`seq=128, hidden=256, experts=8, intermediate=256, warmup=2, iters=10`) passed legacy and sharded A2A with `rank_statuses=0,0`. Legacy median was about `5.47 ms` / `46.4k processed tokens/s` (`23.2k unique tokens/s`), while sharded A2A was about `6.72 ms` / `37.8k processed and unique tokens/s`. This is a synthetic native baseline, not Megatron-LM parity. # Decisions During Execution diff --git a/docs/qwen35-qwen36-megatron-audit.md b/docs/qwen35-qwen36-megatron-audit.md index a7e52237..20d8ecb0 100644 --- a/docs/qwen35-qwen36-megatron-audit.md +++ b/docs/qwen35-qwen36-megatron-audit.md @@ -44,6 +44,8 @@ Megatron 的 TP 会按 head、hidden/intermediate 和 vocab 维度切分权重 当前粗粒度 C++ FFI、grouped MoE 和 activation checkpoint/offload 是有效优化，但尚无 Megatron/Transformer Engine 级别的端到端数据：没有完整模型在同一 GPU、序列长度、microbatch、精度和通信配置下的 tokens/s、显存、扩展效率对照，也没有 FP8/FP4 参数与 fused attention/DeepEP 的 Qwen 路径。 +本次 native benchmark 没有证明 gated A2A 的端到端 step-time 优势：在该小型 workload 上 sharded A2A 的中位 step 反而比 legacy 高约 `23%`。它没有实现 DeepEP 的 fused permutation、GPU-only split planning 或通信计算 overlap，也没有覆盖 H=`2048`/E=`256` 的完整 Qwen3.6 workload。legacy 模式复制输入 batch，因此必须同时报告唯一样本吞吐，不能只看所有 rank 的 processed tokens/s。 + 目标 H20 的 ABI1 环境有 PyTorch 2.12.1、Triton 和 NumPy，但没有 Megatron、Transformer Engine、DeepEP、flash-attn、Apex 或缓存的兼容 prebuilt wheel。本地 Megatron 的 Qwen3.5 35B-A3B 入口强制 TE/flash-attn，且是 full-parameter SFT，不提供 trainable LoRA wrapper；其 `moe_perf` 也固定 TE grouped MLP 和 H100 条件。因此当前不能诚实地产出 matched Megatron-LoRA benchmark，且本工作没有通过 JIT 或自构建依赖绕过该限制。 ## 验证边界 @@ -57,3 +59,29 @@ Megatron 的 TP 会按 head、hidden/intermediate 和 vocab 维度切分权重 3. 将 EP dispatch/combine 替换为 fused/异步路径，并测量通信与计算重叠。 4. 为 LoRA 增加 FP32 accumulation、每 adapter optimizer step、可恢复的 accumulation 状态和 rank-sharded checkpoint。 5. 在固定硬件和 workload 上，与 Megatron-LM 记录 tokens/s、step time、峰值显存、通信占比和 loss 曲线。 + +## Native EP Benchmark Artifact + +`crates/rustrain-qwen3-6/tests/native_ep_bench.cpp` is a dependency-free +synthetic baseline for the native Qwen C ABI. It times complete +`qwen36_train_step` calls (forward, backward, and Adam) with CUDA events and +accepts `BENCH_SEQ`, `BENCH_HIDDEN`, `BENCH_EXPERTS`, `BENCH_INTERMEDIATE`, +`BENCH_WARMUP`, and `BENCH_ITERS`. It reports processed and unique tokens/s, +per-rank step statistics, and free memory. This is not a Megatron-LM +comparison and does not claim DeepEP or Transformer Engine parity. + +Fresh H20 ABI0 run (`seq=128, hidden=256, experts=8, intermediate=256, +warmup=2, iters=10`, two ranks): + +| Mode | Rank | Median step | Mean step | Processed tokens/s | Unique tokens/s | Status | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| Legacy EP (`QWEN36_EP_A2A=0`) | 0 | 5.4752 ms | 5.4902 ms | 46,391 | 23,196 | 0 | +| Legacy EP (`QWEN36_EP_A2A=0`) | 1 | 5.4672 ms | 5.4591 ms | 46,459 | 23,230 | 0 | +| Sharded A2A (`QWEN36_EP_A2A=1`, `QWEN36_EP_A2A_SHARDED=1`) | 0 | 6.7192 ms | 6.7459 ms | 37,802 | 37,802 | 0 | +| Sharded A2A (`QWEN36_EP_A2A=1`, `QWEN36_EP_A2A_SHARDED=1`) | 1 | 6.7261 ms | 6.7458 ms | 37,763 | 37,763 | 0 | + +Each rank processed 127 local tokens and 254 processed tokens per step. In +replicated legacy EP, unique tokens/s is 127 tokens divided by step time; +sharded A2A has 254 unique global tokens per step. These are synthetic native +measurements only; workload, precision, packing, and communication semantics +are not matched to Megatron-LM. From 0f39f07f1e9fb37aebc852081ec2735d49bb3e41 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 08:55:47 +0800 Subject: [PATCH 023/156] fix: normalize dynamic sharded EP LoRA gradients --- .../kernels/qwen3_6_kernels.cpp | 199 ++++++++++----- .../tests/native_ep_smoke.cpp | 227 +++++++++++++++++- .../rustrain-qwen3-6/tests/native_smoke.cpp | 30 ++- docs/plans/qwen-lora-megatron-progress.md | 11 +- docs/qwen35-qwen36-megatron-audit.md | 10 +- 5 files changed, 396 insertions(+), 81 deletions(-) diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 5ece667e..e3fecaba 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -1177,6 +1177,10 @@ static at::Tensor moe_routed_a2a( flat.size(0), at::TensorOptions().device(flat.device()).dtype(at::kLong)); for (int64_t kk = 0; kk < top_k; ++kk) { auto expert_indices = topk_indices.select(-1, kk).contiguous(); + // `received_tokens` preserves the source flattened row index through + // dispatch. Dynamic multi-LoRA uses floor_divide(row, seq) to recover + // the tenant/sample row, so sharded A2A does not need a second host + // metadata exchange for adapter IDs. auto dispatched = Qwen36A2ADispatchFunction::apply( flat, expert_indices, token_indices, expert_count, static_cast(reinterpret_cast(comm))); @@ -1307,8 +1311,10 @@ static at::Tensor moe_forward( "and QWEN36_EP_A2A=1"); } bool use_a2a = false; + const bool sharded_a2a_mode = env_enabled("QWEN36_EP_A2A_SHARDED") && + expert_count < num_experts; if (nccl_comm_v && env_enabled("QWEN36_EP_A2A") && - !expert_gate_up_lora && !expert_down_lora) { + ((!expert_gate_up_lora && !expert_down_lora) || sharded_a2a_mode)) { auto comm = reinterpret_cast(nccl_comm_v); int world = 1, rank = 0; auto err = ncclCommCount(comm, &world); @@ -2184,6 +2190,38 @@ static void reduce_lora_accumulator( accumulator.copy_(reduced); } +static void normalize_lora_accumulator_numerator( + TrainingContext* ctx, at::Tensor& accumulator, + const at::Tensor& global_weight, bool allreduce +) { + if (!accumulator.defined()) return; + TORCH_CHECK(accumulator.scalar_type() == at::kFloat, + "LoRA DP gradient accumulator must be FP32"); + TORCH_CHECK(global_weight.numel() == 1, + "per-adapter LoRA global token weight must be scalar"); + auto numerator = accumulator.contiguous(); + at::Tensor reduced; + if (allreduce) { + TORCH_CHECK(ctx->nccl_comm, + "LoRA gradient all-reduce has no communicator"); + reduced = at::empty_like(numerator); + int dev = numerator.device().index(); + cudaSetDevice(dev); + auto stream = c10::cuda::getCurrentCUDAStream(dev).stream(); + auto err = ncclAllReduce( + numerator.data_ptr(), reduced.data_ptr(), numerator.numel(), + nccl_dtype_for(numerator), ncclSum, ctx->nccl_comm, stream); + TORCH_CHECK(err == ncclSuccess, + "NCCL LoRA numerator all-reduce failed: ", + ncclGetErrorString(err)); + } else { + reduced = numerator; + } + reduced = reduced / global_weight.clamp_min(1.0); + at::NoGradGuard guard; + accumulator.copy_(reduced); +} + static void reduce_lora_accumulator_weighted( TrainingContext* ctx, at::Tensor& accumulator, const at::Tensor& local_weight, const at::Tensor& global_weight, @@ -2226,7 +2264,8 @@ static void reduce_lora_accumulator_weighted( static void synchronize_lora_gradients( TrainingContext* ctx, const at::Tensor& target_mask, double accumulated_token_weight = 0.0, - const at::Tensor* per_adapter_token_counts = nullptr + const at::Tensor* per_adapter_token_counts = nullptr, + std::vector* adapter_has_global_tokens = nullptr ) { const bool sharded_a2a = ctx->nccl_comm && !ctx->data_parallel && env_enabled("QWEN36_EP_A2A_SHARDED"); @@ -2236,9 +2275,6 @@ static void synchronize_lora_gradients( const bool normalization_allreduce = dp_allreduce || sharded_a2a; const bool per_adapter_weighting = per_adapter_token_counts && per_adapter_token_counts->defined(); - TORCH_CHECK(!sharded_a2a || !per_adapter_weighting, - "dynamic multi-LoRA is not supported with sharded EP A2A; " - "source tenant metadata is not implemented"); at::Tensor local_adapter_weights; at::Tensor global_adapter_weights; double scale = 1.0; @@ -2253,7 +2289,7 @@ static void synchronize_lora_gradients( "dynamic LoRA token counts must be finite"); TORCH_CHECK((local_adapter_weights >= 0).all().item(), "dynamic LoRA token counts must be non-negative"); - if (dp_allreduce) { + if (dp_allreduce || sharded_a2a) { global_adapter_weights = at::empty_like(local_adapter_weights); auto stream = c10::cuda::getCurrentCUDAStream( local_adapter_weights.device().index()).stream(); @@ -2268,8 +2304,15 @@ static void synchronize_lora_gradients( } else { global_adapter_weights = local_adapter_weights; } - TORCH_CHECK((global_adapter_weights > 0).all().item(), - "every dynamic LoRA adapter must have at least one global target token"); + if (adapter_has_global_tokens) { + adapter_has_global_tokens->assign(ctx->adapters.size(), 0); + auto global_cpu = global_adapter_weights.to( + at::TensorOptions().device(at::kCPU)); + const auto* counts = global_cpu.data_ptr(); + for (size_t i = 0; i < ctx->adapters.size(); ++i) { + (*adapter_has_global_tokens)[i] = counts[i] > 0.0f ? 1 : 0; + } + } } else if (accumulated_token_weight > 0.0) { double global_weight = accumulated_token_weight; if (normalization_allreduce) { @@ -2321,12 +2364,36 @@ static void synchronize_lora_gradients( {static_cast(adapter_index)}); auto global_weight = global_adapter_weights.index( {static_cast(adapter_index)}); - reduce_lora_accumulator_weighted( - ctx, accum_it->second[pair][0], local_weight, - global_weight, dp_allreduce); - reduce_lora_accumulator_weighted( - ctx, accum_it->second[pair][1], local_weight, - global_weight, dp_allreduce); + TORCH_CHECK(adapter_has_global_tokens && + adapter_has_global_tokens->size() == ctx->adapters.size(), + "dynamic LoRA global-token activity is unavailable"); + if (!(*adapter_has_global_tokens)[adapter_index]) { + if (accum_it->second[pair][0].defined()) + accum_it->second[pair][0].zero_(); + if (accum_it->second[pair][1].defined()) + accum_it->second[pair][1].zero_(); + continue; + } + const bool grouped_expert = table.entries[pair].grouped_expert; + if (sharded_a2a) { + // Sharded A2A restores each source row to a numerator + // before backward. Expert owners already receive all + // source numerators; replicated projections still need + // one all-reduce before division by the global count. + normalize_lora_accumulator_numerator( + ctx, accum_it->second[pair][0], global_weight, + !grouped_expert); + normalize_lora_accumulator_numerator( + ctx, accum_it->second[pair][1], global_weight, + !grouped_expert); + } else { + reduce_lora_accumulator_weighted( + ctx, accum_it->second[pair][0], local_weight, + global_weight, dp_allreduce); + reduce_lora_accumulator_weighted( + ctx, accum_it->second[pair][1], local_weight, + global_weight, dp_allreduce); + } continue; } // Routed-expert tensors are sharded only in EP. Pure DP has @@ -2350,13 +2417,11 @@ static void synchronize_lora_gradients( } } if (per_adapter_weighting) return; - // The gated A2A prototype currently runs on the repository's replicated - // EP batch layout. Every source rank therefore sends an identical token - // batch to the owning expert. Average only the sharded expert parameter - // gradients here; scaling the combine backward would also under-scale the - // activation gradient returned to each source graph. Unequal source-token - // shards will require global-token weighting before this gate can become - // the default dispatcher. + // The replicated-source A2A path sends an identical token batch from every + // EP rank. Average only its sharded expert parameter gradients here; + // scaling the combine backward would also under-scale the activation + // gradient returned to each source graph. Sharded A2A is handled above via + // global token-count weighting and does not enter this branch. const double replicated_a2a_expert_scale = ctx->nccl_comm && env_enabled("QWEN36_EP_A2A") && !sharded_a2a && !ctx->data_parallel && ctx->ep_world_size > 1 @@ -4485,6 +4550,9 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( auto* ctx = reinterpret_cast(ctx_ptr); TORCH_CHECK(!ctx->topology_invalid, "native Qwen context rejected an incompatible TP/DP/EP topology"); + TORCH_CHECK(ctx->adapters.empty(), + "dynamic LoRA adapters require qwen36_train_multi_lora or " + "qwen36_train_multi_lora_selected"); GradientAccumulationFailureGuard accumulation_guard{ctx}; TORCH_CHECK(gradient_scale > 0.0 && std::isfinite(gradient_scale), "gradient_scale must be finite and positive"); @@ -4868,9 +4936,10 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( auto* ctx = reinterpret_cast(ctx_ptr); TORCH_CHECK(!ctx->topology_invalid, "native Qwen context rejected an incompatible TP/DP/EP topology"); - TORCH_CHECK(!env_enabled("QWEN36_EP_A2A_SHARDED"), - "dynamic multi-LoRA is not supported with sharded EP A2A; " - "source tenant metadata is not implemented"); + TORCH_CHECK(!ctx->accumulation_active && + ctx->accumulated_token_weight == 0.0, + "cannot start dynamic multi-LoRA while a fixed-LoRA gradient " + "accumulation window is pending"); GradientAccumulationFailureGuard accumulation_guard{ctx}; if (ctx->nccl_comm) { c10::cuda::set_device(ctx->cuda_device); @@ -4885,6 +4954,15 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( TORCH_CHECK(n_total > 0 && total_adapters == n_total, "n_total must equal the number of registered adapters (n_total=", n_total, ", registered=", total_adapters, ")"); + for (const auto& adapter : ctx->adapters) { + TORCH_CHECK(adapter.rank == lora_rank, + "lora_rank argument must match every registered adapter; adapter=", + adapter.id, " registered_rank=", adapter.rank, + " requested_rank=", lora_rank); + } + TORCH_CHECK(!ctx->has_mtp || env_enabled("QWEN36_DISABLE_MTP"), + "dynamic multi-LoRA with MTP is not supported until main and MTP " + "objectives have independent global token denominators"); TORCH_CHECK(input_ids.dim() == 2 && target_mask.dim() == 2, "multi-LoRA inputs must have shape [batch, seq]"); const int64_t input_batch = input_ids.size(0); @@ -4900,8 +4978,6 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( auto adapter_token_counts = input_batch == 1 ? input_row_token_counts.repeat({total_adapters}) : input_row_token_counts; - const int64_t multi_lora_invocation = ++ctx->multi_lora_invocation; - // Keep the caller's mask intact. Each chunk receives either the // corresponding rows or a repeated batch-1 mask; this also prevents // elide_trivial_attention_mask from leaking the last chunk into the @@ -4951,16 +5027,13 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( ~AdapterRegistryChunkGuard() { restore(); } }; - // Compute N_max from available GPU memory. - // CRITICAL: all workers must agree on n_max to keep NCCL all-reduce in sync. - // Use file-based barrier: rank 0 computes n_max, writes to file, others read. + // Compute N_max from available GPU memory. All workers must agree on + // the chunk schedule to keep later collectives in the same order, so + // rank 0 publishes its value through the existing communicator. size_t free_mem, total_mem; cudaMemGetInfo(&free_mem, &total_mem); int64_t n_max = 0; if (ctx->nccl_comm && ctx->ep_world_size > 1) { - const std::string sync_path = nccl_sync_dir() + "/nmax_sync_" + - std::to_string(ctx->context_sequence) + "_" + - std::to_string(multi_lora_invocation) + ".txt"; if (ctx->ep_rank == 0) { n_max = compute_n_max( (int64_t)free_mem, lora_rank, @@ -4969,33 +5042,19 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( ); n_max = std::min(n_max, total_adapters); if (n_max < 1) n_max = 1; - const std::string temporary_path = sync_path + ".tmp." + - std::to_string(static_cast(getpid())); - FILE* f = fopen(temporary_path.c_str(), "w"); - TORCH_CHECK(f, "failed to create n_max rendezvous file: ", - temporary_path); - TORCH_CHECK(fprintf(f, "%ld\n", (long)n_max) > 0, - "failed to write n_max rendezvous file: ", temporary_path); - fclose(f); - TORCH_CHECK(rename(temporary_path.c_str(), sync_path.c_str()) == 0, - "failed to publish n_max rendezvous file: ", sync_path); - } else { - bool loaded = false; - for (int i = 0; i < 600; i++) { - FILE* f = fopen(sync_path.c_str(), "r"); - if (f) { - const int parsed = fscanf(f, "%ld", (long*)&n_max); - fclose(f); - TORCH_CHECK(parsed == 1, - "invalid n_max rendezvous file: ", sync_path); - loaded = true; - break; - } - usleep(10000); - } - TORCH_CHECK(loaded, - "timed out waiting for n_max rendezvous file: ", sync_path); } + auto published_n_max = at::full( + {1}, n_max, input_ids.options().dtype(at::kLong)); + auto stream = c10::cuda::getCurrentCUDAStream( + input_ids.device().index()).stream(); + auto err = ncclBroadcast( + published_n_max.data_ptr(), + published_n_max.data_ptr(), 1, ncclInt64, 0, + reinterpret_cast(ctx->nccl_comm), stream); + TORCH_CHECK(err == ncclSuccess, "n_max broadcast failed: ", + ncclGetErrorString(err)); + n_max = published_n_max.to( + at::TensorOptions().device(at::kCPU)).item(); } else { n_max = compute_n_max( (int64_t)free_mem, lora_rank, @@ -5082,6 +5141,15 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( loss_val = loss.value.item(); hidden_grad = loss.hidden_grad; } + // independent_samples produces a local per-tenant mean. Restore + // each row to a token-sum numerator before A2A backward so the + // expert owner can combine unequal source shards correctly. + if (ctx->nccl_comm && !ctx->data_parallel && + env_enabled("QWEN36_EP_A2A_SHARDED")) { + auto chunk_token_counts = adapter_token_counts + .narrow(0, start, n).reshape({n, 1, 1}); + hidden_grad.mul_(chunk_token_counts); + } auto t_loss_end = std::chrono::steady_clock::now(); double loss_ms = std::chrono::duration(t_loss_end - t_loss_start).count(); @@ -5125,8 +5193,13 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( if (chunk == num_chunks - 1) { // DP gradient synchronization and Adam belong to the logical // multi-tenant step, never to an activation-memory chunk. + std::vector adapter_has_global_tokens; synchronize_lora_gradients( - ctx, target_mask, 0.0, &adapter_token_counts); + ctx, target_mask, 0.0, &adapter_token_counts, + &adapter_has_global_tokens); + TORCH_CHECK(adapter_has_global_tokens.size() == + ctx->adapters.size(), + "dynamic LoRA global-token activity vector mismatch"); // Adam step. Group tenants by their own logical clock so // newly-added or resumed tenants do not inherit another @@ -5136,7 +5209,10 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( ctx->lora_cache_valid = false; ctx->lora_batch_valid = false; std::map> groups; - for (auto& adapter : ctx->adapters) { + for (size_t adapter_index = 0; + adapter_index < ctx->adapters.size(); ++adapter_index) { + if (!adapter_has_global_tokens[adapter_index]) continue; + auto& adapter = ctx->adapters[adapter_index]; groups[adapter.optimizer_step + 1].push_back(&adapter); } for (auto& [logical_step, adapters] : groups) { @@ -5980,6 +6056,9 @@ __attribute__((visibility("default"))) double qwen36_eval_step(void* ctx_ptr, void* input_ids_ptr, void* target_mask_ptr, void* attention_mask_ptr) { try { auto* ctx = reinterpret_cast(ctx_ptr); + TORCH_CHECK(ctx->adapters.empty(), + "dynamic LoRA adapters require selected multi-LoRA evaluation; " + "ordinary eval_step has no tenant mapping"); auto& input_ids = *reinterpret_cast(input_ids_ptr); auto& target_mask = *reinterpret_cast(target_mask_ptr); if (attention_mask_ptr) diff --git a/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp index c5af117f..8f5e4aa1 100644 --- a/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_ep_smoke.cpp @@ -25,6 +25,17 @@ extern "C" void* qwen36_create_training_context( double, double, double, double, double, int64_t, double, int64_t, const int64_t*, int64_t, const char*); extern "C" int32_t qwen36_init_nccl(void*); +extern "C" int64_t qwen36_add_lora( + void*, int64_t, double, const int64_t*, int64_t, const char*); +extern "C" void* qwen36_get_adapter_lora_tensor( + void*, int64_t, int64_t, const char*, int32_t); +extern "C" int32_t qwen36_set_adapter_lora_tensor( + void*, int64_t, int64_t, const char*, int32_t, void*); +extern "C" void* qwen36_get_adapter_optimizer_tensor( + void*, int64_t, int64_t, const char*, int32_t, int32_t); +extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t); +extern "C" double qwen36_train_multi_lora( + void*, void*, void*, void*, int32_t, int32_t); extern "C" int64_t qwen36_get_lora_count(void*); extern "C" void* qwen36_get_lora_a(void*, int64_t); extern "C" void* qwen36_get_lora_b(void*, int64_t); @@ -113,6 +124,8 @@ int main() { std::strcmp(std::getenv("QWEN36_EP_A2A"), "0") != 0; const int sharded = a2a && std::getenv("QWEN36_EP_A2A_SHARDED") && std::strcmp(std::getenv("QWEN36_EP_A2A_SHARDED"), "0") != 0; + const bool dynamic_only = std::getenv("QWEN36_DYNAMIC_ONLY") && + std::strcmp(std::getenv("QWEN36_DYNAMIC_ONLY"), "0") != 0; // Replicated A2A receives duplicate source rows in rank order, so its BF16 // accumulation is not bit-identical to the full-expert reference. Sharded // A2A uses distinct rows; its optimizer state matches closely, while a @@ -272,15 +285,16 @@ int main() { // qwen36_init_nccl mutates only the distributed context's copied configs. assert(qwen36_init_nccl(distributed_ctx) == 0); + auto long_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kLong); + auto float_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kFloat); + auto bool_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kBool); + if (!dynamic_only) { at::Tensor input_ids; at::Tensor target_mask; at::Tensor attention_mask; at::Tensor reference_input_ids; at::Tensor reference_target_mask; at::Tensor reference_attention_mask; - auto long_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kLong); - auto float_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kFloat); - auto bool_opts = at::TensorOptions().device(at::kCUDA).dtype(at::kBool); if (sharded) { // Deliberately unequal supervised-token counts: rank 0 contributes one // response token, rank 1 contributes three. The reference evaluates @@ -416,6 +430,213 @@ int main() { assert(gate_up_b_adam_diff <= 1e-5); assert(down_a_adam_diff <= 1e-5); assert(down_b_adam_diff <= 1e-5); + } + + if (sharded) { + // Dynamic tenant rows use the source flattened token index that + // sharded A2A already transports. Rank 0 and rank 1 deliberately swap + // token counts [1,2]/[2,1] for the two adapters, exercising the + // all-reduced per-adapter denominator and owner-local expert update. + const int64_t dynamic_targets[] = {0}; + const char* dynamic_modules = + "experts_gate_up_proj,experts_down_proj"; + const int64_t adapter_a = qwen36_add_lora( + distributed_ctx, lora_rank, 1.0, dynamic_targets, 1, + dynamic_modules); + const int64_t adapter_b = qwen36_add_lora( + distributed_ctx, lora_rank, 1.0, dynamic_targets, 1, + dynamic_modules); + const int64_t reference_adapter_a = qwen36_add_lora( + reference_ctx, lora_rank, 1.0, dynamic_targets, 1, + dynamic_modules); + const int64_t reference_adapter_b = qwen36_add_lora( + reference_ctx, lora_rank, 1.0, dynamic_targets, 1, + dynamic_modules); + assert(adapter_a > 0 && adapter_b > adapter_a); + assert(reference_adapter_a > 0 && + reference_adapter_b > reference_adapter_a); + auto dynamic_tensor = [&](void* ctx, int64_t adapter, + const char* module, int b) { + auto* ptr = qwen36_get_adapter_lora_tensor( + ctx, adapter, 0, module, b); + assert(ptr); + return reinterpret_cast(ptr); + }; + auto initialize_dynamic_adapter = [&](int64_t distributed_adapter, + int64_t reference_adapter, + double offset) { + auto set_pair = [&](const char* module, int b, + std::vector shape) { + int64_t numel = 1; + for (int64_t dim : shape) numel *= dim; + auto global = ((at::arange(numel, opts) + offset) * 5e-5) + .reshape(shape).to(at::kBFloat16); + auto local = global.narrow(0, rank, 1).contiguous(); + assert(qwen36_set_adapter_lora_tensor( + distributed_ctx, distributed_adapter, 0, module, b, + &local) == 0); + assert(qwen36_set_adapter_lora_tensor( + reference_ctx, reference_adapter, 0, module, b, + &global) == 0); + }; + set_pair("experts_gate_up_proj", 0, + {experts, lora_rank, hidden}); + set_pair("experts_gate_up_proj", 1, + {experts, 2 * intermediate, lora_rank}); + set_pair("experts_down_proj", 0, + {experts, lora_rank, intermediate}); + set_pair("experts_down_proj", 1, + {experts, hidden, lora_rank}); + }; + initialize_dynamic_adapter( + adapter_a, reference_adapter_a, 11.0); + initialize_dynamic_adapter( + adapter_b, reference_adapter_b, 1011.0); + auto* dynamic_a_gate = dynamic_tensor( + distributed_ctx, adapter_a, "experts_gate_up_proj", 1); + auto* dynamic_b_gate = dynamic_tensor( + distributed_ctx, adapter_b, "experts_gate_up_proj", 1); + auto dynamic_a_before = dynamic_a_gate->clone(); + auto dynamic_b_before = dynamic_b_gate->clone(); + + auto dynamic_ids = at::tensor( + {1, 2, 3, 4, 4, 5, 6, 7}, long_opts).reshape({2, 4}); + auto dynamic_mask = rank == 0 + ? at::tensor({0, 1, 0, 0, 0, 1, 1, 0}, float_opts) + .reshape({2, 4}) + : at::tensor({0, 0, 1, 1, 0, 0, 0, 1}, float_opts) + .reshape({2, 4}); + auto dynamic_attention = at::ones({2, 4}, bool_opts); + const double dynamic_loss = qwen36_train_multi_lora( + distributed_ctx, &dynamic_ids, &dynamic_mask, + &dynamic_attention, 2, static_cast(lora_rank)); + auto reference_dynamic_mask = at::tensor( + {0, 1, 1, 1, 0, 1, 1, 1}, float_opts).reshape({2, 4}); + const double reference_dynamic_loss = dynamic_only + ? qwen36_train_multi_lora( + reference_ctx, &dynamic_ids, &reference_dynamic_mask, + &dynamic_attention, 2, static_cast(lora_rank)) + : -1.0; + c10::cuda::device_synchronize(); + assert(dynamic_loss > 0.0 && std::isfinite(dynamic_loss)); + assert(!dynamic_only || (reference_dynamic_loss > 0.0 && + std::isfinite(reference_dynamic_loss))); + assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_a) == 1); + assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_b) == 1); + const double dynamic_update_a = update_norm( + *dynamic_a_gate, dynamic_a_before); + const double dynamic_update_b = update_norm( + *dynamic_b_gate, dynamic_b_before); + assert(dynamic_update_a > 0.0 && dynamic_update_b > 0.0); + + double dynamic_param_diff = -1.0; + double dynamic_m_diff = -1.0; + double dynamic_v_diff = -1.0; + if (dynamic_only) { + const auto reference_slice = [rank](const at::Tensor& tensor) { + return tensor.narrow(0, rank, 1); + }; + auto optimizer_tensor = [](void* ctx, int64_t adapter, + const char* module, int b, int is_v) { + auto* ptr = qwen36_get_adapter_optimizer_tensor( + ctx, adapter, 0, module, b, is_v); + assert(ptr); + return reinterpret_cast(ptr); + }; + const int64_t distributed_adapters[] = {adapter_a, adapter_b}; + const int64_t reference_adapters[] = { + reference_adapter_a, reference_adapter_b}; + dynamic_param_diff = 0.0; + dynamic_m_diff = 0.0; + dynamic_v_diff = 0.0; + for (int adapter_index = 0; adapter_index < 2; ++adapter_index) { + for (const char* module : { + "experts_gate_up_proj", "experts_down_proj"}) { + auto* distributed_b = dynamic_tensor( + distributed_ctx, distributed_adapters[adapter_index], + module, 1); + auto* reference_b = dynamic_tensor( + reference_ctx, reference_adapters[adapter_index], + module, 1); + dynamic_param_diff = std::max( + dynamic_param_diff, + max_abs_diff(*distributed_b, + reference_slice(*reference_b))); + for (int is_v = 0; is_v < 2; ++is_v) { + auto* distributed_state = optimizer_tensor( + distributed_ctx, distributed_adapters[adapter_index], + module, 1, is_v); + auto* reference_state = optimizer_tensor( + reference_ctx, reference_adapters[adapter_index], + module, 1, is_v); + double diff = max_abs_diff( + *distributed_state, + reference_slice(*reference_state)); + if (is_v) dynamic_v_diff = std::max(dynamic_v_diff, diff); + else dynamic_m_diff = std::max(dynamic_m_diff, diff); + } + } + } + std::fprintf(stderr, + "native_qwen36_dynamic_reference rank=%d " + "param_diff=%0.8e m_diff=%0.8e v_diff=%0.8e\n", + rank, dynamic_param_diff, dynamic_m_diff, dynamic_v_diff); + assert(dynamic_param_diff <= 2e-3); + assert(dynamic_m_diff <= 5e-3); + assert(dynamic_v_diff <= 1e-5); + } + + // A tenant with no supervised tokens on any rank is a valid no-op: + // it must not abort the logical step or advance its private Adam + // clock. This also covers the final chunk when n_max is smaller than + // the registry size. + const int64_t adapter_c = qwen36_add_lora( + distributed_ctx, lora_rank, 1.0, dynamic_targets, 1, + dynamic_modules); + assert(adapter_c > adapter_b); + auto* dynamic_c_gate = dynamic_tensor( + distributed_ctx, adapter_c, "experts_gate_up_proj", 1); + auto dynamic_c_before = dynamic_c_gate->clone(); + auto zero_ids = at::tensor({ + 1, 2, 3, 4, 4, 5, 6, 7, 1, 2, 3, 4}, long_opts) + .reshape({3, 4}); + auto zero_mask = at::tensor({ + 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0}, float_opts) + .reshape({3, 4}); + auto zero_attention = at::ones({3, 4}, bool_opts); + const double zero_loss = qwen36_train_multi_lora( + distributed_ctx, &zero_ids, &zero_mask, + &zero_attention, 3, static_cast(lora_rank)); + assert(zero_loss > 0.0 && std::isfinite(zero_loss)); + assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_a) == 2); + assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_b) == 2); + assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_c) == 0); + assert(update_norm(*dynamic_c_gate, dynamic_c_before) == 0.0); + + // Ordinary single-adapter entry points must not silently pick an + // arbitrary tenant or crash in a batched BMM after registration. + auto ordinary_ids = dynamic_ids.narrow(0, 0, 1).contiguous(); + auto ordinary_mask = dynamic_mask.narrow(0, 0, 1).contiguous(); + auto ordinary_attention = dynamic_attention.narrow(0, 0, 1).contiguous(); + assert(qwen36_train_step( + distributed_ctx, &ordinary_ids, &ordinary_mask, + &ordinary_attention) < 0.0); + assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_a) == 2); + assert(qwen36_get_adapter_step_count(distributed_ctx, adapter_b) == 2); + std::printf( + "native_qwen36_dynamic_sharded rank=%d world=%d loss=%0.8f " + "adapter_steps=[%ld,%ld,%ld] updates=[%0.8e,%0.8e] " + "zero_loss=%0.8f reference_loss=%0.8f " + "param_diff=%0.8e m_diff=%0.8e v_diff=%0.8e\n", + rank, world, dynamic_loss, + (long)qwen36_get_adapter_step_count(distributed_ctx, adapter_a), + (long)qwen36_get_adapter_step_count(distributed_ctx, adapter_b), + (long)qwen36_get_adapter_step_count(distributed_ctx, adapter_c), + dynamic_update_a, dynamic_update_b, zero_loss, + reference_dynamic_loss, dynamic_param_diff, + dynamic_m_diff, dynamic_v_diff); + std::fflush(stdout); + } qwen36_free_training_context(reference_ctx); qwen36_free_training_context(distributed_ctx); diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index 4a363df1..f4deca1f 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -678,16 +678,30 @@ int main() { ctx, adapter_one, 0, "shared_gate_proj", 1) != nullptr); assert(qwen36_get_adapter_lora_tensor( ctx, adapter_two, 0, "shared_gate_proj", 1) != nullptr); - // A tenant may be empty on one DP rank, but it must have at least one - // target token globally. In this single-rank check the second tenant is - // globally empty, so reject the step without advancing either clock. - auto invalid_multi_target_mask = multi_target_mask.clone(); - invalid_multi_target_mask.select(0, 1).zero_(); + // A globally empty tenant is a successful no-op: active tenants still + // update, while the empty tenant preserves parameters, optimizer state, + // and its private Adam clock. + auto zero_target_mask = multi_target_mask.clone(); + zero_target_mask.select(0, 1).zero_(); + auto empty_tenant_b_before = dynamic_b_two->clone(); + auto* empty_tenant_m = reinterpret_cast( + qwen36_get_adapter_optimizer_tensor( + ctx, adapter_two, 0, "shared_gate_proj", 1, 0)); + auto* empty_tenant_v = reinterpret_cast( + qwen36_get_adapter_optimizer_tensor( + ctx, adapter_two, 0, "shared_gate_proj", 1, 1)); + assert(empty_tenant_m && empty_tenant_v); + auto empty_tenant_m_before = empty_tenant_m->clone(); + auto empty_tenant_v_before = empty_tenant_v->clone(); assert(qwen36_train_multi_lora( - ctx, &multi_input_ids, &invalid_multi_target_mask, &multi_attention_mask, - 2, rank) < 0.0); - assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 2); + ctx, &multi_input_ids, &zero_target_mask, &multi_attention_mask, + 2, rank) > 0.0); + c10::cuda::device_synchronize(); + assert(qwen36_get_adapter_step_count(ctx, adapter_one) == 3); assert(qwen36_get_adapter_step_count(ctx, adapter_two) == 1); + assert((*dynamic_b_two - empty_tenant_b_before).abs().max().item() == 0.0); + assert((*empty_tenant_m - empty_tenant_m_before).abs().max().item() == 0.0); + assert((*empty_tenant_v - empty_tenant_v_before).abs().max().item() == 0.0); qwen36_free_training_context(ctx); // Dense Qwen3.5 variants use the same per-sample activation path for diff --git a/docs/plans/qwen-lora-megatron-progress.md b/docs/plans/qwen-lora-megatron-progress.md index 0acc20b3..b8ba4786 100644 --- a/docs/plans/qwen-lora-megatron-progress.md +++ b/docs/plans/qwen-lora-megatron-progress.md @@ -14,7 +14,7 @@ timestamp: 2026-07-17T00:00:00Z Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP parity against a full-expert reference, variable-split EP A2A with fixed-LoRA data sharding, dense replicated-DP smoke with per-tenant token weighting, TP-only latent-rank-sharded LoRA smoke, dynamic batch logical-step update, ABI11 FP32 gradient storage/aggregation, standard Adam bias correction, per-tenant optimizer-step restore, selected-tenant isolation, same-topology rank-aware checkpointing, and 5D topology mapping. -Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axis TP+DP/EP, PP/CP, DeepEP/TE prebuilt integration, dynamic multi-LoRA source metadata under sharded A2A, cross-topology checkpoint resharding, and matched Megatron throughput. +Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axis TP+DP/EP, PP/CP, DeepEP/TE prebuilt integration, server-side source-sharded dynamic dispatch, heterogeneous dynamic adapter signatures, dynamic+MTP objective normalization, cross-topology checkpoint resharding, and matched Megatron throughput. Native direct dynamic source metadata under sharded A2A is now implemented and full-reference smoke-tested. # Durable Milestones @@ -28,7 +28,7 @@ Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axi - `76eace6`: direct native context creation also rejects PP/CP sizes greater than one, so unsupported pipeline/context parallelism cannot silently fall back to replicated single-stage execution. - `b7050ea`: fixed standard Adam bias correction, separated pure-DP gradient NCCL from MoE activation collectives, made dynamic-MoE DP token weighting numerical, scoped n_max rendezvous by invocation with atomic publication, and added full-expert EP parity plus Adam oracles. - `bfa3d8e`: added a gated variable-split NCCL dispatch/inverse-combine prototype with differentiable forward/backward collectives; default legacy EP remains unchanged. -- Working tree after `bfa3d8e`: `QWEN36_EP_A2A_SHARDED=1` disjoins EP input rows, all-reduces global fixed-LoRA token numerators, keeps grouped expert LoRA local after A2A, and explicitly rejects dynamic multi-LoRA until source tenant metadata exists. +- Working tree after `bfa3d8e`: `QWEN36_EP_A2A_SHARDED=1` disjoins EP input rows, all-reduces global fixed-LoRA token numerators, and keeps grouped expert LoRA local after A2A. Native dynamic multi-LoRA now reuses the transported source row for tenant mapping, all-reduces per-tenant counts, skips zero-global-token tenants, and rejects ordinary single-adapter entry points while a dynamic registry is active. - H20 `123.57.26.97:28004`: ABI8 native smoke passed grouped/fallback parity, GDN, dense/MoE LoRA, dynamic adapters, and step setter validation. - H20 `123.57.26.97:28004`: ABI9 native smoke passed selected-tenant training with a positive selected update, exactly zero unselected update, independent clocks (`2` vs `1`), and registry preservation after an unknown ID. - H20 `123.57.26.97:28004`: ABI10 two-rank TP native smoke passed on both ranks for MoE, dense MLP, GDN/linear attention, dynamic multi-LoRA, and selected-tenant isolation. Rank-local LoRA tensors used distinct rank `4` slices for global rank `8`; losses matched and both shards had positive updates. @@ -38,16 +38,17 @@ Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axi - H20 `123.57.26.97:28004`: gated A2A EP2 smoke passed in legacy, replicated-source, and sharded-source modes (`rank_statuses=0,0`). Sharded mode used token counts `[1,3]`; weighted global loss differed from the full-batch reference by `9.6e-7`, expert m/v maxima were `1.22e-5` / `3.92e-9`, and every standard Adam oracle was zero. One distributed BF16 parameter landed in the adjacent rounding bin (`9.61e-4`), bounded separately by the smoke. - Target runtime probe: PyTorch 2.5.1+cu121, ABI0; Transformer Engine, flash-attn, DeepEP, Triton, and DeepSpeed are not importable. - H20 `123.57.26.97:28004`: `native_ep_bench.cpp` fresh ABI0 benchmark (`seq=128, hidden=256, experts=8, intermediate=256, warmup=2, iters=10`) passed legacy and sharded A2A with `rank_statuses=0,0`. Legacy median was about `5.47 ms` / `46.4k processed tokens/s` (`23.2k unique tokens/s`), while sharded A2A was about `6.72 ms` / `37.8k processed and unique tokens/s`. This is a synthetic native baseline, not Megatron-LM parity. +- H20 `123.57.26.97:28004`: fresh ABI1 dynamic sharded native smoke passed on both ranks (`rank_statuses=0,0`) against a full-expert reference with complementary source masks. Dynamic grouped-expert parameter/m/v maxima were `1.53e-5` / `4.88e-5` / `5.75e-8`; it also exercised a third tenant with zero global target tokens, clocks `[2,2,0]`, and explicit rejection of ordinary `train_step`. The server path still broadcasts replicated source batches, and dynamic+MTP is explicitly rejected until its two objective denominators are separated. # Decisions During Execution - Keep TP and EP communicators separate; do not reuse the existing EP `LayerConfig.nccl_comm` for LoRA TP deltas. -- Scope multi-LoRA `n_max` rendezvous files per native context; reusing one filename across sessions can give ranks different chunk schedules and deadlock TP collectives. +- Publish multi-LoRA `n_max` directly from rank 0 with `ncclBroadcast`; filesystem rendezvous can reuse stale files across process restarts and give ranks different chunk schedules. - Do not enable Qwen TP/PP/CP by merely relaxing runtime validation. - Treat Exa/Jina dependency search failures as missing evidence, not as proof that a package is compatible. # Verification -Passed: `cargo check -p rustrain-qwen3-6 -p rustrain-server -p rustrain-ipc` (with the repository PyTorch 2.12.1 host venv), `cargo test -p rustrain-server --lib` (6), Qwen unit tests (3), Qwen integration (6), remote ABI8 smoke, remote ABI9 selected-tenant native smoke, remote ABI10 single-rank native smoke, remote ABI10 two-rank TP native smoke, and remote ABI11 single/TP2/DP2/EP2 native smoke with numerical Adam and parity oracles. +Passed: `cargo check -p rustrain-qwen3-6 -p rustrain-server -p rustrain-ipc` (with the repository PyTorch 2.12.1 host venv), `cargo test -p rustrain-qwen3-6 --lib` (3), `cargo test -p rustrain-server --lib` (6), Qwen integration tests, remote ABI8 smoke, remote ABI9 selected-tenant native smoke, remote ABI10 single-rank native smoke, remote ABI10 two-rank TP native smoke, remote ABI11 single/TP2/DP2/EP2 native smoke with numerical Adam and parity oracles, and the ABI1 dynamic sharded full-reference smoke described above. -Not run: Megatron-style base-model TP, multi-axis TP+DP/EP, PP/CP, dynamic multi-LoRA source metadata under sharded A2A, cross-topology resharding, and matched Megatron performance benchmark. The target host lacks importable Megatron/Transformer Engine/DeepEP/flash-attn prebuilt packages, so no dependency installation or JIT workaround was used. +Not run: Megatron-style base-model TP, multi-axis TP+DP/EP, PP/CP, server-side source-sharded dynamic dispatch, heterogeneous dynamic adapter signatures, dynamic+MTP, cross-topology resharding, and matched Megatron performance benchmark. The target host lacks importable Megatron/Transformer Engine/DeepEP/flash-attn prebuilt packages, so no dependency installation or JIT workaround was used. diff --git a/docs/qwen35-qwen36-megatron-audit.md b/docs/qwen35-qwen36-megatron-audit.md index 20d8ecb0..3b4c5ff5 100644 --- a/docs/qwen35-qwen36-megatron-audit.md +++ b/docs/qwen35-qwen36-megatron-audit.md @@ -5,7 +5,7 @@ ## 结论 - 模型语义：Qwen3.5 dense、Qwen3.6 dense/MoE 的 native forward/backward 路径已经覆盖 hybrid full attention、GDN、MoE、MTP 和 LoRA 目标模块；已有配置解析、集成测试及 H20 native smoke 证据。 -- 已实现并可验证的分布式子集：MoE expert parallel，以及 replicated LoRA 的 data parallel；梯度累积和 dynamic multi-LoRA 已有 logical-step 边界。DP 动态租户按 adapter token count 加权，纯 DP 不把 world communicator 传入 MoE activation reduction。 +- 已实现并可验证的分布式子集：MoE expert parallel，以及 replicated LoRA 的 data parallel；梯度累积和 dynamic multi-LoRA 已有 logical-step 边界。DP 动态租户按 adapter token count 加权，sharded A2A native 路径会保留 source flattened row 来恢复租户，并按全局租户 token count 归一化。 - 性能：MoE grouped dispatch 相对逐 expert matmul 的已有 microbenchmark 为约 3.70x（E=32, N=4096, H=2048, I=768，结果误差为 0）；这不是端到端训练吞吐或 Megatron 对比。 - 已实现 LoRA latent-rank 的 TP-only 子集，但尚未实现 frozen base-weight tensor parallel、pipeline parallel、context parallel，以及 TP/PP/CP 与 EP/DP 的组合。当前训练上下文仍由单个进程持有完整 dense 权重和完整层栈。 - 因此当前实现不能宣称“Megatron-LM 级别”。它是一个计算集中在 C++ 的 LoRA/EP/DP 子集，离 Megatron 的完整并行和通信重叠仍有实质差距。 @@ -19,10 +19,10 @@ | Qwen3.6 MoE | 已实现 | grouped dispatch、EP smoke；完整模型仍需目标 GPU/权重运行 | | MTP | 已实现 | C++ hidden gradient 检查和集成测试；可通过环境变量关闭 | | fixed LoRA | 已实现 | attention/GDN/MLP/shared/routed expert 目标模块 | -| dynamic multi-LoRA | 已实现子集 | 请求按 adapter 分组，单个 logical step 统一 backward/Adam；每个 adapter 独立 optimizer clock 和 m/v，DP 按 token count 加权 | +| dynamic multi-LoRA | 已实现子集 | 请求按 adapter 分组，单个 logical step 统一 backward/Adam；每个 adapter 独立 optimizer clock 和 m/v，DP 与 native sharded A2A 按全局 token count 加权；零全局 token 租户跳过更新；dynamic+MTP 暂拒绝 | | microbatch accumulation | 已实现子集 | non-final microbatch 只 backward，final microbatch 才 optimizer；FP32 accumulator 存储/聚合，autograd leaf backward 仍为 BF16 | | replicated data parallel | 已实现 | logical-step 边界同步 replicated LoRA；EP expert 参数不走该 reduction | -| expert parallel | 已实现子集 | 默认 routed-output all-reduce；gated variable-split A2A 已验证 fixed-LoRA data sharding，但 dynamic tenant metadata、GPU-only split planning 和 overlap 未实现 | +| expert parallel | 已实现子集 | 默认 routed-output all-reduce；gated variable-split A2A 已验证 fixed-LoRA 和 native dynamic-LoRA data sharding；GPU-only split planning、异步 overlap 和 DeepEP backend 未实现 | | tensor parallel | LoRA-only 子集 | latent rank 分片和独立 TP communicator 已验证；attention/MLP/LM-head base 权重仍不切分 | | pipeline parallel | 未实现于 Qwen native | 没有 stage 切分、microbatch scheduler 或 activation send/recv | | context parallel | 未实现于 Qwen native | 没有 ring attention、跨 rank KV/索引合并 | @@ -34,7 +34,7 @@ Megatron 的 TP 会按 head、hidden/intermediate 和 vocab 维度切分权重，并在线性层边界执行必要的 reduce-scatter/all-reduce；PP 会把层分到不同 stage 并使用 1F1B 等调度；CP 会在序列和 attention state 上做跨 rank 通信。当前 Qwen native `TrainingContext` 仍加载完整模型并在一个 C++ forward 中执行全部层，因此仅增加 `tensor_model_parallel_size` 等配置不能得到正确的 TP/PP/CP。 -当前 DP/EP 也不是完整 Megatron 语义：DP 同步 replicated LoRA 梯度并按租户 token count 归一化，expert 参数留在 EP rank；EP 默认使用 routed-output all-reduce，实验 gate 已加入 variable-split dispatch/inverse combine 和 fixed-LoRA data sharding。该 gate 仍逐 top-k 执行 host-visible count sync，没有 fused permutation、dynamic tenant source metadata、异步 overlap 或 DeepEP backend。 +当前 DP/EP 也不是完整 Megatron 语义：DP 同步 replicated LoRA 梯度并按租户 token count 归一化，expert 参数留在 EP rank；EP 默认使用 routed-output all-reduce，实验 gate 已加入 variable-split dispatch/inverse combine 和 fixed/dynamic LoRA data sharding。dynamic native path 的 source row metadata 已随 token index 传输，但仍逐 top-k 执行 host-visible count sync，没有 fused permutation、GPU-only split planning、异步 overlap 或 DeepEP backend。server 的 `TrainMultiLora` 目前向各 worker 广播相同 batch，因此不能把它当作 source-sharded 服务吞吐。 ### 优化器与恢复 @@ -50,7 +50,7 @@ Megatron 的 TP 会按 head、hidden/intermediate 和 vocab 维度切分权重 ## 验证边界 -已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试，以及 H20 ABI11 的单卡、TP2、EP2 和 DP2 native smoke。DP2 的 weighted m/grouped/v/Adam BF16 delta oracle 分别达到 `2.43e-8`、`2.27e-8`、`7.33e-8` 和 `0`；legacy EP2 与 full-expert reference 的 loss、LoRA、Adam state 和标准 Adam 首步 oracle 在两 rank 均为零差异。Replicated A2A 与 sharded A2A 也均通过两 rank full-expert reference；sharded token counts `[1,3]` 的加权 loss 与 global reference 相差约 `9.6e-7`，m/v 最大差 `1.22e-5` / `3.92e-9`，Adam oracle 差为 `0`。没有完成 Qwen3.5/3.6 完整大模型的长时间训练、跨节点通信、PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖且有直接 oracle 的子集，而不是所有并行配置。 +已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试，以及 H20 ABI11 的单卡、TP2、EP2 和 DP2 native smoke。DP2 的 weighted m/grouped/v/Adam BF16 delta oracle 分别达到 `2.43e-8`、`2.27e-8`、`7.33e-8` 和 `0`；legacy EP2 与 full-expert reference 的 loss、LoRA、Adam state 和标准 Adam 首步 oracle 在两 rank 均为零差异。Replicated A2A 与 fixed-LoRA sharded A2A 也均通过两 rank full-expert reference；sharded token counts `[1,3]` 的加权 loss 与 global reference 相差约 `9.6e-7`，m/v 最大差 `1.22e-5` / `3.92e-9`，Adam oracle 差为 `0`。新增 H20 ABI1 dynamic sharded full-reference smoke 在两 rank 返回 `0`：互补 source masks 的 dynamic grouped-expert 参数最大差 `1.53e-5`，m/v 最大差 `4.88e-5` / `5.75e-8`；两次有效租户更新后 step 为 `[2,2]`，全局零 token 租户保持 step `0` 且参数无更新，普通 `train_step` 被显式拒绝。没有完成 Qwen3.5/3.6 完整大模型的长时间训练、跨节点通信、PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖且有直接 oracle 的子集，而不是所有并行配置。 ## 继续达到 Megatron 级别所需的最小工作包 From 54bcce5a2e8729c935d97fb4f2ce5391faa58ec4 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 09:56:56 +0800 Subject: [PATCH 024/156] feat: add frozen dense MLP tensor parallelism --- .../kernels/qwen3_6_kernels.cpp | 230 ++++++++++++++++-- crates/rustrain-qwen3-6/src/kernel.rs | 90 ++++++- crates/rustrain-qwen3-6/src/session.rs | 71 +++++- .../rustrain-qwen3-6/tests/native_smoke.cpp | 2 +- .../tests/native_tp_mlp_smoke.cpp | 225 +++++++++++++++++ crates/rustrain-server/src/session.rs | 55 ++++- docs/plans/qwen-lora-megatron-progress.md | 10 +- docs/qwen35-qwen36-megatron-audit.md | 12 +- 8 files changed, 644 insertions(+), 51 deletions(-) create mode 100644 crates/rustrain-qwen3-6/tests/native_tp_mlp_smoke.cpp diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index e3fecaba..b078f24f 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -194,6 +194,28 @@ struct NcclAllReduceFunction : public torch::autograd::Function { + static at::Tensor forward(torch::autograd::AutogradContext* ctx, + at::Tensor input, int64_t comm_ptr, int64_t stream_ptr) { + ctx->saved_data["comm"] = comm_ptr; + ctx->saved_data["stream"] = stream_ptr; + return input; + } + + static std::vector backward( + torch::autograd::AutogradContext* ctx, + std::vector grad_output) { + auto comm = reinterpret_cast(ctx->saved_data["comm"].toInt()); + auto stream = reinterpret_cast(ctx->saved_data["stream"].toInt()); + auto grad_input = NcclAllReduceFunction::allreduce( + grad_output[0], comm, stream); + return {grad_input, at::Tensor(), at::Tensor()}; + } +}; + struct FusedSwiGLUFunction : public torch::autograd::Function { static at::Tensor forward(torch::autograd::AutogradContext* ctx, at::Tensor gate, at::Tensor up, double limit) { @@ -1626,6 +1648,13 @@ static inline int64_t weight_count_for_layer(const LayerConfig& cfg) { enum class LoraSegment : uint8_t { Attention, Mlp }; +static bool is_mlp_lora_target(const std::string& name) { + return name == "gate_proj" || name == "up_proj" || name == "down_proj" || + name == "shared_gate_proj" || name == "shared_up_proj" || + name == "shared_down_proj" || name == "experts_gate_up_proj" || + name == "experts_down_proj"; +} + struct LoraProjectionSpec { const char* name; int64_t weight_index; @@ -1692,6 +1721,10 @@ static int64_t lora_pair_index(const LayerConfig& cfg, const char* name) { static RoutedExpertLora routed_expert_lora( TrainingContext* ctx, int64_t layer_idx, const LayerConfig& cfg); +static at::Tensor tp_allreduce_base_mlp( + TrainingContext* ctx, const at::Tensor& local_output); +static at::Tensor tp_copy_base_mlp_input( + TrainingContext* ctx, const at::Tensor& input); static at::Tensor forward_single_layer( TrainingContext* ctx, const at::Tensor& hidden, at::Tensor** w, const LayerConfig* cfg, @@ -1763,7 +1796,9 @@ static at::Tensor forward_single_layer( lora_pair_index(*cfg, "up_proj"), *w[9]); auto down = apply_multi_lora(ctx, layer_idx, lora_pair_index(*cfg, "down_proj"), *w[10]); - auto mlp_out = dense_mlp_forward(post_attn, gate, up, down, kind); + auto mlp_input = tp_copy_base_mlp_input(ctx, post_attn); + auto mlp_out = tp_allreduce_base_mlp( + ctx, dense_mlp_forward(mlp_input, gate, up, down, kind)); return hidden + attn_output + mlp_out; } } else { @@ -1825,7 +1860,9 @@ static at::Tensor forward_single_layer( lora_pair_index(*cfg, "up_proj"), *w[12]); auto down = apply_multi_lora(ctx, layer_idx, lora_pair_index(*cfg, "down_proj"), *w[13]); - auto mlp_out = dense_mlp_forward(post_attn, gate, up, down, kind); + auto mlp_input = tp_copy_base_mlp_input(ctx, post_attn); + auto mlp_out = tp_allreduce_base_mlp( + ctx, dense_mlp_forward(mlp_input, gate, up, down, kind)); return hidden + attn_output + mlp_out; } } @@ -1964,6 +2001,10 @@ struct TrainingContext { cudaStream_t tp_stream = nullptr; int tp_world_size = 1; int tp_rank = 0; + // Frozen dense SwiGLU TP: gate/up are output-sharded and down is + // input-sharded by the Rust weight loader. The local row contribution + // is reduced over the TP communicator before the residual add. + bool base_tp_mlp = false; // Set when a legacy NCCL setter supplies an incompatible mixed topology. // Training entry points reject the context before touching parameters. bool topology_invalid = false; @@ -2096,6 +2137,32 @@ static at::Tensor tp_allreduce_lora_delta( (int64_t)reinterpret_cast(ctx->tp_stream)); } +static at::Tensor tp_allreduce_base_mlp( + TrainingContext* ctx, const at::Tensor& local_output +) { + if (!ctx || !ctx->base_tp_mlp || ctx->tp_world_size <= 1) + return local_output; + TORCH_CHECK(ctx->tp_comm, + "base MLP TP communicator is not initialized for TP_SIZE=", + ctx->tp_world_size); + return NcclAllReduceFunction::apply( + local_output, (int64_t)ctx->tp_comm, + (int64_t)reinterpret_cast(ctx->tp_stream)); +} + +static at::Tensor tp_copy_base_mlp_input( + TrainingContext* ctx, const at::Tensor& input +) { + if (!ctx || !ctx->base_tp_mlp || ctx->tp_world_size <= 1) + return input; + TORCH_CHECK(ctx->tp_comm, + "base MLP TP communicator is not initialized for TP_SIZE=", + ctx->tp_world_size); + return TpCopyToRegionFunction::apply( + input, (int64_t)ctx->tp_comm, + (int64_t)reinterpret_cast(ctx->tp_stream)); +} + static at::Tensor initialize_lora_a( TrainingContext* ctx, const at::TensorOptions& options, int64_t experts, int64_t global_rank, int64_t in_features @@ -2709,17 +2776,20 @@ static at::Tensor dense_mlp_forward_batched( const int64_t up_pair = lora_pair_index(cfg, "up_proj"); const int64_t down_pair = lora_pair_index(cfg, "down_proj"); - auto gate_out = at::matmul(hidden, gate_proj.t()); - auto up_out = at::matmul(hidden, up_proj.t()); + auto mlp_input = tp_copy_base_mlp_input(ctx, hidden); + auto gate_out = at::matmul(mlp_input, gate_proj.t()); + auto up_out = at::matmul(mlp_input, up_proj.t()); gate_out = add_batched_lora( - ctx, gate_out, hidden, lora_batch_entry(ctx, layer_idx, gate_pair)); + ctx, gate_out, mlp_input, lora_batch_entry(ctx, layer_idx, gate_pair)); up_out = add_batched_lora( - ctx, up_out, hidden, lora_batch_entry(ctx, layer_idx, up_pair)); + ctx, up_out, mlp_input, lora_batch_entry(ctx, layer_idx, up_pair)); auto activated = fused_swiglu_op(gate_out, up_out, 0.0); auto result = at::matmul(activated, down_proj.t()); result = add_batched_lora( ctx, result, activated, lora_batch_entry(ctx, layer_idx, down_pair)); - return result.to(compute_type); + // Base TP uses a row-parallel down projection. Reduce its local hidden + // contribution before the residual add. + return tp_allreduce_base_mlp(ctx, result.to(compute_type)); } __attribute__((noinline, visibility("default"))) @@ -2869,7 +2939,9 @@ at::Tensor compute_mlp_only( lora_pair_index(cfg, "up_proj"), *ctx->weight_ptrs[w_offset+mlp_start+1]); auto down = apply_multi_lora(ctx, layer_idx, lora_pair_index(cfg, "down_proj"), *ctx->weight_ptrs[w_offset+mlp_start+2]); - return dense_mlp_forward(post_attn, gate, up, down, kind); + auto mlp_input = tp_copy_base_mlp_input(ctx, post_attn); + return tp_allreduce_base_mlp( + ctx, dense_mlp_forward(mlp_input, gate, up, down, kind)); } } @@ -4295,7 +4367,7 @@ static at::Tensor mtp_compute_loss( extern "C" { __attribute__((visibility("default"))) int64_t qwen36_kernel_abi_version() { - return 11; + return 13; } // Create training context — called once at startup @@ -4321,6 +4393,7 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( ctx->vocab_size = vocab_size; ctx->rms_eps = rms_eps; ctx->step_count = 0; ctx->lora_scaling = lora_scaling; ctx->num_layers = num_layers; + ctx->has_mtp = false; ctx->use_checkpoint = false; ctx->group_size = 4; const char* tp_size_env = getenv("TP_SIZE"); if (!tp_size_env) tp_size_env = getenv("RUSTRAIN_TP_SIZE"); @@ -4504,9 +4577,83 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( } } +__attribute__((visibility("default"))) int32_t qwen36_set_base_tp_mlp( + void* ctx_ptr, int32_t enabled +) { + try { + auto* ctx = reinterpret_cast(ctx_ptr); + TORCH_CHECK(ctx, "null training context"); + if (!enabled) { + TORCH_CHECK(!ctx->base_tp_mlp, + "base dense MLP TP cannot be disabled after enablement because " + "the context owns TP-sharded weights"); + return 0; + } + if (ctx->base_tp_mlp) return 0; + TORCH_CHECK(ctx->tp_world_size > 1, + "base dense MLP TP requires TP_SIZE>1"); + TORCH_CHECK(!ctx->has_mtp, + "base dense MLP TP requires MTP to be disabled"); + for (const auto& adapter : ctx->adapters) { + TORCH_CHECK(!adapter.target_modules.empty(), + "base dense MLP TP does not support an existing dynamic LoRA " + "adapter targeting all modules"); + for (const auto& name : adapter.target_modules) { + TORCH_CHECK(!is_mlp_lora_target(name), + "base dense MLP TP does not support existing dynamic MLP LoRA target ", + name, "; use attention-only targets"); + } + } + + int64_t weight_offset = 0; + for (int64_t layer = 0; layer < ctx->num_layers; ++layer) { + const auto& cfg = ctx->layer_configs[layer]; + TORCH_CHECK(cfg.num_experts == 0, + "base dense MLP TP currently supports dense models only"); + TORCH_CHECK(cfg.intermediate_size > 0 && + cfg.intermediate_size % ctx->tp_world_size == 0, + "dense intermediate_size must be divisible by TP_SIZE"); + const int64_t local_intermediate = + cfg.intermediate_size / ctx->tp_world_size; + const int64_t mlp_start = cfg.layer_type == 0 ? 8 : 11; + auto* gate = ctx->weight_ptrs[weight_offset + mlp_start]; + auto* up = ctx->weight_ptrs[weight_offset + mlp_start + 1]; + auto* down = ctx->weight_ptrs[weight_offset + mlp_start + 2]; + TORCH_CHECK(gate && up && down && + gate->dim() == 2 && up->dim() == 2 && down->dim() == 2, + "base dense MLP TP requires matrix gate/up/down weights"); + TORCH_CHECK(gate->size(0) == local_intermediate && + up->size(0) == local_intermediate && + down->size(1) == local_intermediate && + gate->size(1) == up->size(1) && + gate->size(1) == down->size(0), + "base dense MLP TP received inconsistent local weight shapes: gate=", + gate->sizes(), " up=", up->sizes(), " down=", down->sizes(), + " expected local intermediate=", local_intermediate); + + const int64_t lora_offset = ctx->lora_layer_offset[layer]; + auto projections = lora_projection_table(cfg); + for (int64_t pair = 0; pair < projections.count; ++pair) { + if (projections.entries[pair].segment == LoraSegment::Mlp) { + TORCH_CHECK(!ctx->lora_active[lora_offset + pair], + "base dense MLP TP does not yet support MLP LoRA target ", + projections.entries[pair].name, + "; use explicit attention-only targets"); + } + } + weight_offset += weight_count_for_layer(cfg); + } + ctx->base_tp_mlp = true; + return 0; + } catch (const std::exception& e) { + fprintf(stderr, "[q36] set_base_tp_mlp FAILED: %s\n", e.what()); + return -1; + } +} + // Set MTP weights on an existing training context. // Called after create_training_context if MTP is enabled. -__attribute__((visibility("default"))) void qwen36_set_mtp_weights( +__attribute__((visibility("default"))) int32_t qwen36_set_mtp_weights( void* ctx_ptr, void* mtp_fc_ptr, void* mtp_pre_fc_norm_emb_ptr, @@ -4515,25 +4662,35 @@ __attribute__((visibility("default"))) void qwen36_set_mtp_weights( void** mtp_layer_weight_ptrs, int64_t num_mtp_layer_weights, void* mtp_layer_configs_ptr, int64_t num_mtp_layers ) { - auto* ctx = reinterpret_cast(ctx_ptr); - ctx->has_mtp = true; - ctx->mtp_fc = reinterpret_cast(mtp_fc_ptr); - ctx->mtp_pre_fc_norm_emb = reinterpret_cast(mtp_pre_fc_norm_emb_ptr); - ctx->mtp_pre_fc_norm_hidden = reinterpret_cast(mtp_pre_fc_norm_hidden_ptr); - ctx->mtp_norm = reinterpret_cast(mtp_norm_ptr); + try { + auto* ctx = reinterpret_cast(ctx_ptr); + TORCH_CHECK(ctx, "null training context"); + TORCH_CHECK(!ctx->base_tp_mlp, + "MTP cannot be enabled after base dense MLP TP because MTP weights are not sharded"); + TORCH_CHECK(!ctx->has_mtp, "MTP weights are already configured"); + ctx->has_mtp = true; + ctx->mtp_fc = reinterpret_cast(mtp_fc_ptr); + ctx->mtp_pre_fc_norm_emb = reinterpret_cast(mtp_pre_fc_norm_emb_ptr); + ctx->mtp_pre_fc_norm_hidden = reinterpret_cast(mtp_pre_fc_norm_hidden_ptr); + ctx->mtp_norm = reinterpret_cast(mtp_norm_ptr); + + auto** wp = reinterpret_cast(mtp_layer_weight_ptrs); + for (int64_t i = 0; i < num_mtp_layer_weights; i++) { + ctx->mtp_layer_weights.push_back(wp[i]); + } - auto** wp = reinterpret_cast(mtp_layer_weight_ptrs); - for (int64_t i = 0; i < num_mtp_layer_weights; i++) { - ctx->mtp_layer_weights.push_back(wp[i]); - } + auto* lcfgs = reinterpret_cast(mtp_layer_configs_ptr); + for (int64_t i = 0; i < num_mtp_layers; i++) { + ctx->mtp_layer_configs.push_back(lcfgs[i]); + } - auto* lcfgs = reinterpret_cast(mtp_layer_configs_ptr); - for (int64_t i = 0; i < num_mtp_layers; i++) { - ctx->mtp_layer_configs.push_back(lcfgs[i]); + fprintf(stderr, "[q36_ctx] MTP set: %ld MTP layers, %ld MTP weight pointers\n", + (long)num_mtp_layers, (long)num_mtp_layer_weights); + return 0; + } catch (const std::exception& e) { + fprintf(stderr, "[q36] set_mtp_weights FAILED: %s\n", e.what()); + return -1; } - - fprintf(stderr, "[q36_ctx] MTP set: %ld MTP layers, %ld MTP weight pointers\n", - (long)num_mtp_layers, (long)num_mtp_layer_weights); } // One training micro-step. Non-final micro-steps accumulate scaled leaf @@ -5746,6 +5903,16 @@ int64_t qwen36_add_lora( while (std::getline(ss, item, ',')) adapter.target_modules.insert(item); } + if (ctx->base_tp_mlp) { + TORCH_CHECK(!adapter.target_modules.empty(), + "base MLP tensor parallelism does not support a dynamic LoRA " + "adapter targeting all modules; use explicit attention-only targets"); + for (const auto& name : adapter.target_modules) { + TORCH_CHECK(!is_mlp_lora_target(name), + "base MLP tensor parallelism does not yet support dynamic MLP LoRA target ", name, + "; use attention-only targets until projection-axis LoRA collectives are implemented"); + } + } for (auto layer : adapter.target_layers) { TORCH_CHECK(layer >= 0 && layer < ctx->num_layers, "dynamic LoRA target layer out of range: ", layer, @@ -6056,6 +6223,17 @@ __attribute__((visibility("default"))) double qwen36_eval_step(void* ctx_ptr, void* input_ids_ptr, void* target_mask_ptr, void* attention_mask_ptr) { try { auto* ctx = reinterpret_cast(ctx_ptr); + struct EvalCacheGuard { + TrainingContext* ctx; + ~EvalCacheGuard() { + // Evaluation builds projection caches under no-grad. They + // must never be reused by the following training step. + if (!ctx) return; + ctx->lora_cache_valid = false; + ctx->lora_batch_valid = false; + } + } cache_guard{ctx}; + TORCH_CHECK(ctx, "null training context"); TORCH_CHECK(ctx->adapters.empty(), "dynamic LoRA adapters require selected multi-LoRA evaluation; " "ordinary eval_step has no tenant mapping"); diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs index c9c48a4f..7bdc1cc5 100644 --- a/crates/rustrain-qwen3-6/src/kernel.rs +++ b/crates/rustrain-qwen3-6/src/kernel.rs @@ -73,11 +73,12 @@ type FnSetMtpWeights = unsafe extern "C" fn( i64, // num_mtp_layer_weights *mut c_void, // mtp_layer_configs_ptr i64, // num_mtp_layers -); +) -> i32; type FnSetCheckpoint = unsafe extern "C" fn(*mut c_void, i32, i64); type FnSetNcclComm = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, i32, i32); type FnInitNccl = unsafe extern "C" fn(*mut c_void) -> i32; type FnSetCudaDevice = unsafe extern "C" fn(i32); +type FnSetBaseTpMlp = unsafe extern "C" fn(*mut c_void, i32) -> i32; type FnAddLora = unsafe extern "C" fn(*mut c_void, i64, f64, *const i64, i64, *const i8) -> i64; type FnRemoveLora = unsafe extern "C" fn(*mut c_void, i64) -> i32; type FnListLora = unsafe extern "C" fn(*mut c_void, *mut i64, i64) -> i64; @@ -144,6 +145,7 @@ struct KernelHandles { set_nccl_comm: FnSetNcclComm, init_nccl: FnInitNccl, set_cuda_device: FnSetCudaDevice, + set_base_tp_mlp: FnSetBaseTpMlp, add_lora: FnAddLora, remove_lora: FnRemoveLora, list_lora: FnListLora, @@ -194,7 +196,7 @@ unsafe fn load_kernels() -> Option { }}; } let abi_version: FnKernelAbiVersion = sym!("qwen36_kernel_abi_version"); - if abi_version() != 11 { + if abi_version() != 13 { return None; } Some(KernelHandles { @@ -222,6 +224,7 @@ unsafe fn load_kernels() -> Option { set_nccl_comm: sym!("qwen36_set_nccl_comm"), init_nccl: sym!("qwen36_init_nccl"), set_cuda_device: sym!("qwen36_set_cuda_device"), + set_base_tp_mlp: sym!("qwen36_set_base_tp_mlp"), add_lora: sym!("qwen36_add_lora"), remove_lora: sym!("qwen36_remove_lora"), list_lora: sym!("qwen36_list_lora"), @@ -252,6 +255,38 @@ fn get_ptr(weights: &std::collections::BTreeMap, name: &str) -> } } +/// Return the rank-local frozen dense MLP shard for TP, or `None` when the +/// tensor is not one of gate/up/down. This runs during CPU weight loading. +pub fn shard_dense_mlp_weight_for_tp( + name: &str, + tensor: &Tensor, + tp_size: usize, + tp_rank: usize, +) -> Result> { + let dim = if name.ends_with(".mlp.gate_proj.weight") || name.ends_with(".mlp.up_proj.weight") { + 0 + } else if name.ends_with(".mlp.down_proj.weight") { + 1 + } else { + return Ok(None); + }; + if tp_size <= 1 || tp_rank >= tp_size { + bail!("invalid dense MLP TP shard: tp_rank={tp_rank}, tp_size={tp_size}"); + } + let full = *tensor + .size() + .get(dim as usize) + .ok_or_else(|| anyhow::anyhow!("base TP MLP weight {name} has no dimension {dim}"))?; + if full <= 0 || full % tp_size as i64 != 0 { + bail!( + "base TP MLP weight {name} dimension {dim}={full} is not divisible by TP_SIZE={tp_size}" + ); + } + let shard = full / tp_size as i64; + let start = tp_rank as i64 * shard; + Ok(Some(tensor.narrow(dim, start, shard).contiguous())) +} + pub fn build_weight_ptrs( weights: &std::collections::BTreeMap, config: &crate::config::Qwen36RuntimeConfig, @@ -468,6 +503,7 @@ impl CppTrainingContext { eps: f64, lora_scaling: f64, lora_rank: i64, + base_tp_mlp: bool, target_layers: &[usize], target_modules: &[Qwen36LoraTargetModule], expert_start: usize, @@ -548,6 +584,11 @@ impl CppTrainingContext { if ptr.is_null() { bail!("C++ create_training_context returned null"); } + let base_tp_status = unsafe { (kh.set_base_tp_mlp)(ptr, i32::from(base_tp_mlp)) }; + if base_tp_status != 0 { + unsafe { (kh.free_ctx)(ptr) }; + bail!("C++ base dense MLP TP configuration failed"); + } let lora_count = unsafe { (kh.get_lora_count)(ptr) }; Ok(Self { ptr, lora_count }) } @@ -750,7 +791,7 @@ impl CppTrainingContext { let lc_ptr = mtp_layer_configs.as_ptr() as *mut c_void; std::mem::forget(mtp_layer_configs); - unsafe { + let status = unsafe { (kh.set_mtp_weights)( self.ptr, mtp_fc_ptr, @@ -761,7 +802,10 @@ impl CppTrainingContext { wp_len as i64, lc_ptr, config.mtp_num_hidden_layers as i64, - ); + ) + }; + if status != 0 { + bail!("C++ set_mtp_weights rejected the requested context state"); } Ok(()) } @@ -1103,3 +1147,41 @@ impl Drop for CppTrainingContext { } } } + +#[cfg(test)] +mod tests { + use super::shard_dense_mlp_weight_for_tp; + use tch::{Kind, Tensor}; + + #[test] + fn dense_mlp_tp_shards_matching_intermediate_axes() { + let gate = Tensor::arange(48, (Kind::Float, tch::Device::Cpu)).reshape([12, 4]); + let down = Tensor::arange(48, (Kind::Float, tch::Device::Cpu)).reshape([4, 12]); + let gate_rank_one = + shard_dense_mlp_weight_for_tp("model.layers.0.mlp.gate_proj.weight", &gate, 2, 1) + .unwrap() + .unwrap(); + let down_rank_one = + shard_dense_mlp_weight_for_tp("model.layers.0.mlp.down_proj.weight", &down, 2, 1) + .unwrap() + .unwrap(); + assert_eq!(gate_rank_one.size(), [6, 4]); + assert_eq!(down_rank_one.size(), [4, 6]); + assert_eq!(gate_rank_one.double_value(&[0, 0]), 24.0); + assert_eq!(down_rank_one.double_value(&[0, 0]), 6.0); + } + + #[test] + fn dense_mlp_tp_ignores_non_mlp_weights_and_rejects_bad_shapes() { + let tensor = Tensor::zeros([5, 4], (Kind::Float, tch::Device::Cpu)); + assert!( + shard_dense_mlp_weight_for_tp("model.layers.0.self_attn.q_proj.weight", &tensor, 2, 0) + .unwrap() + .is_none() + ); + assert!( + shard_dense_mlp_weight_for_tp("model.layers.0.mlp.up_proj.weight", &tensor, 2, 0) + .is_err() + ); + } +} diff --git a/crates/rustrain-qwen3-6/src/session.rs b/crates/rustrain-qwen3-6/src/session.rs index 83f1599e..5b69981d 100644 --- a/crates/rustrain-qwen3-6/src/session.rs +++ b/crates/rustrain-qwen3-6/src/session.rs @@ -304,7 +304,10 @@ fn train_impl( ); } if lora_config.rank % tp_size as i64 != 0 { - bail!("LoRA rank {} must be divisible by TP_SIZE={tp_size}", lora_config.rank); + bail!( + "LoRA rank {} must be divisible by TP_SIZE={tp_size}", + lora_config.rank + ); } unsafe { std::env::set_var("TP_SIZE", tp_size.to_string()); @@ -325,6 +328,45 @@ fn train_impl( ); } + // Dense base-weight TP follows Megatron's ColumnParallel/RowParallel MLP: + // gate/up rows are owned by one TP rank and down columns are owned by the + // same rank. Attention and embeddings remain replicated in this first + // unit. MoE, MTP, and mixed TP/DP/EP are rejected rather than silently + // running with an unsharded path. + let base_tp_mlp = tp_size > 1 && !runtime_config.is_moe; + if base_tp_mlp { + if runtime_config.mtp_num_hidden_layers > 0 { + bail!("base dense TP MLP currently requires MTP to be disabled"); + } + if runtime_config.intermediate_size <= 0 + || runtime_config.intermediate_size % tp_size as i64 != 0 + { + bail!( + "dense intermediate_size={} must be divisible by TP_SIZE={tp_size}", + runtime_config.intermediate_size + ); + } + if lora_config.target_modules.is_empty() + || lora_config.target_modules.iter().any(|module| { + matches!( + module.cpp_name(), + "gate_proj" + | "up_proj" + | "down_proj" + | "shared_gate_proj" + | "shared_up_proj" + | "shared_down_proj" + | "experts_gate_up_proj" + | "experts_down_proj" + ) + }) + { + bail!( + "base dense TP MLP currently requires explicit attention-only LoRA target_modules" + ); + } + } + // Build needed weight set let needed = build_needed_weights(&runtime_config, &lora_config, shard_ref); @@ -371,7 +413,24 @@ fn train_impl( ); } else { for (name, tensor) in &weights { - weights_gpu.insert(name.clone(), tensor.to_device(device).to_kind(compute_kind)); + let local_shard = if base_tp_mlp { + crate::kernel::shard_dense_mlp_weight_for_tp(name, tensor, tp_size, rank % tp_size)? + } else { + None + }; + let gpu_tensor = local_shard + .as_ref() + .unwrap_or(tensor) + .to_device(device) + .to_kind(compute_kind); + weights_gpu.insert(name.clone(), gpu_tensor); + } + if base_tp_mlp { + info!( + tp_size, + tp_rank = rank % tp_size, + "base dense MLP TP enabled: gate/up row shards and down column shards" + ); } } @@ -429,6 +488,7 @@ fn train_impl( config.train.adam_eps as f64, lora_config.alpha as f64 / lora_config.rank as f64, // lora scaling = alpha / rank lora_config.rank as i64, + base_tp_mlp, &lora_config.target_layers, &lora_config.target_modules, expert_start, @@ -458,12 +518,7 @@ fn train_impl( // Set MTP weights if available if runtime_config.mtp_num_hidden_layers > 0 { - ctx.set_mtp_weights( - &weights_gpu, - &runtime_config, - expert_start, - expert_count, - )?; + ctx.set_mtp_weights(&weights_gpu, &runtime_config, expert_start, expert_count)?; info!( "C++ TrainingContext: MTP weights set ({} layers)", runtime_config.mtp_num_hidden_layers diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index f4deca1f..d3533a61 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -366,7 +366,7 @@ static int run_dynamic_dp_smoke( } int main() { - assert(qwen36_kernel_abi_version() == 11); + assert(qwen36_kernel_abi_version() == 13); const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); const int process_rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0"); const int local_rank = std::atoi(std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); diff --git a/crates/rustrain-qwen3-6/tests/native_tp_mlp_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_mlp_smoke.cpp new file mode 100644 index 00000000..52da9995 --- /dev/null +++ b/crates/rustrain-qwen3-6/tests/native_tp_mlp_smoke.cpp @@ -0,0 +1,225 @@ +#include +#include + +#include +#include +#include +#include +#include + +struct LayerConfig { + int64_t layer_type, num_heads, num_kv_heads, head_dim; + int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel; + double partial_rotary_factor, rope_theta, rms_eps; + int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count; + int64_t intermediate_size; + int32_t norm_topk_prob; + void* nccl_comm; + void* nccl_stream; +}; + +extern "C" void qwen36_set_cuda_device(int32_t); +extern "C" void* qwen36_create_training_context( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*); +extern "C" int32_t qwen36_init_nccl(void*); +extern "C" int32_t qwen36_set_base_tp_mlp(void*, int32_t); +extern "C" int32_t qwen36_set_mtp_weights( + void*, void*, void*, void*, void*, void**, int64_t, void*, int64_t); +extern "C" int64_t qwen36_add_lora( + void*, int64_t, double, const int64_t*, int64_t, const char*); +extern "C" void* qwen36_get_lora_a(void*, int64_t); +extern "C" void* qwen36_get_lora_b(void*, int64_t); +extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*); +extern "C" void* qwen36_get_lora_grad_accumulator(void*, int64_t, int32_t); +extern "C" int32_t qwen36_abort_gradient_accumulation(void*); +extern "C" double qwen36_eval_step(void*, void*, void*, void*); +extern "C" double qwen36_train_step(void*, void*, void*, void*); +extern "C" double qwen36_train_micro_step( + void*, void*, void*, void*, double, int32_t); +extern "C" void qwen36_free_training_context(void*); + +static at::Tensor deterministic(std::initializer_list shape, double scale) { + int64_t count = 1; + for (int64_t dim : shape) count *= dim; + return ((at::arange(count, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)) + .remainder(17) - 8.0) * scale) + .reshape(shape).to(at::kBFloat16); +} + +static std::vector pointers(std::vector& tensors) { + std::vector result; + result.reserve(tensors.size()); + for (auto& tensor : tensors) result.push_back(&tensor); + return result; +} + +int main() { + const int process_rank = std::atoi(std::getenv("RANK")); + const int world = std::atoi(std::getenv("WORLD_SIZE")); + assert(world == 2 && process_rank >= 0 && process_rank < world); + qwen36_set_cuda_device(process_rank); + + constexpr int64_t hidden = 8; + constexpr int64_t intermediate = 12; + constexpr int64_t vocab = 16; + constexpr int64_t lora_rank = 4; + constexpr int64_t local_lora_rank = lora_rank / 2; + + std::vector full_weights; + full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + full_weights.push_back(deterministic({2 * hidden, hidden}, 0.01)); + full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + full_weights.push_back(deterministic({hidden, hidden}, 0.012)); + full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + full_weights.push_back(deterministic({hidden, hidden}, 0.008)); + full_weights.push_back(deterministic({hidden, hidden}, 0.011)); + full_weights.push_back(deterministic({intermediate, hidden}, 0.009)); + full_weights.push_back(deterministic({intermediate, hidden}, 0.007)); + full_weights.push_back(deterministic({hidden, intermediate}, 0.01)); + for (auto& weight : full_weights) weight.set_requires_grad(false); + + const int64_t intermediate_start = process_rank * (intermediate / world); + std::vector local_weights(full_weights.begin(), full_weights.begin() + 8); + local_weights.push_back(full_weights[8].narrow(0, intermediate_start, intermediate / world).contiguous()); + local_weights.push_back(full_weights[9].narrow(0, intermediate_start, intermediate / world).contiguous()); + local_weights.push_back(full_weights[10].narrow(1, intermediate_start, intermediate / world).contiguous()); + assert(local_weights[8].sizes() == at::IntArrayRef({intermediate / 2, hidden})); + assert(local_weights[9].sizes() == at::IntArrayRef({intermediate / 2, hidden})); + assert(local_weights[10].sizes() == at::IntArrayRef({hidden, intermediate / 2})); + + auto embed = deterministic({vocab, hidden}, 0.02); + auto final_norm = at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)); + auto lm_head = deterministic({vocab, hidden}, 0.015); + LayerConfig config{}; + config.layer_type = 0; + config.num_heads = 1; + config.num_kv_heads = 1; + config.head_dim = hidden; + config.partial_rotary_factor = 1.0; + config.rope_theta = 10000.0; + config.rms_eps = 1e-5; + config.intermediate_size = intermediate; + + const int64_t target_layer = 0; + auto local_ptrs = pointers(local_weights); + setenv("TP_SIZE", "2", 1); + unsetenv("RUSTRAIN_DATA_PARALLEL"); + void* rejected = qwen36_create_training_context( + local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head, + &config, 1, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + &target_layer, 1, "gate_proj"); + assert(rejected && qwen36_set_base_tp_mlp(rejected, 1) != 0); + qwen36_free_training_context(rejected); + + void* dynamic_rejected = qwen36_create_training_context( + local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head, + &config, 1, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + &target_layer, 1, "q_proj"); + assert(dynamic_rejected); + assert(qwen36_add_lora( + dynamic_rejected, lora_rank, 1.0, &target_layer, 1, "gate_proj") > 0); + assert(qwen36_set_base_tp_mlp(dynamic_rejected, 1) != 0); + qwen36_free_training_context(dynamic_rejected); + + void* distributed = qwen36_create_training_context( + local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head, + &config, 1, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + &target_layer, 1, "q_proj"); + assert(distributed && qwen36_set_base_tp_mlp(distributed, 1) == 0); + assert(qwen36_set_base_tp_mlp(distributed, 0) != 0); + assert(qwen36_set_mtp_weights( + distributed, nullptr, nullptr, nullptr, nullptr, nullptr, 0, nullptr, 0) != 0); + assert(qwen36_add_lora( + distributed, lora_rank, 1.0, &target_layer, 1, nullptr) < 0); + assert(qwen36_init_nccl(distributed) == 0); + + setenv("TP_SIZE", "1", 1); + auto full_ptrs = pointers(full_weights); + void* reference = qwen36_create_training_context( + full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head, + &config, 1, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + &target_layer, 1, "q_proj"); + assert(reference); + + auto full_a = deterministic({lora_rank, hidden}, 0.002); + auto full_b = deterministic({2 * hidden, lora_rank}, 0.001); + auto local_a = full_a.narrow(0, process_rank * local_lora_rank, local_lora_rank).contiguous(); + auto local_b = full_b.narrow(1, process_rank * local_lora_rank, local_lora_rank).contiguous(); + assert(qwen36_set_lora_tensor(distributed, 0, 0, &local_a) == 0); + assert(qwen36_set_lora_tensor(distributed, 0, 1, &local_b) == 0); + assert(qwen36_set_lora_tensor(reference, 0, 0, &full_a) == 0); + assert(qwen36_set_lora_tensor(reference, 0, 1, &full_b) == 0); + auto local_a_before = local_a.clone(); + auto local_b_before = local_b.clone(); + + auto input_ids = at::tensor({1, 2, 3}, + at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 3}); + auto target_mask = at::ones({1, 3}, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); + auto attention_mask = at::ones({1, 3}, + at::TensorOptions().device(at::kCUDA).dtype(at::kBool)); + + const double distributed_eval = qwen36_eval_step( + distributed, &input_ids, &target_mask, &attention_mask); + const double reference_eval = qwen36_eval_step( + reference, &input_ids, &target_mask, &attention_mask); + assert(distributed_eval > 0.0 && reference_eval > 0.0); + assert(std::abs(distributed_eval - reference_eval) < 1e-4); + + const double distributed_micro = qwen36_train_micro_step( + distributed, &input_ids, &target_mask, &attention_mask, 1.0, 0); + const double reference_micro = qwen36_train_micro_step( + reference, &input_ids, &target_mask, &attention_mask, 1.0, 0); + assert(distributed_micro > 0.0 && reference_micro > 0.0); + auto* distributed_b_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(distributed, 0, 1)); + auto* reference_b_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(reference, 0, 1)); + assert(distributed_b_grad && reference_b_grad); + auto reference_b_grad_slice = reference_b_grad->narrow( + 1, process_rank * local_lora_rank, local_lora_rank); + const double b_grad_diff = + (*distributed_b_grad - reference_b_grad_slice).abs().max().item(); + assert(distributed_b_grad->abs().max().item() > 0.0); + assert(b_grad_diff < 1e-4); + assert(qwen36_abort_gradient_accumulation(distributed) == 0); + assert(qwen36_abort_gradient_accumulation(reference) == 0); + + const double distributed_loss = qwen36_train_step( + distributed, &input_ids, &target_mask, &attention_mask); + const double reference_loss = qwen36_train_step( + reference, &input_ids, &target_mask, &attention_mask); + assert(distributed_loss > 0.0 && reference_loss > 0.0); + assert(std::abs(distributed_loss - reference_loss) < 1e-4); + + auto* updated_a = reinterpret_cast(qwen36_get_lora_a(distributed, 0)); + auto* updated_b = reinterpret_cast(qwen36_get_lora_b(distributed, 0)); + auto* reference_a = reinterpret_cast(qwen36_get_lora_a(reference, 0)); + auto* reference_b = reinterpret_cast(qwen36_get_lora_b(reference, 0)); + auto reference_a_slice = reference_a->narrow( + 0, process_rank * local_lora_rank, local_lora_rank); + auto reference_b_slice = reference_b->narrow( + 1, process_rank * local_lora_rank, local_lora_rank); + const double a_diff = (*updated_a - reference_a_slice).abs().max().item(); + const double b_diff = (*updated_b - reference_b_slice).abs().max().item(); + const double a_update = (*updated_a - local_a_before).abs().max().item(); + const double b_update = (*updated_b - local_b_before).abs().max().item(); + std::printf( + "base_tp_mlp_smoke rank=%d eval_diff=%0.8e loss_diff=%0.8e b_grad_diff=%0.8e a_diff=%0.8e b_diff=%0.8e\n", + process_rank, std::abs(distributed_eval - reference_eval), + std::abs(distributed_loss - reference_loss), b_grad_diff, a_diff, b_diff); + assert(a_diff < 1e-5 && b_diff < 1e-5); + assert(a_update > 0.0 || b_update > 0.0); + + qwen36_free_training_context(reference); + qwen36_free_training_context(distributed); + return 0; +} diff --git a/crates/rustrain-server/src/session.rs b/crates/rustrain-server/src/session.rs index d188b19e..d34fffbd 100644 --- a/crates/rustrain-server/src/session.rs +++ b/crates/rustrain-server/src/session.rs @@ -320,12 +320,12 @@ impl TrainingSession for Qwen36Session { .and_then(|s| s.parse::().ok()) .unwrap_or(1); if tp_size > 1 { - if ep_world_size != tp_size || req.rank < 0 || req.rank as usize >= ep_world_size { + if ep_world_size != tp_size || ep_rank >= ep_world_size { return Err(anyhow!( "native Qwen server TP-only mode requires WORLD_SIZE=TP_SIZE and a valid global rank (world={}, tp={}, rank={})", ep_world_size, tp_size, - req.rank + ep_rank )); } unsafe { @@ -334,6 +334,41 @@ impl TrainingSession for Qwen36Session { } let is_ep = ep_world_size > 1 && runtime_config.is_moe && tp_size == 1; let is_data_parallel = ep_world_size > 1 && !runtime_config.is_moe && tp_size == 1; + let base_tp_mlp = tp_size > 1 && !runtime_config.is_moe; + if base_tp_mlp { + if runtime_config.mtp_num_hidden_layers > 0 { + return Err(anyhow!( + "base dense TP MLP currently requires MTP to be disabled" + )); + } + if runtime_config.intermediate_size <= 0 + || runtime_config.intermediate_size % tp_size as i64 != 0 + { + return Err(anyhow!( + "dense intermediate_size={} must be divisible by TP_SIZE={tp_size}", + runtime_config.intermediate_size + )); + } + if req.target_modules.is_empty() + || req.target_modules.iter().any(|name| { + matches!( + name.as_str(), + "gate_proj" + | "up_proj" + | "down_proj" + | "shared_gate_proj" + | "shared_up_proj" + | "shared_down_proj" + | "experts_gate_up_proj" + | "experts_down_proj" + ) + }) + { + return Err(anyhow!( + "base dense TP MLP currently requires explicit attention-only LoRA target_modules" + )); + } + } // Compute expert shard let (expert_start, expert_count) = if is_ep { @@ -371,7 +406,20 @@ impl TrainingSession for Qwen36Session { .to_kind(self.compute_kind); weights.insert(name, narrowed); } else { - let t = tensor.to_device(self.device); + let local_shard = if base_tp_mlp { + rustrain_qwen3_6::kernel::shard_dense_mlp_weight_for_tp( + &name, + &tensor, + tp_size, + ep_rank % tp_size, + )? + } else { + None + }; + let t = local_shard + .as_ref() + .unwrap_or(&tensor) + .to_device(self.device); let processed = t.to_kind(self.compute_kind); weights.insert(name, processed); } @@ -409,6 +457,7 @@ impl TrainingSession for Qwen36Session { req.eps, lora_scaling, req.rank, + base_tp_mlp, &all_layers, &target_modules, expert_start, diff --git a/docs/plans/qwen-lora-megatron-progress.md b/docs/plans/qwen-lora-megatron-progress.md index b8ba4786..a9db0caa 100644 --- a/docs/plans/qwen-lora-megatron-progress.md +++ b/docs/plans/qwen-lora-megatron-progress.md @@ -12,9 +12,9 @@ timestamp: 2026-07-17T00:00:00Z # Current State -Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP parity against a full-expert reference, variable-split EP A2A with fixed-LoRA data sharding, dense replicated-DP smoke with per-tenant token weighting, TP-only latent-rank-sharded LoRA smoke, dynamic batch logical-step update, ABI11 FP32 gradient storage/aggregation, standard Adam bias correction, per-tenant optimizer-step restore, selected-tenant isolation, same-topology rank-aware checkpointing, and 5D topology mapping. +Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP parity against a full-expert reference, variable-split EP A2A with fixed-LoRA data sharding, dense replicated-DP smoke with per-tenant token weighting, TP-only latent-rank-sharded LoRA smoke, frozen dense SwiGLU MLP base-weight TP, dynamic batch logical-step update, ABI11 FP32 gradient storage/aggregation, standard Adam bias correction, per-tenant optimizer-step restore, selected-tenant isolation, same-topology rank-aware checkpointing, and 5D topology mapping. -Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axis TP+DP/EP, PP/CP, DeepEP/TE prebuilt integration, server-side source-sharded dynamic dispatch, heterogeneous dynamic adapter signatures, dynamic+MTP objective normalization, cross-topology checkpoint resharding, and matched Megatron throughput. Native direct dynamic source metadata under sharded A2A is now implemented and full-reference smoke-tested. +Not yet verified or implemented: base-weight TP for attention/GDN/MoE/embedding/LM-head, MLP-targeted LoRA under base TP, multi-axis TP+DP/EP, PP/CP, DeepEP/TE prebuilt integration, server-side source-sharded dynamic dispatch, heterogeneous dynamic adapter signatures, dynamic+MTP objective normalization, cross-topology checkpoint resharding, and matched Megatron throughput. Native direct dynamic source metadata under sharded A2A is implemented and full-reference smoke-tested. # Durable Milestones @@ -39,16 +39,18 @@ Not yet verified or implemented: Megatron-style frozen base-weight TP, multi-axi - Target runtime probe: PyTorch 2.5.1+cu121, ABI0; Transformer Engine, flash-attn, DeepEP, Triton, and DeepSpeed are not importable. - H20 `123.57.26.97:28004`: `native_ep_bench.cpp` fresh ABI0 benchmark (`seq=128, hidden=256, experts=8, intermediate=256, warmup=2, iters=10`) passed legacy and sharded A2A with `rank_statuses=0,0`. Legacy median was about `5.47 ms` / `46.4k processed tokens/s` (`23.2k unique tokens/s`), while sharded A2A was about `6.72 ms` / `37.8k processed and unique tokens/s`. This is a synthetic native baseline, not Megatron-LM parity. - H20 `123.57.26.97:28004`: fresh ABI1 dynamic sharded native smoke passed on both ranks (`rank_statuses=0,0`) against a full-expert reference with complementary source masks. Dynamic grouped-expert parameter/m/v maxima were `1.53e-5` / `4.88e-5` / `5.75e-8`; it also exercised a third tenant with zero global target tokens, clocks `[2,2,0]`, and explicit rejection of ordinary `train_step`. The server path still broadcasts replicated source batches, and dynamic+MTP is explicitly rejected until its two objective denominators are separated. +- H20 `123.57.26.97:28004`: fresh ABI13 dense base-MLP TP2 smoke passed on both ranks. CLI and server use one CPU sharding helper for gate/up rows and matching down columns; the per-context native flag validates local shapes without process-global env state. C++ reduces the row-parallel output and all-reduces column-parallel input dgrad. Eval/train loss differed from the replicated full-weight reference by `1.84e-5`; FP32 LoRA gradient-accumulator maxima were `3.05e-5` / `4.58e-5`, and the largest post-Adam LoRA slice difference was `9.31e-9`. The smoke also rejects MLP LoRA targets, prevents incompatible TP state transitions, and covers eval-to-train cache invalidation. # Decisions During Execution - Keep TP and EP communicators separate; do not reuse the existing EP `LayerConfig.nccl_comm` for LoRA TP deltas. - Publish multi-LoRA `n_max` directly from rank 0 with `ncclBroadcast`; filesystem rendezvous can reuse stale files across process restarts and give ranks different chunk schedules. - Do not enable Qwen TP/PP/CP by merely relaxing runtime validation. +- Treat dense base-MLP TP as one accepted slice, not full model TP. Until projection-aware LoRA collectives exist, reject gate/up/down LoRA targets instead of applying the replicated-projection reduction rule to disjoint output shards. - Treat Exa/Jina dependency search failures as missing evidence, not as proof that a package is compatible. # Verification -Passed: `cargo check -p rustrain-qwen3-6 -p rustrain-server -p rustrain-ipc` (with the repository PyTorch 2.12.1 host venv), `cargo test -p rustrain-qwen3-6 --lib` (3), `cargo test -p rustrain-server --lib` (6), Qwen integration tests, remote ABI8 smoke, remote ABI9 selected-tenant native smoke, remote ABI10 single-rank native smoke, remote ABI10 two-rank TP native smoke, remote ABI11 single/TP2/DP2/EP2 native smoke with numerical Adam and parity oracles, and the ABI1 dynamic sharded full-reference smoke described above. +Passed: `cargo check -p rustrain-qwen3-6 -p rustrain-server -p rustrain-ipc` (with the repository host venv), `cargo test -p rustrain-qwen3-6 --lib` (5), `cargo test -p rustrain-server --lib` (6), Qwen integration tests, remote ABI8 smoke, remote ABI9 selected-tenant native smoke, remote ABI10 single-rank native smoke, remote ABI10 two-rank TP native smoke, remote ABI11 single/TP2/DP2/EP2 native smoke with numerical Adam and parity oracles, ABI1 dynamic sharded full-reference smoke, and ABI13 dense base-MLP TP2 parity smoke. -Not run: Megatron-style base-model TP, multi-axis TP+DP/EP, PP/CP, server-side source-sharded dynamic dispatch, heterogeneous dynamic adapter signatures, dynamic+MTP, cross-topology resharding, and matched Megatron performance benchmark. The target host lacks importable Megatron/Transformer Engine/DeepEP/flash-attn prebuilt packages, so no dependency installation or JIT workaround was used. +Not run: full-model base TP beyond dense MLP, multi-axis TP+DP/EP, PP/CP, server-side source-sharded dynamic dispatch, heterogeneous dynamic adapter signatures, dynamic+MTP, cross-topology resharding, and matched Megatron performance benchmark. The target host lacks importable Megatron/Transformer Engine/DeepEP/flash-attn prebuilt packages, so no dependency installation or JIT workaround was used. diff --git a/docs/qwen35-qwen36-megatron-audit.md b/docs/qwen35-qwen36-megatron-audit.md index 3b4c5ff5..41793b7d 100644 --- a/docs/qwen35-qwen36-megatron-audit.md +++ b/docs/qwen35-qwen36-megatron-audit.md @@ -7,7 +7,7 @@ - 模型语义：Qwen3.5 dense、Qwen3.6 dense/MoE 的 native forward/backward 路径已经覆盖 hybrid full attention、GDN、MoE、MTP 和 LoRA 目标模块；已有配置解析、集成测试及 H20 native smoke 证据。 - 已实现并可验证的分布式子集：MoE expert parallel，以及 replicated LoRA 的 data parallel；梯度累积和 dynamic multi-LoRA 已有 logical-step 边界。DP 动态租户按 adapter token count 加权，sharded A2A native 路径会保留 source flattened row 来恢复租户，并按全局租户 token count 归一化。 - 性能：MoE grouped dispatch 相对逐 expert matmul 的已有 microbenchmark 为约 3.70x（E=32, N=4096, H=2048, I=768，结果误差为 0）；这不是端到端训练吞吐或 Megatron 对比。 -- 已实现 LoRA latent-rank 的 TP-only 子集，但尚未实现 frozen base-weight tensor parallel、pipeline parallel、context parallel，以及 TP/PP/CP 与 EP/DP 的组合。当前训练上下文仍由单个进程持有完整 dense 权重和完整层栈。 +- 已实现 LoRA latent-rank TP-only，以及 frozen dense SwiGLU MLP 的 gate/up row shard、down column shard 和输出 all-reduce。attention/GDN/MoE/vocab 仍复制，MLP LoRA 在 base TP 下暂拒绝；PP/CP 和 TP 与 EP/DP 的组合仍未实现。 - 因此当前实现不能宣称“Megatron-LM 级别”。它是一个计算集中在 C++ 的 LoRA/EP/DP 子集，离 Megatron 的完整并行和通信重叠仍有实质差距。 ## 当前能力矩阵 @@ -23,7 +23,7 @@ | microbatch accumulation | 已实现子集 | non-final microbatch 只 backward，final microbatch 才 optimizer；FP32 accumulator 存储/聚合，autograd leaf backward 仍为 BF16 | | replicated data parallel | 已实现 | logical-step 边界同步 replicated LoRA；EP expert 参数不走该 reduction | | expert parallel | 已实现子集 | 默认 routed-output all-reduce；gated variable-split A2A 已验证 fixed-LoRA 和 native dynamic-LoRA data sharding；GPU-only split planning、异步 overlap 和 DeepEP backend 未实现 | -| tensor parallel | LoRA-only 子集 | latent rank 分片和独立 TP communicator 已验证；attention/MLP/LM-head base 权重仍不切分 | +| tensor parallel | dense MLP 子集 | latent rank 分片和独立 TP communicator 已验证；dense gate/up/down base 权重按 intermediate 维切分并有 TP2 full-reference smoke；attention/GDN/MoE/LM-head 仍不切分，MLP LoRA/MTP 暂拒绝 | | pipeline parallel | 未实现于 Qwen native | 没有 stage 切分、microbatch scheduler 或 activation send/recv | | context parallel | 未实现于 Qwen native | 没有 ring attention、跨 rank KV/索引合并 | | distributed checkpoint | 已实现子集 | same-topology TP rank-sharded v3 已验证；跨 topology reshard 和 PP/CP 未实现 | @@ -32,7 +32,9 @@ ### 并行语义 -Megatron 的 TP 会按 head、hidden/intermediate 和 vocab 维度切分权重，并在线性层边界执行必要的 reduce-scatter/all-reduce；PP 会把层分到不同 stage 并使用 1F1B 等调度；CP 会在序列和 attention state 上做跨 rank 通信。当前 Qwen native `TrainingContext` 仍加载完整模型并在一个 C++ forward 中执行全部层，因此仅增加 `tensor_model_parallel_size` 等配置不能得到正确的 TP/PP/CP。 +Megatron 的通用 MLP 使用 ColumnParallel fc1、local gated activation 和 RowParallel fc2，并进一步提供 fused fc1/activation、sequence parallel、通信重叠和 sharded-state 支持。当前 Qwen native 已补上算法等价的 separate gate/up row shard 与 down column shard，但仍是两次独立 GEMM，且 attention/GDN/MoE/vocab 不分片。Megatron-LM 本仓库也没有直接的 Qwen3.5/3.6 模型实现，精确模型入口依赖外部 bridge，因此这里比较的是成熟的通用并行基础设施，不是同模型端到端实现。 + +PP 会把层分到不同 stage 并使用 1F1B 等调度；CP 会在序列和 attention state 上做跨 rank 通信。当前 Qwen native `TrainingContext` 仍在每个进程执行完整层栈，因此仅增加 PP/CP 配置不能得到正确语义。 当前 DP/EP 也不是完整 Megatron 语义：DP 同步 replicated LoRA 梯度并按租户 token count 归一化，expert 参数留在 EP rank；EP 默认使用 routed-output all-reduce，实验 gate 已加入 variable-split dispatch/inverse combine 和 fixed/dynamic LoRA data sharding。dynamic native path 的 source row metadata 已随 token index 传输，但仍逐 top-k 执行 host-visible count sync，没有 fused permutation、GPU-only split planning、异步 overlap 或 DeepEP backend。server 的 `TrainMultiLora` 目前向各 worker 广播相同 batch，因此不能把它当作 source-sharded 服务吞吐。 @@ -50,12 +52,12 @@ Megatron 的 TP 会按 head、hidden/intermediate 和 vocab 维度切分权重 ## 验证边界 -已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试，以及 H20 ABI11 的单卡、TP2、EP2 和 DP2 native smoke。DP2 的 weighted m/grouped/v/Adam BF16 delta oracle 分别达到 `2.43e-8`、`2.27e-8`、`7.33e-8` 和 `0`；legacy EP2 与 full-expert reference 的 loss、LoRA、Adam state 和标准 Adam 首步 oracle 在两 rank 均为零差异。Replicated A2A 与 fixed-LoRA sharded A2A 也均通过两 rank full-expert reference；sharded token counts `[1,3]` 的加权 loss 与 global reference 相差约 `9.6e-7`，m/v 最大差 `1.22e-5` / `3.92e-9`，Adam oracle 差为 `0`。新增 H20 ABI1 dynamic sharded full-reference smoke 在两 rank 返回 `0`：互补 source masks 的 dynamic grouped-expert 参数最大差 `1.53e-5`，m/v 最大差 `4.88e-5` / `5.75e-8`；两次有效租户更新后 step 为 `[2,2]`，全局零 token 租户保持 step `0` 且参数无更新，普通 `train_step` 被显式拒绝。没有完成 Qwen3.5/3.6 完整大模型的长时间训练、跨节点通信、PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖且有直接 oracle 的子集，而不是所有并行配置。 +已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试，以及 H20 ABI11 的单卡、TP2、EP2 和 DP2 native smoke。DP2 的 weighted m/grouped/v/Adam BF16 delta oracle 分别达到 `2.43e-8`、`2.27e-8`、`7.33e-8` 和 `0`；legacy EP2 与 full-expert reference 的 loss、LoRA、Adam state 和标准 Adam 首步 oracle 在两 rank 均为零差异。Replicated A2A 与 fixed-LoRA sharded A2A 也均通过两 rank full-expert reference；sharded token counts `[1,3]` 的加权 loss 与 global reference 相差约 `9.6e-7`，m/v 最大差 `1.22e-5` / `3.92e-9`，Adam oracle 差为 `0`。H20 ABI1 dynamic sharded full-reference smoke 在两 rank 返回 `0`：dynamic grouped-expert 参数最大差 `1.53e-5`，m/v 最大差 `4.88e-5` / `5.75e-8`。新增 ABI13 dense base-MLP TP2 smoke 对 gate/up/down 使用半尺寸本地权重，验证 row-parallel forward sum 与 column-parallel input dgrad sum；eval/train loss 与完整权重参考相差 `1.84e-5`，并直接比较 FP32 LoRA accumulator 和更新切片。没有完成完整大模型长时间训练、跨节点通信、完整 base TP、PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖且有直接 oracle 的子集，而不是所有并行配置。 ## 继续达到 Megatron 级别所需的最小工作包 1. 建立 5D TP/PP/DP/EP/CP topology，并让 launcher、NCCL process groups 和 checkpoint 使用同一 rank 映射。 -2. 为 Qwen full/GDN attention、dense MLP、MoE、LM-head/CE 实现 TP shard 和对应 collective；为 PP 实现 stage forward/backward 与 1F1B scheduler；为 CP 实现 ring attention/state exchange。 +2. 补齐 Qwen full/GDN attention、MoE、LM-head/CE 的 TP shard，并为 dense MLP 增加 projection-aware LoRA、fused gate/up FC1 和 MTP 支持；为 PP 实现 stage forward/backward 与 1F1B scheduler；为 CP 实现 ring attention/state exchange。 3. 将 EP dispatch/combine 替换为 fused/异步路径，并测量通信与计算重叠。 4. 为 LoRA 增加 FP32 accumulation、每 adapter optimizer step、可恢复的 accumulation 状态和 rank-sharded checkpoint。 5. 在固定硬件和 workload 上，与 Megatron-LM 记录 tokens/s、step time、峰值显存、通信占比和 loss 曲线。 From e5bd9c41897fe21a3eda7101188dbc8e0d19d32c Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 12:33:47 +0800 Subject: [PATCH 025/156] feat: add qwen attention tensor parallelism --- .../kernels/qwen3_6_kernels.cpp | 405 ++++++++++++-- crates/rustrain-qwen3-6/src/kernel.rs | 100 +++- crates/rustrain-qwen3-6/src/session.rs | 71 ++- .../rustrain-qwen3-6/tests/native_smoke.cpp | 2 +- .../tests/native_tp_attention_smoke.cpp | 523 ++++++++++++++++++ .../tests/native_tp_latent_smoke.cpp | 298 ++++++++++ crates/rustrain-server/src/checkpoint.rs | 471 +++++++++++++++- crates/rustrain-server/src/session.rs | 314 +++++++++-- docs/plans/qwen-lora-megatron-progress.md | 13 +- docs/qwen35-qwen36-megatron-audit.md | 26 +- 10 files changed, 2075 insertions(+), 148 deletions(-) create mode 100644 crates/rustrain-qwen3-6/tests/native_tp_attention_smoke.cpp create mode 100644 crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index b078f24f..51583d1d 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -1108,6 +1108,12 @@ struct RoutedExpertLora { double scaling = 0.0; }; +enum class LoraTpLayout : uint8_t { + LatentRank, + ColumnParallel, + RowParallel, +}; + // Per-sample adapter projection used by the batched multi-LoRA path. This is // intentionally separate from RoutedExpertLora: routed experts carry one // A/B pair per local expert, while dense/shared projections carry one pair per @@ -1115,7 +1121,8 @@ struct RoutedExpertLora { struct LoraBatchEntry { at::Tensor a_stack; // [N, rank, in] at::Tensor b_stack; // [N, out, rank] - at::Tensor scaling; // [N, 1, 1] + at::Tensor scaling; // [N, 1, 1] + LoraTpLayout layout = LoraTpLayout::LatentRank; }; static const LoraBatchEntry* lora_batch_entry( @@ -1128,7 +1135,7 @@ static at::Tensor dense_mlp_forward_batched( static at::Tensor lora_activation_delta( TrainingContext* ctx, const at::Tensor& x, const at::Tensor& A, const at::Tensor& B, - const at::Tensor& scaling); + const at::Tensor& scaling, LoraTpLayout layout); static at::Tensor add_batched_lora( TrainingContext* ctx, const at::Tensor& base, const at::Tensor& input, @@ -1136,7 +1143,8 @@ static at::Tensor add_batched_lora( ) { if (!entry) return base; return base + lora_activation_delta( - ctx, input, entry->a_stack, entry->b_stack, entry->scaling); + ctx, input, entry->a_stack, entry->b_stack, entry->scaling, + entry->layout); } // Per-token routed-expert LoRA. Dynamic adapters add a leading sample axis to @@ -1725,6 +1733,11 @@ static at::Tensor tp_allreduce_base_mlp( TrainingContext* ctx, const at::Tensor& local_output); static at::Tensor tp_copy_base_mlp_input( TrainingContext* ctx, const at::Tensor& input); +static at::Tensor tp_allreduce_base_attention( + TrainingContext* ctx, const at::Tensor& local_output); +static at::Tensor tp_copy_base_attention_input( + TrainingContext* ctx, const at::Tensor& input); +static bool base_tp_attention_enabled(const TrainingContext* ctx); static at::Tensor forward_single_layer( TrainingContext* ctx, const at::Tensor& hidden, at::Tensor** w, const LayerConfig* cfg, @@ -1740,6 +1753,8 @@ static at::Tensor forward_single_layer( if (cfg->layer_type == 0) { // Full attention auto q_proj = *w[2], q_norm = *w[3], k_proj = *w[4], k_norm = *w[5], v_proj = *w[6], o_proj = *w[7]; + TORCH_CHECK(!base_tp_attention_enabled(ctx) || use_batched, + "base full-attention TP requires the activation-level LoRA path"); if (use_batched) { // Activation-level LoRA: pass base weights, apply delta inside attention attn_output = full_attention_batched( @@ -2005,6 +2020,10 @@ struct TrainingContext { // input-sharded by the Rust weight loader. The local row contribution // is reduced over the TP communicator before the residual add. bool base_tp_mlp = false; + // Frozen full-attention TP: Q/K/V own disjoint head bundles and O owns + // the matching input columns. GDN layers remain replicated until their + // state/head bundle partition has a dedicated implementation. + bool base_tp_attention = false; // Set when a legacy NCCL setter supplies an incompatible mixed topology. // Training entry points reject the context before touching parameters. bool topology_invalid = false; @@ -2137,6 +2156,22 @@ static at::Tensor tp_allreduce_lora_delta( (int64_t)reinterpret_cast(ctx->tp_stream)); } +static at::Tensor tp_copy_lora_input( + TrainingContext* ctx, const at::Tensor& input +) { + if (!ctx || ctx->tp_world_size <= 1) return input; + TORCH_CHECK(ctx->tp_comm, + "LoRA TP communicator is not initialized for TP_SIZE=", + ctx->tp_world_size); + return TpCopyToRegionFunction::apply( + input, (int64_t)ctx->tp_comm, + (int64_t)reinterpret_cast(ctx->tp_stream)); +} + +static bool base_tp_attention_enabled(const TrainingContext* ctx) { + return ctx && ctx->base_tp_attention && ctx->tp_world_size > 1; +} + static at::Tensor tp_allreduce_base_mlp( TrainingContext* ctx, const at::Tensor& local_output ) { @@ -2163,6 +2198,50 @@ static at::Tensor tp_copy_base_mlp_input( (int64_t)reinterpret_cast(ctx->tp_stream)); } +static at::Tensor tp_allreduce_base_attention( + TrainingContext* ctx, const at::Tensor& local_output +) { + if (!ctx || !ctx->base_tp_attention || ctx->tp_world_size <= 1) + return local_output; + TORCH_CHECK(ctx->tp_comm, + "base attention TP communicator is not initialized for TP_SIZE=", + ctx->tp_world_size); + return NcclAllReduceFunction::apply( + local_output, (int64_t)ctx->tp_comm, + (int64_t)reinterpret_cast(ctx->tp_stream)); +} + +static at::Tensor tp_copy_base_attention_input( + TrainingContext* ctx, const at::Tensor& input +) { + if (!ctx || !ctx->base_tp_attention || ctx->tp_world_size <= 1) + return input; + TORCH_CHECK(ctx->tp_comm, + "base attention TP communicator is not initialized for TP_SIZE=", + ctx->tp_world_size); + return TpCopyToRegionFunction::apply( + input, (int64_t)ctx->tp_comm, + (int64_t)reinterpret_cast(ctx->tp_stream)); +} + +static LoraTpLayout lora_tp_layout( + const TrainingContext* ctx, int64_t layer_idx, int64_t pair_idx +) { + if (!ctx || !ctx->base_tp_attention || ctx->tp_world_size <= 1 || + layer_idx < 0 || layer_idx >= ctx->num_layers) + return LoraTpLayout::LatentRank; + const auto& cfg = ctx->layer_configs[layer_idx]; + if (cfg.layer_type != 0) return LoraTpLayout::LatentRank; + auto table = lora_projection_table(cfg); + TORCH_CHECK(pair_idx >= 0 && pair_idx < table.count, + "invalid LoRA pair for TP layout"); + const std::string name(table.entries[pair_idx].name); + if (name == "q_proj" || name == "k_proj" || name == "v_proj") + return LoraTpLayout::ColumnParallel; + if (name == "o_proj") return LoraTpLayout::RowParallel; + return LoraTpLayout::LatentRank; +} + static at::Tensor initialize_lora_a( TrainingContext* ctx, const at::TensorOptions& options, int64_t experts, int64_t global_rank, int64_t in_features @@ -2227,6 +2306,80 @@ static ncclDataType_t nccl_dtype_for(const at::Tensor& tensor) { } } +static void tp_broadcast_lora_parameter( + TrainingContext* ctx, at::Tensor& tensor +) { + if (!ctx || ctx->tp_world_size <= 1 || !tensor.defined()) return; + TORCH_CHECK(ctx->tp_comm, + "LoRA TP communicator is not initialized for parameter broadcast"); + TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous(), + "LoRA TP parameter broadcast requires a contiguous CUDA tensor"); + const int dev = tensor.device().index(); + cudaSetDevice(dev); + auto stream = c10::cuda::getCurrentCUDAStream(dev).stream(); + auto err = ncclBroadcast( + tensor.data_ptr(), tensor.data_ptr(), tensor.numel(), + nccl_dtype_for(tensor), 0, ctx->tp_comm, stream); + TORCH_CHECK(err == ncclSuccess, + "NCCL LoRA parameter broadcast failed: ", ncclGetErrorString(err)); +} + +static void synchronize_adapter_replicated_lora_parameters( + TrainingContext* ctx, TrainingContext::LoRAAdapter& adapter); + +static void synchronize_fixed_replicated_lora_parameters(TrainingContext* ctx) { + if (!ctx || !ctx->base_tp_attention || ctx->tp_world_size <= 1) return; + for (int64_t layer = 0; layer < ctx->num_layers; ++layer) { + const int64_t offset = ctx->lora_layer_offset[layer]; + const int64_t pairs = lora_pair_count(ctx->layer_configs[layer]); + for (int64_t pair = 0; pair < pairs; ++pair) { + const int64_t slot = offset + pair; + if (!legacy_lora_slot_active(ctx, slot)) continue; + const auto layout = lora_tp_layout(ctx, layer, pair); + if (layout == LoraTpLayout::ColumnParallel) + tp_broadcast_lora_parameter(ctx, ctx->lora_a[slot]); + else if (layout == LoraTpLayout::RowParallel) + tp_broadcast_lora_parameter(ctx, ctx->lora_b[slot]); + } + } + for (auto& adapter : ctx->adapters) + synchronize_adapter_replicated_lora_parameters(ctx, adapter); +} + +static void synchronize_adapter_replicated_lora_parameters( + TrainingContext* ctx, TrainingContext::LoRAAdapter& adapter +) { + if (!ctx || !ctx->base_tp_attention || ctx->tp_world_size <= 1) return; + if (!ctx->tp_comm) return; // qwen36_init_nccl synchronizes deferred adapters. + for (auto& [layer, pairs] : adapter.params) { + for (int64_t pair = 0; pair < static_cast(pairs.size()); ++pair) { + auto& [a, b] = pairs[pair]; + if (!a.requires_grad() && !b.requires_grad()) continue; + const auto layout = lora_tp_layout(ctx, layer, pair); + if (layout == LoraTpLayout::ColumnParallel) + tp_broadcast_lora_parameter(ctx, a); + else if (layout == LoraTpLayout::RowParallel) + tp_broadcast_lora_parameter(ctx, b); + } + } +} + +static void tp_sum_replicated_lora_accumulator( + TrainingContext* ctx, at::Tensor& accumulator, + LoraTpLayout layout, bool is_a +) { + const bool replicated = + (layout == LoraTpLayout::ColumnParallel && is_a) || + (layout == LoraTpLayout::RowParallel && !is_a); + if (!replicated || !accumulator.defined() || ctx->tp_world_size <= 1) return; + TORCH_CHECK(ctx->tp_comm, + "LoRA TP communicator is not initialized for replicated gradient sum"); + auto reduced = NcclAllReduceFunction::allreduce( + accumulator, ctx->tp_comm, ctx->tp_stream); + at::NoGradGuard guard; + accumulator.copy_(reduced); +} + static void reduce_lora_accumulator( TrainingContext* ctx, at::Tensor& accumulator, double scale, bool allreduce @@ -2401,20 +2554,23 @@ static void synchronize_lora_gradients( // Dynamic multi-LoRA currently contributes one independently-normalized // row per tenant. Preserve that contract while weighting replicated DP // ranks by the selected batch's token count. - if (!dp_allreduce) return; - auto shifted_mask = target_mask.narrow(1, 1, target_mask.size(1) - 1) - .to(at::kFloat).sum().reshape({1}); - auto global_mask = at::empty_like(shifted_mask); - auto stream = c10::cuda::getCurrentCUDAStream( - shifted_mask.device().index()).stream(); - auto err = ncclAllReduce( - shifted_mask.data_ptr(), global_mask.data_ptr(), 1, - ncclFloat, ncclSum, ctx->nccl_comm, stream); - TORCH_CHECK(err == ncclSuccess, "NCCL token-count all-reduce failed: ", - ncclGetErrorString(err)); - const double local_tokens = shifted_mask.item(); - const double global_tokens = global_mask.item(); - scale = local_tokens / std::max(global_tokens, 1.0); + if (!dp_allreduce) { + if (!ctx->base_tp_attention) return; + } else { + auto shifted_mask = target_mask.narrow(1, 1, target_mask.size(1) - 1) + .to(at::kFloat).sum().reshape({1}); + auto global_mask = at::empty_like(shifted_mask); + auto stream = c10::cuda::getCurrentCUDAStream( + shifted_mask.device().index()).stream(); + auto err = ncclAllReduce( + shifted_mask.data_ptr(), global_mask.data_ptr(), 1, + ncclFloat, ncclSum, ctx->nccl_comm, stream); + TORCH_CHECK(err == ncclSuccess, "NCCL token-count all-reduce failed: ", + ncclGetErrorString(err)); + const double local_tokens = shifted_mask.item(); + const double global_tokens = global_mask.item(); + scale = local_tokens / std::max(global_tokens, 1.0); + } } for (size_t adapter_index = 0; adapter_index < ctx->adapters.size(); ++adapter_index) { @@ -2483,6 +2639,20 @@ static void synchronize_lora_gradients( } } } + // Projection-aware TP keeps one LoRA factor replicated. Its gradient is + // the sum of disjoint output-head (column) or input-column (row) + // contributions and is synchronized once at the optimizer boundary. + for (auto& adapter : ctx->adapters) { + for (auto& [layer_idx, pairs] : adapter.grad_accum) { + for (int64_t pair = 0; pair < static_cast(pairs.size()); ++pair) { + const auto layout = lora_tp_layout(ctx, layer_idx, pair); + tp_sum_replicated_lora_accumulator( + ctx, pairs[pair][0], layout, true); + tp_sum_replicated_lora_accumulator( + ctx, pairs[pair][1], layout, false); + } + } + } if (per_adapter_weighting) return; // The replicated-source A2A path sends an identical token batch from every // EP rank. Average only its sharded expert parameter gradients here; @@ -2519,6 +2689,17 @@ static void synchronize_lora_gradients( dp_allreduce || sharded_a2a); } } + for (int64_t layer = 0; layer < ctx->num_layers; ++layer) { + const int64_t offset = ctx->lora_layer_offset[layer]; + const int64_t pairs = lora_pair_count(ctx->layer_configs[layer]); + for (int64_t pair = 0; pair < pairs; ++pair) { + const auto layout = lora_tp_layout(ctx, layer, pair); + tp_sum_replicated_lora_accumulator( + ctx, ctx->grad_accum_a[offset + pair], layout, true); + tp_sum_replicated_lora_accumulator( + ctx, ctx->grad_accum_b[offset + pair], layout, false); + } + } } static void elide_trivial_attention_mask(TrainingContext* ctx) { @@ -2700,7 +2881,8 @@ static void prepare_lora_batch(TrainingContext* ctx) { auto scaling = scaling_cpu.to(a_stack.device()).to(at::kBFloat16); // [N, 1, 1] ctx->lora_batch_cache[lora_cache_key(layer_idx, pair_idx)] = { - a_stack, b_stack, scaling + a_stack, b_stack, scaling, + lora_tp_layout(ctx, layer_idx, pair_idx) }; } } @@ -2725,7 +2907,8 @@ static void prepare_fixed_lora_batch(TrainingContext* ctx) { {1, 1, 1}, ctx->lora_scaling, at::TensorOptions().dtype(a.scalar_type()).device(a.device())); ctx->lora_batch_cache[lora_cache_key(layer_idx, pair_idx)] = { - a.unsqueeze(0), b.unsqueeze(0), scaling}; + a.unsqueeze(0), b.unsqueeze(0), scaling, + lora_tp_layout(ctx, layer_idx, pair_idx)}; } } ctx->lora_batch_valid = true; @@ -2739,7 +2922,8 @@ static at::Tensor lora_activation_delta( const at::Tensor& x, // [N, seq, in] const at::Tensor& A, // [N, rank, in] const at::Tensor& B, // [N, out, rank] - const at::Tensor& scaling // [N, 1, 1] + const at::Tensor& scaling, // [N, 1, 1] + LoraTpLayout layout ) { // Cast to compute dtype (BF16) auto kind = x.scalar_type(); @@ -2751,11 +2935,19 @@ static at::Tensor lora_activation_delta( B_c = B_c.expand({x.size(0), B_c.size(1), B_c.size(2)}); s_c = s_c.expand({x.size(0), 1, 1}); } + // Latent-rank TP sums local deltas in forward. Its replicated input must + // likewise sum the rank-local input-gradient contributions in backward; + // otherwise a later sharded LoRA branch feeds only a partial dgrad into + // preceding replicated layers. + auto lora_input = layout == LoraTpLayout::LatentRank + ? tp_copy_lora_input(ctx, x) : x; // Ax = A @ x^T → [N, rank, seq] - auto Ax = at::bmm(A_c, x.transpose(-2, -1)); + auto Ax = at::bmm(A_c, lora_input.transpose(-2, -1)); // delta = B @ Ax → [N, out, seq] → transpose → [N, seq, out] auto delta = at::bmm(B_c, Ax).transpose(-2, -1); - return tp_allreduce_lora_delta(ctx, delta * s_c); + auto scaled = delta * s_c; + return layout == LoraTpLayout::LatentRank + ? tp_allreduce_lora_delta(ctx, scaled) : scaled; } static const LoraBatchEntry* lora_batch_entry( @@ -2846,6 +3038,8 @@ at::Tensor compute_attn_only( } // Legacy path: weight-level LoRA + TORCH_CHECK(!ctx->base_tp_attention || cfg.layer_type != 0, + "base full-attention TP requires the activation-level LoRA path"); int64_t lora_count = lora_pair_count(cfg); int64_t la_offset = ctx->lora_layer_offset[layer_idx]; bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size(); @@ -3099,24 +3293,35 @@ static at::Tensor full_attention_batched( ) { // Compute Q/K/V with base weight, then add LoRA delta if present int64_t batch = hidden.size(0), seq = hidden.size(1); + auto projection_input = tp_copy_base_attention_input(ctx, hidden); + if (ctx->base_tp_attention) { + TORCH_CHECK(num_heads % ctx->tp_world_size == 0 && + num_kv_heads % ctx->tp_world_size == 0, + "full-attention heads must be divisible by TP_SIZE"); + num_heads /= ctx->tp_world_size; + num_kv_heads /= ctx->tp_world_size; + } int64_t qkv_dim = num_heads * head_dim; - auto q = at::matmul(hidden, q_proj.t()); - auto k = at::matmul(hidden, k_proj.t()); - auto v = at::matmul(hidden, v_proj.t()); + auto q = at::matmul(projection_input, q_proj.t()); + auto k = at::matmul(projection_input, k_proj.t()); + auto v = at::matmul(projection_input, v_proj.t()); // Apply activation-level LoRA: q += B@(A@hidden) * scaling auto it_q = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 0)); if (it_q != ctx->lora_batch_cache.end()) { - q = q + lora_activation_delta(ctx, hidden, it_q->second.a_stack, it_q->second.b_stack, it_q->second.scaling); + q = q + lora_activation_delta(ctx, projection_input, it_q->second.a_stack, + it_q->second.b_stack, it_q->second.scaling, it_q->second.layout); } auto it_k = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 1)); if (it_k != ctx->lora_batch_cache.end()) { - k = k + lora_activation_delta(ctx, hidden, it_k->second.a_stack, it_k->second.b_stack, it_k->second.scaling); + k = k + lora_activation_delta(ctx, projection_input, it_k->second.a_stack, + it_k->second.b_stack, it_k->second.scaling, it_k->second.layout); } auto it_v = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 2)); if (it_v != ctx->lora_batch_cache.end()) { - v = v + lora_activation_delta(ctx, hidden, it_v->second.a_stack, it_v->second.b_stack, it_v->second.scaling); + v = v + lora_activation_delta(ctx, projection_input, it_v->second.a_stack, + it_v->second.b_stack, it_v->second.scaling, it_v->second.layout); } // Reshape Q: [batch, seq, num_heads, head_dim*2] → split into q and gate @@ -3186,9 +3391,10 @@ static at::Tensor full_attention_batched( auto it_o = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 3)); if (it_o != ctx->lora_batch_cache.end()) { result = result + lora_activation_delta(ctx, attn_flat, - it_o->second.a_stack, it_o->second.b_stack, it_o->second.scaling); + it_o->second.a_stack, it_o->second.b_stack, it_o->second.scaling, + it_o->second.layout); } - return result; + return tp_allreduce_base_attention(ctx, result); } static at::Tensor linear_attention_batched( @@ -3212,7 +3418,8 @@ static at::Tensor linear_attention_batched( auto qkv = at::matmul(hidden, in_proj_qkv.t()); auto it_qkv = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 0)); if (it_qkv != ctx->lora_batch_cache.end()) { - qkv = qkv + lora_activation_delta(ctx, hidden, it_qkv->second.a_stack, it_qkv->second.b_stack, it_qkv->second.scaling); + qkv = qkv + lora_activation_delta(ctx, hidden, it_qkv->second.a_stack, + it_qkv->second.b_stack, it_qkv->second.scaling, it_qkv->second.layout); } // DIAG: dump after QKV projection @@ -3284,18 +3491,21 @@ static at::Tensor linear_attention_batched( auto b = at::matmul(hidden, in_proj_b.t()); auto it_a = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 2)); if (it_a != ctx->lora_batch_cache.end()) { - a = a + lora_activation_delta(ctx, hidden, it_a->second.a_stack, it_a->second.b_stack, it_a->second.scaling); + a = a + lora_activation_delta(ctx, hidden, it_a->second.a_stack, + it_a->second.b_stack, it_a->second.scaling, it_a->second.layout); } auto it_b = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 3)); if (it_b != ctx->lora_batch_cache.end()) { - b = b + lora_activation_delta(ctx, hidden, it_b->second.a_stack, it_b->second.b_stack, it_b->second.scaling); + b = b + lora_activation_delta(ctx, hidden, it_b->second.a_stack, + it_b->second.b_stack, it_b->second.scaling, it_b->second.layout); } // Z projection + LoRA delta auto z = at::matmul(hidden, in_proj_z.t()); auto it_z = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 1)); if (it_z != ctx->lora_batch_cache.end()) { - z = z + lora_activation_delta(ctx, hidden, it_z->second.a_stack, it_z->second.b_stack, it_z->second.scaling); + z = z + lora_activation_delta(ctx, hidden, it_z->second.a_stack, + it_z->second.b_stack, it_z->second.scaling, it_z->second.layout); } z = z.reshape({batch, seq, num_v_heads, head_v_dim}); @@ -3407,7 +3617,8 @@ static at::Tensor linear_attention_batched( // out_proj LoRA delta: result += B@(A@gated) * scaling auto it_op = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 4)); if (it_op != ctx->lora_batch_cache.end()) { - result = result + lora_activation_delta(ctx, gated, it_op->second.a_stack, it_op->second.b_stack, it_op->second.scaling); + result = result + lora_activation_delta(ctx, gated, it_op->second.a_stack, + it_op->second.b_stack, it_op->second.scaling, it_op->second.layout); } return result; @@ -4367,14 +4578,16 @@ static at::Tensor mtp_compute_loss( extern "C" { __attribute__((visibility("default"))) int64_t qwen36_kernel_abi_version() { - return 13; + return 14; } +static constexpr int32_t QWEN36_CONTEXT_BASE_TP_ATTENTION = 1 << 0; + // Create training context — called once at startup // lora_rank: LoRA rank (from config) // target_layers: array of layer indices to apply LoRA (nullptr = all layers) // num_target_layers: length of target_layers array -__attribute__((visibility("default"))) void* qwen36_create_training_context( +static void* qwen36_create_training_context_impl( void** weight_ptrs, int64_t num_weight_ptrs, void* embed_ptr, void* final_norm_ptr, void* lm_head_ptr, void* layer_configs_ptr, int64_t num_layers, @@ -4383,7 +4596,8 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( int64_t vocab_size, double rms_eps, int64_t lora_rank, const int64_t* target_layers, int64_t num_target_layers, - const char* target_modules_str + const char* target_modules_str, + int32_t context_flags ) { try { auto* ctx = new TrainingContext(); @@ -4393,6 +4607,8 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( ctx->vocab_size = vocab_size; ctx->rms_eps = rms_eps; ctx->step_count = 0; ctx->lora_scaling = lora_scaling; ctx->num_layers = num_layers; + ctx->base_tp_attention = + (context_flags & QWEN36_CONTEXT_BASE_TP_ATTENTION) != 0; ctx->has_mtp = false; ctx->use_checkpoint = false; ctx->group_size = 4; const char* tp_size_env = getenv("TP_SIZE"); @@ -4449,6 +4665,52 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( ctx->layer_configs.push_back(lcfgs[i]); } + if (ctx->base_tp_attention) { + TORCH_CHECK(ctx->tp_world_size > 1, + "base full-attention TP requires TP_SIZE>1"); + int64_t weight_offset = 0; + for (int64_t layer = 0; layer < num_layers; ++layer) { + const auto& cfg = ctx->layer_configs[layer]; + if (cfg.layer_type == 0) { + TORCH_CHECK(cfg.num_heads > 0 && cfg.num_kv_heads > 0 && + cfg.head_dim > 0, + "invalid full-attention head configuration at layer ", layer); + TORCH_CHECK(cfg.num_heads % cfg.num_kv_heads == 0 && + cfg.num_heads % ctx->tp_world_size == 0 && + cfg.num_kv_heads % ctx->tp_world_size == 0, + "full-attention heads must preserve GQA groups and be divisible by TP_SIZE at layer ", + layer); + const int64_t rotary_dim = static_cast( + cfg.head_dim * cfg.partial_rotary_factor); + TORCH_CHECK(rotary_dim >= 0 && rotary_dim <= cfg.head_dim && + rotary_dim % 2 == 0, + "full-attention rotary dimension must be even and within head_dim at layer ", + layer, ": rotary_dim=", rotary_dim, + " head_dim=", cfg.head_dim); + auto* q = ctx->weight_ptrs[weight_offset + 2]; + auto* k = ctx->weight_ptrs[weight_offset + 4]; + auto* v = ctx->weight_ptrs[weight_offset + 6]; + auto* o = ctx->weight_ptrs[weight_offset + 7]; + TORCH_CHECK(q && k && v && o && q->dim() == 2 && + k->dim() == 2 && v->dim() == 2 && o->dim() == 2, + "base full-attention TP requires matrix Q/K/V/O weights at layer ", + layer); + const int64_t local_heads = cfg.num_heads / ctx->tp_world_size; + const int64_t local_kv_heads = cfg.num_kv_heads / ctx->tp_world_size; + TORCH_CHECK(q->size(0) == local_heads * cfg.head_dim * 2 && + k->size(0) == local_kv_heads * cfg.head_dim && + v->size(0) == local_kv_heads * cfg.head_dim && + o->size(1) == local_heads * cfg.head_dim && + q->size(1) == k->size(1) && q->size(1) == v->size(1) && + q->size(1) == o->size(0), + "base full-attention TP received inconsistent local weight shapes at layer ", + layer, ": q=", q->sizes(), " k=", k->sizes(), + " v=", v->sizes(), " o=", o->sizes()); + } + weight_offset += weight_count_for_layer(cfg); + } + } + // Build target layer set std::set target_set; const bool all_target_layers = !target_layers || num_target_layers == 0; @@ -4539,8 +4801,17 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( b = at::zeros({experts, out_f, local_lora_rank}, opts); } else { int64_t out_f = base->size(0), in_f = base->size(1); - a = initialize_lora_a(ctx, opts, 0, lora_rank, in_f); - b = at::zeros({out_f, local_lora_rank}, opts); + const auto layout = lora_tp_layout(ctx, i, k); + if (layout == LoraTpLayout::ColumnParallel) { + a = at::randn({lora_rank, in_f}, opts) * 0.01; + b = at::zeros({out_f, lora_rank}, opts); + } else if (layout == LoraTpLayout::RowParallel) { + a = at::randn({lora_rank, in_f}, opts) * 0.01; + b = at::zeros({out_f, lora_rank}, opts); + } else { + a = initialize_lora_a(ctx, opts, 0, lora_rank, in_f); + b = at::zeros({out_f, local_lora_rank}, opts); + } } a.set_requires_grad(active); b.set_requires_grad(active); @@ -4577,6 +4848,42 @@ __attribute__((visibility("default"))) void* qwen36_create_training_context( } } +__attribute__((visibility("default"))) void* qwen36_create_training_context( + void** weight_ptrs, int64_t num_weight_ptrs, + void* embed_ptr, void* final_norm_ptr, void* lm_head_ptr, + void* layer_configs_ptr, int64_t num_layers, + int32_t compute_type, + double lora_scaling, double lr, double beta1, double beta2, double eps, + int64_t vocab_size, double rms_eps, + int64_t lora_rank, + const int64_t* target_layers, int64_t num_target_layers, + const char* target_modules_str +) { + return qwen36_create_training_context_impl( + weight_ptrs, num_weight_ptrs, embed_ptr, final_norm_ptr, lm_head_ptr, + layer_configs_ptr, num_layers, compute_type, lora_scaling, lr, beta1, + beta2, eps, vocab_size, rms_eps, lora_rank, target_layers, + num_target_layers, target_modules_str, 0); +} + +__attribute__((visibility("default"))) void* qwen36_create_training_context_ex( + void** weight_ptrs, int64_t num_weight_ptrs, + void* embed_ptr, void* final_norm_ptr, void* lm_head_ptr, + void* layer_configs_ptr, int64_t num_layers, + int32_t compute_type, + double lora_scaling, double lr, double beta1, double beta2, double eps, + int64_t vocab_size, double rms_eps, + int64_t lora_rank, + const int64_t* target_layers, int64_t num_target_layers, + const char* target_modules_str, int32_t context_flags +) { + return qwen36_create_training_context_impl( + weight_ptrs, num_weight_ptrs, embed_ptr, final_norm_ptr, lm_head_ptr, + layer_configs_ptr, num_layers, compute_type, lora_scaling, lr, beta1, + beta2, eps, vocab_size, rms_eps, lora_rank, target_layers, + num_target_layers, target_modules_str, context_flags); +} + __attribute__((visibility("default"))) int32_t qwen36_set_base_tp_mlp( void* ctx_ptr, int32_t enabled ) { @@ -4665,8 +4972,8 @@ __attribute__((visibility("default"))) int32_t qwen36_set_mtp_weights( try { auto* ctx = reinterpret_cast(ctx_ptr); TORCH_CHECK(ctx, "null training context"); - TORCH_CHECK(!ctx->base_tp_mlp, - "MTP cannot be enabled after base dense MLP TP because MTP weights are not sharded"); + TORCH_CHECK(!ctx->base_tp_mlp && !ctx->base_tp_attention, + "MTP cannot be enabled after frozen base TP because MTP weights are not sharded"); TORCH_CHECK(!ctx->has_mtp, "MTP weights are already configured"); ctx->has_mtp = true; ctx->mtp_fc = reinterpret_cast(mtp_fc_ptr); @@ -5634,6 +5941,7 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( lc.nccl_stream = layer_stream; } } + synchronize_fixed_replicated_lora_parameters(ctx); return 0; } @@ -5799,6 +6107,7 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( } } + synchronize_fixed_replicated_lora_parameters(ctx); return 0; } @@ -5992,8 +6301,15 @@ int64_t qwen36_add_lora( TORCH_CHECK(base->dim() == 2, "dynamic LoRA projection must be a matrix: ", projection.name); int64_t out_f = base->size(0), in_f = base->size(1); - a = initialize_lora_a(ctx, opts, 0, rank, in_f); - b = at::zeros({out_f, local_rank}, opts); + const auto layout = lora_tp_layout(ctx, i, k); + if (layout == LoraTpLayout::ColumnParallel || + layout == LoraTpLayout::RowParallel) { + a = at::randn({rank, in_f}, opts) * 0.01; + b = at::zeros({out_f, rank}, opts); + } else { + a = initialize_lora_a(ctx, opts, 0, rank, in_f); + b = at::zeros({out_f, local_rank}, opts); + } } } else { a = at::zeros({}, opts); @@ -6018,6 +6334,7 @@ int64_t qwen36_add_lora( adapter.adam_state[i] = std::move(adam_states); adapter.grad_accum[i] = std::move(grad_accumulators); } + synchronize_adapter_replicated_lora_parameters(ctx, adapter); int64_t id = adapter.id; ctx->adapters.push_back(std::move(adapter)); ctx->lora_cache_valid = false; diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs index 7bdc1cc5..9bf5a9e9 100644 --- a/crates/rustrain-qwen3-6/src/kernel.rs +++ b/crates/rustrain-qwen3-6/src/kernel.rs @@ -4,7 +4,7 @@ //! Rust only handles: weight loading, data loading, training loop orchestration. use crate::lora::Qwen36LoraTargetModule; -use anyhow::{Result, bail}; +use anyhow::{bail, Result}; use std::ffi::c_void; use std::sync::OnceLock; use tch::{Kind, Tensor}; @@ -31,6 +31,7 @@ type FnCreateCtx = unsafe extern "C" fn( *const i64, i64, *const i8, + i32, ) -> *mut c_void; type FnKernelAbiVersion = unsafe extern "C" fn() -> i64; type FnTrainStep = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, *mut c_void) -> f64; @@ -196,11 +197,11 @@ unsafe fn load_kernels() -> Option { }}; } let abi_version: FnKernelAbiVersion = sym!("qwen36_kernel_abi_version"); - if abi_version() != 13 { + if abi_version() != 14 { return None; } Some(KernelHandles { - create_ctx: sym!("qwen36_create_training_context"), + create_ctx: sym!("qwen36_create_training_context_ex"), train_step: sym!("qwen36_train_step"), train_micro_step: sym!("qwen36_train_micro_step"), train_multi_lora: sym!("qwen36_train_multi_lora"), @@ -287,6 +288,44 @@ pub fn shard_dense_mlp_weight_for_tp( Ok(Some(tensor.narrow(dim, start, shard).contiguous())) } +/// Return the rank-local frozen full-attention shard for TP. Q/K/V own +/// contiguous output-head bundles; O owns the matching input columns. +pub fn shard_full_attention_weight_for_tp( + name: &str, + tensor: &Tensor, + tp_size: usize, + tp_rank: usize, +) -> Result> { + let dim = if name.ends_with(".self_attn.q_proj.weight") + || name.ends_with(".self_attn.k_proj.weight") + || name.ends_with(".self_attn.v_proj.weight") + { + 0 + } else if name.ends_with(".self_attn.o_proj.weight") { + 1 + } else { + return Ok(None); + }; + if tp_size <= 1 || tp_rank >= tp_size { + bail!("invalid full-attention TP shard: tp_rank={tp_rank}, tp_size={tp_size}"); + } + let full = *tensor + .size() + .get(dim as usize) + .ok_or_else(|| anyhow::anyhow!("base TP attention weight {name} has no dimension {dim}"))?; + if full <= 0 || full % tp_size as i64 != 0 { + bail!( + "base TP attention weight {name} dimension {dim}={full} is not divisible by TP_SIZE={tp_size}" + ); + } + let shard = full / tp_size as i64; + Ok(Some( + tensor + .narrow(dim, tp_rank as i64 * shard, shard) + .contiguous(), + )) +} + pub fn build_weight_ptrs( weights: &std::collections::BTreeMap, config: &crate::config::Qwen36RuntimeConfig, @@ -503,6 +542,7 @@ impl CppTrainingContext { eps: f64, lora_scaling: f64, lora_rank: i64, + base_tp_attention: bool, base_tp_mlp: bool, target_layers: &[usize], target_modules: &[Qwen36LoraTargetModule], @@ -579,6 +619,7 @@ impl CppTrainingContext { tl_ptr, tl_len, modules_ptr, + i32::from(base_tp_attention), ) }; if ptr.is_null() { @@ -1150,7 +1191,7 @@ impl Drop for CppTrainingContext { #[cfg(test)] mod tests { - use super::shard_dense_mlp_weight_for_tp; + use super::{shard_dense_mlp_weight_for_tp, shard_full_attention_weight_for_tp}; use tch::{Kind, Tensor}; #[test] @@ -1174,14 +1215,55 @@ mod tests { #[test] fn dense_mlp_tp_ignores_non_mlp_weights_and_rejects_bad_shapes() { let tensor = Tensor::zeros([5, 4], (Kind::Float, tch::Device::Cpu)); - assert!( - shard_dense_mlp_weight_for_tp("model.layers.0.self_attn.q_proj.weight", &tensor, 2, 0) - .unwrap() - .is_none() - ); + assert!(shard_dense_mlp_weight_for_tp( + "model.layers.0.self_attn.q_proj.weight", + &tensor, + 2, + 0 + ) + .unwrap() + .is_none()); assert!( shard_dense_mlp_weight_for_tp("model.layers.0.mlp.up_proj.weight", &tensor, 2, 0) .is_err() ); } + + #[test] + fn full_attention_tp_shards_head_and_output_axes() { + let q = Tensor::arange(96, (Kind::Float, tch::Device::Cpu)).reshape([24, 4]); + let o = Tensor::arange(48, (Kind::Float, tch::Device::Cpu)).reshape([4, 12]); + let q_rank_one = + shard_full_attention_weight_for_tp("model.layers.0.self_attn.q_proj.weight", &q, 2, 1) + .unwrap() + .unwrap(); + let o_rank_one = + shard_full_attention_weight_for_tp("model.layers.0.self_attn.o_proj.weight", &o, 2, 1) + .unwrap() + .unwrap(); + assert_eq!(q_rank_one.size(), [12, 4]); + assert_eq!(o_rank_one.size(), [4, 6]); + assert_eq!(q_rank_one.double_value(&[0, 0]), 48.0); + assert_eq!(o_rank_one.double_value(&[0, 0]), 6.0); + } + + #[test] + fn full_attention_tp_ignores_other_weights_and_rejects_bad_shapes() { + let tensor = Tensor::zeros([5, 4], (Kind::Float, tch::Device::Cpu)); + assert!(shard_full_attention_weight_for_tp( + "model.layers.0.mlp.gate_proj.weight", + &tensor, + 2, + 0, + ) + .unwrap() + .is_none()); + assert!(shard_full_attention_weight_for_tp( + "model.layers.0.self_attn.q_proj.weight", + &tensor, + 2, + 0, + ) + .is_err()); + } } diff --git a/crates/rustrain-qwen3-6/src/session.rs b/crates/rustrain-qwen3-6/src/session.rs index 5b69981d..8a9767f5 100644 --- a/crates/rustrain-qwen3-6/src/session.rs +++ b/crates/rustrain-qwen3-6/src/session.rs @@ -3,15 +3,15 @@ use std::collections::{BTreeMap, HashSet}; use std::env; -use anyhow::{Context, Result, anyhow, bail}; +use anyhow::{anyhow, bail, Context, Result}; use tch::{Kind, Tensor}; use tracing::info; use crate::config::{ - LayerType, Qwen36RuntimeConfig, read_qwen36_runtime_config, resolve_qwen36_model_path, + read_qwen36_runtime_config, resolve_qwen36_model_path, LayerType, Qwen36RuntimeConfig, }; use crate::lora::{ - Qwen36AdapterArtifact, Qwen36LoraConfig, Qwen36LoraTargetModule, validate_lora_targets, + validate_lora_targets, Qwen36AdapterArtifact, Qwen36LoraConfig, Qwen36LoraTargetModule, }; use crate::sft::SftDataset; use rustrain_checkpoint::safetensors::read_safetensors_dir_filtered; @@ -328,16 +328,42 @@ fn train_impl( ); } - // Dense base-weight TP follows Megatron's ColumnParallel/RowParallel MLP: - // gate/up rows are owned by one TP rank and down columns are owned by the - // same rank. Attention and embeddings remain replicated in this first - // unit. MoE, MTP, and mixed TP/DP/EP are rejected rather than silently - // running with an unsharded path. + // Full attention follows Megatron's Q/K/V ColumnParallel and O + // RowParallel layout. Dense MLP additionally shards gate/up rows and down + // columns. GDN, MoE experts, embeddings, and the LM head remain replicated. + let base_tp_attention = tp_size > 1; let base_tp_mlp = tp_size > 1 && !runtime_config.is_moe; - if base_tp_mlp { + if base_tp_attention { if runtime_config.mtp_num_hidden_layers > 0 { - bail!("base dense TP MLP currently requires MTP to be disabled"); + bail!("frozen base TP currently requires MTP to be disabled"); + } + if runtime_config.num_attention_heads <= 0 + || runtime_config.num_attention_heads % tp_size as i64 != 0 + || runtime_config.num_key_value_heads <= 0 + || runtime_config.num_key_value_heads % tp_size as i64 != 0 + || runtime_config.num_attention_heads % runtime_config.num_key_value_heads != 0 + { + bail!( + "full-attention heads (q={}, kv={}) must preserve GQA groups and be divisible by TP_SIZE={tp_size}", + runtime_config.num_attention_heads, + runtime_config.num_key_value_heads + ); + } + let rotary_dim = + (runtime_config.head_dim as f64 * runtime_config.partial_rotary_factor) as i64; + if runtime_config.head_dim <= 0 + || rotary_dim < 0 + || rotary_dim > runtime_config.head_dim + || rotary_dim % 2 != 0 + { + bail!( + "full-attention head_dim={} and partial_rotary_factor={} produce invalid rotary_dim={rotary_dim}", + runtime_config.head_dim, + runtime_config.partial_rotary_factor + ); } + } + if base_tp_mlp { if runtime_config.intermediate_size <= 0 || runtime_config.intermediate_size % tp_size as i64 != 0 { @@ -413,8 +439,23 @@ fn train_impl( ); } else { for (name, tensor) in &weights { - let local_shard = if base_tp_mlp { - crate::kernel::shard_dense_mlp_weight_for_tp(name, tensor, tp_size, rank % tp_size)? + let local_shard = if base_tp_attention { + let attention_shard = crate::kernel::shard_full_attention_weight_for_tp( + name, + tensor, + tp_size, + rank % tp_size, + )?; + if attention_shard.is_some() || !base_tp_mlp { + attention_shard + } else { + crate::kernel::shard_dense_mlp_weight_for_tp( + name, + tensor, + tp_size, + rank % tp_size, + )? + } } else { None }; @@ -425,11 +466,12 @@ fn train_impl( .to_kind(compute_kind); weights_gpu.insert(name.clone(), gpu_tensor); } - if base_tp_mlp { + if base_tp_attention { info!( tp_size, tp_rank = rank % tp_size, - "base dense MLP TP enabled: gate/up row shards and down column shards" + base_tp_mlp, + "frozen base TP enabled: full-attention head shards and optional dense MLP shards" ); } } @@ -488,6 +530,7 @@ fn train_impl( config.train.adam_eps as f64, lora_config.alpha as f64 / lora_config.rank as f64, // lora scaling = alpha / rank lora_config.rank as i64, + base_tp_attention, base_tp_mlp, &lora_config.target_layers, &lora_config.target_modules, diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index d3533a61..ed57a659 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -366,7 +366,7 @@ static int run_dynamic_dp_smoke( } int main() { - assert(qwen36_kernel_abi_version() == 13); + assert(qwen36_kernel_abi_version() == 14); const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); const int process_rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0"); const int local_rank = std::atoi(std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); diff --git a/crates/rustrain-qwen3-6/tests/native_tp_attention_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_attention_smoke.cpp new file mode 100644 index 00000000..569cb2d3 --- /dev/null +++ b/crates/rustrain-qwen3-6/tests/native_tp_attention_smoke.cpp @@ -0,0 +1,523 @@ +#include + +#include +#include +#include +#include +#include +#include + +struct LayerConfig { + int64_t layer_type, num_heads, num_kv_heads, head_dim; + int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel; + double partial_rotary_factor, rope_theta, rms_eps; + int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count; + int64_t intermediate_size; + int32_t norm_topk_prob; + void* nccl_comm; + void* nccl_stream; +}; + +extern "C" void qwen36_set_cuda_device(int32_t); +extern "C" void* qwen36_create_training_context( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*); +extern "C" void* qwen36_create_training_context_ex( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*, int32_t); +extern "C" int32_t qwen36_init_nccl(void*); +extern "C" int32_t qwen36_set_mtp_weights( + void*, void*, void*, void*, void*, void**, int64_t, void*, int64_t); +extern "C" int64_t qwen36_add_lora( + void*, int64_t, double, const int64_t*, int64_t, const char*); +extern "C" void* qwen36_get_adapter_lora_tensor( + void*, int64_t, int64_t, const char*, int32_t); +extern "C" int32_t qwen36_set_adapter_lora_tensor( + void*, int64_t, int64_t, const char*, int32_t, void*); +extern "C" void* qwen36_get_lora_a(void*, int64_t); +extern "C" void* qwen36_get_lora_b(void*, int64_t); +extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*); +extern "C" void* qwen36_get_lora_grad_accumulator(void*, int64_t, int32_t); +extern "C" int32_t qwen36_abort_gradient_accumulation(void*); +extern "C" int64_t qwen36_export_optimizer_state( + void*, void**, void**, int64_t); +extern "C" double qwen36_eval_step(void*, void*, void*, void*); +extern "C" double qwen36_train_step(void*, void*, void*, void*); +extern "C" double qwen36_train_micro_step( + void*, void*, void*, void*, double, int32_t); +extern "C" double qwen36_train_multi_lora_selected( + void*, void*, void*, void*, const int64_t*, int32_t, int32_t); +extern "C" void qwen36_free_training_context(void*); + +static at::Tensor deterministic(std::initializer_list shape, double scale) { + int64_t count = 1; + for (int64_t dim : shape) count *= dim; + return ((at::arange(count, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)) + .remainder(19) - 9.0) * scale) + .reshape(shape).to(at::kBFloat16); +} + +static std::vector pointers(std::vector& tensors) { + std::vector result; + result.reserve(tensors.size()); + for (auto& tensor : tensors) result.push_back(&tensor); + return result; +} + +static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) { + return (lhs - rhs).abs().max().item(); +} + +int main() { + const int rank = std::atoi(std::getenv("RANK")); + const int world = std::atoi(std::getenv("WORLD_SIZE")); + assert(world == 2 && rank >= 0 && rank < world); + qwen36_set_cuda_device(rank); + + constexpr int64_t hidden = 8; + constexpr int64_t heads = 4; + constexpr int64_t kv_heads = 2; + constexpr int64_t head_dim = 2; + constexpr int64_t intermediate = 12; + constexpr int64_t vocab = 16; + constexpr int64_t lora_rank = 4; + + std::vector full_weights; + full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + full_weights.push_back(at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + full_weights.push_back(deterministic({2 * heads * head_dim, hidden}, 0.010)); + full_weights.push_back(at::ones({head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + full_weights.push_back(deterministic({kv_heads * head_dim, hidden}, 0.012)); + full_weights.push_back(at::ones({head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + full_weights.push_back(deterministic({kv_heads * head_dim, hidden}, 0.008)); + full_weights.push_back(deterministic({hidden, heads * head_dim}, 0.011)); + full_weights.push_back(deterministic({intermediate, hidden}, 0.009)); + full_weights.push_back(deterministic({intermediate, hidden}, 0.007)); + full_weights.push_back(deterministic({hidden, intermediate}, 0.010)); + for (auto& weight : full_weights) weight.set_requires_grad(false); + + const int64_t local_heads = heads / world; + const int64_t local_kv_heads = kv_heads / world; + std::vector local_weights; + local_weights.push_back(full_weights[0]); + local_weights.push_back(full_weights[1]); + local_weights.push_back(full_weights[2].narrow( + 0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim).contiguous()); + local_weights.push_back(full_weights[3]); + local_weights.push_back(full_weights[4].narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim).contiguous()); + local_weights.push_back(full_weights[5]); + local_weights.push_back(full_weights[6].narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim).contiguous()); + local_weights.push_back(full_weights[7].narrow( + 1, rank * local_heads * head_dim, local_heads * head_dim).contiguous()); + local_weights.insert( + local_weights.end(), full_weights.begin() + 8, full_weights.end()); + + auto embed = deterministic({vocab, hidden}, 0.020); + auto final_norm = at::ones({hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)); + auto lm_head = deterministic({vocab, hidden}, 0.015); + LayerConfig config{}; + config.layer_type = 0; + config.num_heads = heads; + config.num_kv_heads = kv_heads; + config.head_dim = head_dim; + config.partial_rotary_factor = 1.0; + config.rope_theta = 10000.0; + config.rms_eps = 1e-5; + config.intermediate_size = intermediate; + + const int64_t target_layer = 0; + auto local_ptrs = pointers(local_weights); + setenv("TP_SIZE", "2", 1); + unsetenv("RUSTRAIN_DATA_PARALLEL"); + void* distributed = qwen36_create_training_context_ex( + local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head, + &config, 1, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + &target_layer, 1, "q_proj,k_proj,v_proj,o_proj", 1); + assert(distributed); + assert(qwen36_set_mtp_weights( + distributed, nullptr, nullptr, nullptr, nullptr, nullptr, 0, nullptr, 0) != 0); + assert(qwen36_init_nccl(distributed) == 0); + + setenv("TP_SIZE", "1", 1); + auto full_ptrs = pointers(full_weights); + void* reference = qwen36_create_training_context( + full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head, + &config, 1, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + &target_layer, 1, "q_proj,k_proj,v_proj,o_proj"); + assert(reference); + + auto q_a = deterministic({lora_rank, hidden}, 0.0020); + auto q_b = deterministic({2 * heads * head_dim, lora_rank}, 0.0010); + auto k_a = deterministic({lora_rank, hidden}, 0.0018); + auto k_b = deterministic({kv_heads * head_dim, lora_rank}, 0.0011); + auto v_a = deterministic({lora_rank, hidden}, 0.0016); + auto v_b = deterministic({kv_heads * head_dim, lora_rank}, 0.0009); + auto o_a = deterministic({lora_rank, heads * head_dim}, 0.0015); + auto o_b = deterministic({hidden, lora_rank}, 0.0012); + auto local_q_b = q_b.narrow( + 0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim).contiguous(); + auto local_o_a = o_a.narrow( + 1, rank * local_heads * head_dim, local_heads * head_dim).contiguous(); + auto local_k_b = k_b.narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim).contiguous(); + auto local_v_b = v_b.narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim).contiguous(); + assert(qwen36_set_lora_tensor(distributed, 0, 0, &q_a) == 0); + assert(qwen36_set_lora_tensor(distributed, 0, 1, &local_q_b) == 0); + assert(qwen36_set_lora_tensor(distributed, 1, 0, &k_a) == 0); + assert(qwen36_set_lora_tensor(distributed, 1, 1, &local_k_b) == 0); + assert(qwen36_set_lora_tensor(distributed, 2, 0, &v_a) == 0); + assert(qwen36_set_lora_tensor(distributed, 2, 1, &local_v_b) == 0); + assert(qwen36_set_lora_tensor(distributed, 3, 0, &local_o_a) == 0); + assert(qwen36_set_lora_tensor(distributed, 3, 1, &o_b) == 0); + assert(qwen36_set_lora_tensor(reference, 0, 0, &q_a) == 0); + assert(qwen36_set_lora_tensor(reference, 0, 1, &q_b) == 0); + assert(qwen36_set_lora_tensor(reference, 1, 0, &k_a) == 0); + assert(qwen36_set_lora_tensor(reference, 1, 1, &k_b) == 0); + assert(qwen36_set_lora_tensor(reference, 2, 0, &v_a) == 0); + assert(qwen36_set_lora_tensor(reference, 2, 1, &v_b) == 0); + assert(qwen36_set_lora_tensor(reference, 3, 0, &o_a) == 0); + assert(qwen36_set_lora_tensor(reference, 3, 1, &o_b) == 0); + + auto input_ids = at::tensor({1, 2, 3}, + at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 3}); + auto target_mask = at::ones({1, 3}, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); + auto attention_mask = at::ones({1, 3}, + at::TensorOptions().device(at::kCUDA).dtype(at::kBool)); + + const double distributed_eval = qwen36_eval_step( + distributed, &input_ids, &target_mask, &attention_mask); + const double reference_eval = qwen36_eval_step( + reference, &input_ids, &target_mask, &attention_mask); + assert(distributed_eval > 0.0 && reference_eval > 0.0); + + const double distributed_micro = qwen36_train_micro_step( + distributed, &input_ids, &target_mask, &attention_mask, 1.0, 0); + const double reference_micro = qwen36_train_micro_step( + reference, &input_ids, &target_mask, &attention_mask, 1.0, 0); + assert(distributed_micro > 0.0 && reference_micro > 0.0); + auto* local_q_b_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(distributed, 0, 1)); + auto* full_q_b_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(reference, 0, 1)); + auto* local_k_b_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(distributed, 1, 1)); + auto* full_k_b_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(reference, 1, 1)); + auto* local_v_b_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(distributed, 2, 1)); + auto* full_v_b_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(reference, 2, 1)); + auto* local_o_a_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(distributed, 3, 0)); + auto* full_o_a_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(reference, 3, 0)); + assert(local_q_b_grad && full_q_b_grad && local_k_b_grad && full_k_b_grad); + assert(local_v_b_grad && full_v_b_grad && local_o_a_grad && full_o_a_grad); + const double q_b_grad_diff = max_diff( + *local_q_b_grad, + full_q_b_grad->narrow( + 0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim)); + const double o_a_grad_diff = max_diff( + *local_o_a_grad, + full_o_a_grad->narrow( + 1, rank * local_heads * head_dim, local_heads * head_dim)); + const double k_b_grad_diff = max_diff( + *local_k_b_grad, + full_k_b_grad->narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim)); + const double v_b_grad_diff = max_diff( + *local_v_b_grad, + full_v_b_grad->narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim)); + assert(qwen36_abort_gradient_accumulation(distributed) == 0); + assert(qwen36_abort_gradient_accumulation(reference) == 0); + + const double distributed_loss = qwen36_train_step( + distributed, &input_ids, &target_mask, &attention_mask); + const double reference_loss = qwen36_train_step( + reference, &input_ids, &target_mask, &attention_mask); + assert(distributed_loss > 0.0 && reference_loss > 0.0); + + auto* updated_q_a = reinterpret_cast(qwen36_get_lora_a(distributed, 0)); + auto* updated_q_b = reinterpret_cast(qwen36_get_lora_b(distributed, 0)); + auto* updated_k_a = reinterpret_cast(qwen36_get_lora_a(distributed, 1)); + auto* updated_k_b = reinterpret_cast(qwen36_get_lora_b(distributed, 1)); + auto* updated_v_a = reinterpret_cast(qwen36_get_lora_a(distributed, 2)); + auto* updated_v_b = reinterpret_cast(qwen36_get_lora_b(distributed, 2)); + auto* updated_o_a = reinterpret_cast(qwen36_get_lora_a(distributed, 3)); + auto* updated_o_b = reinterpret_cast(qwen36_get_lora_b(distributed, 3)); + auto* reference_q_a = reinterpret_cast(qwen36_get_lora_a(reference, 0)); + auto* reference_q_b = reinterpret_cast(qwen36_get_lora_b(reference, 0)); + auto* reference_k_a = reinterpret_cast(qwen36_get_lora_a(reference, 1)); + auto* reference_k_b = reinterpret_cast(qwen36_get_lora_b(reference, 1)); + auto* reference_v_a = reinterpret_cast(qwen36_get_lora_a(reference, 2)); + auto* reference_v_b = reinterpret_cast(qwen36_get_lora_b(reference, 2)); + auto* reference_o_a = reinterpret_cast(qwen36_get_lora_a(reference, 3)); + auto* reference_o_b = reinterpret_cast(qwen36_get_lora_b(reference, 3)); + assert(updated_q_a && updated_q_b && updated_k_a && updated_k_b); + assert(updated_v_a && updated_v_b && updated_o_a && updated_o_b); + assert(reference_q_a && reference_q_b && reference_k_a && reference_k_b); + assert(reference_v_a && reference_v_b && reference_o_a && reference_o_b); + + constexpr int64_t optimizer_count = 14; + std::vector local_m(optimizer_count), local_v(optimizer_count); + std::vector full_m(optimizer_count), full_v(optimizer_count); + assert(qwen36_export_optimizer_state( + distributed, local_m.data(), local_v.data(), optimizer_count) == + optimizer_count); + assert(qwen36_export_optimizer_state( + reference, full_m.data(), full_v.data(), optimizer_count) == + optimizer_count); + double optimizer_m_diff = 0.0; + double optimizer_v_diff = 0.0; + double adam_error = 0.0; + auto observe_optimizer = [&](int64_t index, const at::Tensor& expected_m, + const at::Tensor& expected_v, + const at::Tensor& updated, + const at::Tensor& before) { + auto* m = reinterpret_cast(local_m[index]); + auto* v = reinterpret_cast(local_v[index]); + assert(m && v); + optimizer_m_diff = std::max(optimizer_m_diff, max_diff(*m, expected_m)); + optimizer_v_diff = std::max(optimizer_v_diff, max_diff(*v, expected_v)); + auto expected_param = ( + before.to(at::kFloat) - 1e-3 * + (*m / (1.0 - 0.9)) / + (((*v / (1.0 - 0.999)).sqrt()) + 1e-8)) + .to(at::kBFloat16); + adam_error = std::max(adam_error, max_diff(updated, expected_param)); + }; + auto state = [](std::vector& tensors, int64_t index) -> at::Tensor& { + auto* tensor = reinterpret_cast(tensors[index]); + assert(tensor); + return *tensor; + }; + observe_optimizer(0, state(full_m, 0), state(full_v, 0), *updated_q_a, q_a); + observe_optimizer(1, + state(full_m, 1).narrow( + 0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim), + state(full_v, 1).narrow( + 0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim), + *updated_q_b, local_q_b); + observe_optimizer(2, state(full_m, 2), state(full_v, 2), *updated_k_a, k_a); + observe_optimizer(3, + state(full_m, 3).narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim), + state(full_v, 3).narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim), + *updated_k_b, local_k_b); + observe_optimizer(4, state(full_m, 4), state(full_v, 4), *updated_v_a, v_a); + observe_optimizer(5, + state(full_m, 5).narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim), + state(full_v, 5).narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim), + *updated_v_b, local_v_b); + observe_optimizer(6, + state(full_m, 6).narrow( + 1, rank * local_heads * head_dim, local_heads * head_dim), + state(full_v, 6).narrow( + 1, rank * local_heads * head_dim, local_heads * head_dim), + *updated_o_a, local_o_a); + observe_optimizer(7, state(full_m, 7), state(full_v, 7), *updated_o_b, o_b); + const double q_a_diff = max_diff(*updated_q_a, *reference_q_a); + const double q_b_diff = max_diff( + *updated_q_b, + reference_q_b->narrow( + 0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim)); + const double o_a_diff = max_diff( + *updated_o_a, + reference_o_a->narrow( + 1, rank * local_heads * head_dim, local_heads * head_dim)); + const double o_b_diff = max_diff(*updated_o_b, *reference_o_b); + const double k_a_diff = max_diff(*updated_k_a, *reference_k_a); + const double k_b_diff = max_diff( + *updated_k_b, reference_k_b->narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim)); + const double v_a_diff = max_diff(*updated_v_a, *reference_v_a); + const double v_b_diff = max_diff( + *updated_v_b, reference_v_b->narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim)); + std::printf( + "base_tp_attention_smoke rank=%d eval_diff=%0.8e loss_diff=%0.8e q_b_grad_diff=%0.8e k_b_grad_diff=%0.8e v_b_grad_diff=%0.8e o_a_grad_diff=%0.8e m_diff=%0.8e v_diff=%0.8e adam_error=%0.8e q_a_diff=%0.8e q_b_diff=%0.8e k_a_diff=%0.8e k_b_diff=%0.8e v_a_diff=%0.8e v_b_diff=%0.8e o_a_diff=%0.8e o_b_diff=%0.8e\n", + rank, std::abs(distributed_eval - reference_eval), + std::abs(distributed_loss - reference_loss), q_b_grad_diff, + k_b_grad_diff, v_b_grad_diff, o_a_grad_diff, optimizer_m_diff, + optimizer_v_diff, adam_error, q_a_diff, q_b_diff, k_a_diff, + k_b_diff, v_a_diff, v_b_diff, o_a_diff, o_b_diff); + std::fflush(stdout); + // Row-parallel BF16 rounds each local matmul before the NCCL sum, while + // the reference rounds once after the full matmul. + assert(std::abs(distributed_eval - reference_eval) < 5e-3); + assert(std::abs(distributed_loss - reference_loss) < 5e-3); + assert(q_b_grad_diff < 3e-4 && k_b_grad_diff < 3e-4); + assert(v_b_grad_diff < 5e-4 && o_a_grad_diff < 5e-4); + assert(optimizer_m_diff < 5e-5 && optimizer_v_diff < 5e-8); + assert(adam_error < 1e-8); + assert(std::max({q_a_diff, q_b_diff, k_a_diff, k_b_diff, + v_a_diff, v_b_diff, o_a_diff, o_b_diff}) <= 2e-3); + + qwen36_free_training_context(reference); + qwen36_free_training_context(distributed); + + setenv("TP_SIZE", "2", 1); + void* dynamic_distributed = qwen36_create_training_context_ex( + local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head, + &config, 1, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + &target_layer, 1, "v_proj", 1); + assert(dynamic_distributed && qwen36_init_nccl(dynamic_distributed) == 0); + const int64_t distributed_id = qwen36_add_lora( + dynamic_distributed, lora_rank, lora_rank, + &target_layer, 1, "q_proj,k_proj,v_proj,o_proj"); + assert(distributed_id > 0); + + setenv("TP_SIZE", "1", 1); + void* dynamic_reference = qwen36_create_training_context( + full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head, + &config, 1, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + &target_layer, 1, "v_proj"); + assert(dynamic_reference); + const int64_t reference_id = qwen36_add_lora( + dynamic_reference, lora_rank, lora_rank, + &target_layer, 1, "q_proj,k_proj,v_proj,o_proj"); + assert(reference_id > 0); + + assert(qwen36_set_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "q_proj", 0, &q_a) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "q_proj", 1, &local_q_b) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "k_proj", 0, &k_a) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "k_proj", 1, &local_k_b) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "v_proj", 0, &v_a) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "v_proj", 1, &local_v_b) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "o_proj", 0, &local_o_a) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "o_proj", 1, &o_b) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "q_proj", 0, &q_a) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "q_proj", 1, &q_b) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "k_proj", 0, &k_a) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "k_proj", 1, &k_b) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "v_proj", 0, &v_a) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "v_proj", 1, &v_b) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "o_proj", 0, &o_a) == 0); + assert(qwen36_set_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "o_proj", 1, &o_b) == 0); + + const double dynamic_distributed_loss = qwen36_train_multi_lora_selected( + dynamic_distributed, &input_ids, &target_mask, &attention_mask, + &distributed_id, 1, lora_rank); + const double dynamic_reference_loss = qwen36_train_multi_lora_selected( + dynamic_reference, &input_ids, &target_mask, &attention_mask, + &reference_id, 1, lora_rank); + assert(dynamic_distributed_loss > 0.0 && dynamic_reference_loss > 0.0); + + auto* dynamic_q_a = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "q_proj", 0)); + auto* dynamic_q_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "q_proj", 1)); + auto* dynamic_k_a = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "k_proj", 0)); + auto* dynamic_k_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "k_proj", 1)); + auto* dynamic_v_a = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "v_proj", 0)); + auto* dynamic_v_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "v_proj", 1)); + auto* dynamic_o_a = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "o_proj", 0)); + auto* dynamic_o_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_distributed, distributed_id, 0, "o_proj", 1)); + auto* dynamic_reference_q_a = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "q_proj", 0)); + auto* dynamic_reference_q_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "q_proj", 1)); + auto* dynamic_reference_k_a = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "k_proj", 0)); + auto* dynamic_reference_k_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "k_proj", 1)); + auto* dynamic_reference_v_a = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "v_proj", 0)); + auto* dynamic_reference_v_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "v_proj", 1)); + auto* dynamic_reference_o_a = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "o_proj", 0)); + auto* dynamic_reference_o_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor( + dynamic_reference, reference_id, 0, "o_proj", 1)); + assert(dynamic_q_a && dynamic_q_b && dynamic_o_a && dynamic_o_b); + assert(dynamic_k_a && dynamic_k_b && dynamic_v_a && dynamic_v_b); + assert(dynamic_reference_q_a && dynamic_reference_q_b && + dynamic_reference_o_a && dynamic_reference_o_b); + assert(dynamic_reference_k_a && dynamic_reference_k_b && + dynamic_reference_v_a && dynamic_reference_v_b); + const double dynamic_q_a_diff = max_diff(*dynamic_q_a, *dynamic_reference_q_a); + const double dynamic_q_b_diff = max_diff( + *dynamic_q_b, dynamic_reference_q_b->narrow( + 0, rank * 2 * local_heads * head_dim, 2 * local_heads * head_dim)); + const double dynamic_o_a_diff = max_diff( + *dynamic_o_a, dynamic_reference_o_a->narrow( + 1, rank * local_heads * head_dim, local_heads * head_dim)); + const double dynamic_o_b_diff = max_diff(*dynamic_o_b, *dynamic_reference_o_b); + const double dynamic_k_a_diff = max_diff(*dynamic_k_a, *dynamic_reference_k_a); + const double dynamic_k_b_diff = max_diff( + *dynamic_k_b, dynamic_reference_k_b->narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim)); + const double dynamic_v_a_diff = max_diff(*dynamic_v_a, *dynamic_reference_v_a); + const double dynamic_v_b_diff = max_diff( + *dynamic_v_b, dynamic_reference_v_b->narrow( + 0, rank * local_kv_heads * head_dim, local_kv_heads * head_dim)); + std::printf( + "dynamic_tp_attention_smoke rank=%d loss_diff=%0.8e q_a_diff=%0.8e q_b_diff=%0.8e k_a_diff=%0.8e k_b_diff=%0.8e v_a_diff=%0.8e v_b_diff=%0.8e o_a_diff=%0.8e o_b_diff=%0.8e\n", + rank, + std::abs(dynamic_distributed_loss - dynamic_reference_loss), + dynamic_q_a_diff, dynamic_q_b_diff, dynamic_k_a_diff, dynamic_k_b_diff, + dynamic_v_a_diff, dynamic_v_b_diff, dynamic_o_a_diff, dynamic_o_b_diff); + std::fflush(stdout); + assert(std::abs(dynamic_distributed_loss - dynamic_reference_loss) < 5e-3); + assert(dynamic_q_a_diff < 5e-5 && dynamic_q_b_diff < 5e-5); + assert(dynamic_k_a_diff < 5e-5 && dynamic_k_b_diff < 5e-5); + assert(dynamic_v_a_diff < 5e-5 && dynamic_v_b_diff < 5e-5); + assert(dynamic_o_a_diff < 5e-5 && dynamic_o_b_diff < 5e-5); + + qwen36_free_training_context(dynamic_reference); + qwen36_free_training_context(dynamic_distributed); + return 0; +} diff --git a/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp new file mode 100644 index 00000000..f675963c --- /dev/null +++ b/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp @@ -0,0 +1,298 @@ +#include + +#include +#include +#include +#include +#include +#include + +struct LayerConfig { + int64_t layer_type, num_heads, num_kv_heads, head_dim; + int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel; + double partial_rotary_factor, rope_theta, rms_eps; + int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count; + int64_t intermediate_size; + int32_t norm_topk_prob; + void* nccl_comm; + void* nccl_stream; +}; + +extern "C" void qwen36_set_cuda_device(int32_t); +extern "C" void* qwen36_create_training_context( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*); +extern "C" void* qwen36_create_training_context_ex( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*, int32_t); +extern "C" int32_t qwen36_init_nccl(void*); +extern "C" void* qwen36_get_lora_a(void*, int64_t); +extern "C" void* qwen36_get_lora_b(void*, int64_t); +extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*); +extern "C" void* qwen36_get_lora_grad_accumulator(void*, int64_t, int32_t); +extern "C" int32_t qwen36_abort_gradient_accumulation(void*); +extern "C" double qwen36_eval_step(void*, void*, void*, void*); +extern "C" double qwen36_train_step(void*, void*, void*, void*); +extern "C" double qwen36_train_micro_step( + void*, void*, void*, void*, double, int32_t); +extern "C" int64_t qwen36_export_optimizer_state( + void*, void**, void**, int64_t); +extern "C" void qwen36_free_training_context(void*); + +static at::Tensor deterministic( + std::initializer_list shape, double scale, int64_t offset = 0 +) { + int64_t count = 1; + for (int64_t dim : shape) count *= dim; + return ((at::arange(count, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)) + .add(offset).remainder(23) - 11.0) * scale) + .reshape(shape).to(at::kBFloat16); +} + +static void append_gdn_layer( + std::vector& weights, int64_t hidden, int64_t intermediate, + int64_t state_dim, int64_t layer +) { + const int64_t qkv_dim = 3 * state_dim; + const auto ones = at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16); + weights.push_back(at::ones({hidden}, ones)); + weights.push_back(at::ones({hidden}, ones)); + weights.push_back(deterministic({qkv_dim, hidden}, 0.0020, layer)); + weights.push_back(deterministic({state_dim, hidden}, 0.0022, layer + 1)); + weights.push_back(deterministic({1, hidden}, 0.0015, layer + 2)); + weights.push_back(deterministic({1, hidden}, 0.0017, layer + 3)); + weights.push_back(deterministic({1}, 0.0010, layer + 4)); + weights.push_back(deterministic({1}, 0.0010, layer + 5)); + weights.push_back(deterministic({qkv_dim, 1, 4}, 0.0012, layer + 6)); + weights.push_back(at::ones({state_dim}, ones)); + weights.push_back(deterministic({hidden, state_dim}, 0.0020, layer + 7)); + weights.push_back(deterministic({intermediate, hidden}, 0.0020, layer + 8)); + weights.push_back(deterministic({intermediate, hidden}, 0.0018, layer + 9)); + weights.push_back(deterministic({hidden, intermediate}, 0.0020, layer + 10)); +} + +static std::vector pointers(std::vector& tensors) { + std::vector result; + result.reserve(tensors.size()); + for (auto& tensor : tensors) result.push_back(&tensor); + return result; +} + +static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) { + return (lhs - rhs).abs().max().item(); +} + +int main() { + const int rank = std::atoi(std::getenv("RANK")); + const int world = std::atoi(std::getenv("WORLD_SIZE")); + assert(world == 2 && rank >= 0 && rank < world); + qwen36_set_cuda_device(rank); + + constexpr int64_t hidden = 8; + constexpr int64_t intermediate = 12; + constexpr int64_t state_dim = 128; + constexpr int64_t qkv_dim = 3 * state_dim; + constexpr int64_t vocab = 16; + constexpr int64_t lora_rank = 4; + constexpr int64_t local_rank = lora_rank / 2; + constexpr int64_t slots_per_layer = 8; + + std::vector weights; + append_gdn_layer(weights, hidden, intermediate, state_dim, 0); + append_gdn_layer(weights, hidden, intermediate, state_dim, 1); + for (auto& weight : weights) weight.set_requires_grad(false); + auto weight_ptrs = pointers(weights); + + auto embed = deterministic({vocab, hidden}, 0.0030); + auto final_norm = at::ones( + {hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)); + auto lm_head = deterministic({vocab, hidden}, 0.0025); + LayerConfig configs[2]{}; + for (auto& config : configs) { + config.layer_type = 1; + config.num_k_heads = 1; + config.key_dim = state_dim; + config.num_v_heads = 1; + config.val_dim = state_dim; + config.conv_kernel = 4; + config.rms_eps = 1e-5; + config.intermediate_size = intermediate; + } + const int64_t target_layers[2] = {0, 1}; + + setenv("TP_SIZE", "2", 1); + unsetenv("RUSTRAIN_DATA_PARALLEL"); + void* distributed = qwen36_create_training_context_ex( + weight_ptrs.data(), weight_ptrs.size(), &embed, &final_norm, &lm_head, + configs, 2, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + target_layers, 2, "in_proj_qkv", 1); + assert(distributed && qwen36_init_nccl(distributed) == 0); + + setenv("TP_SIZE", "1", 1); + void* reference = qwen36_create_training_context( + weight_ptrs.data(), weight_ptrs.size(), &embed, &final_norm, &lm_head, + configs, 2, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + target_layers, 2, "in_proj_qkv"); + assert(reference); + + std::vector full_a; + std::vector full_b; + for (int64_t layer = 0; layer < 2; ++layer) { + full_a.push_back(deterministic( + {lora_rank, hidden}, 0.0010, 20 + layer)); + full_b.push_back(deterministic( + {qkv_dim, lora_rank}, 0.0008, 30 + layer)); + auto local_a = full_a.back().narrow( + 0, rank * local_rank, local_rank).contiguous(); + auto local_b = full_b.back().narrow( + 1, rank * local_rank, local_rank).contiguous(); + const int64_t slot = layer * slots_per_layer; + assert(qwen36_set_lora_tensor(distributed, slot, 0, &local_a) == 0); + assert(qwen36_set_lora_tensor(distributed, slot, 1, &local_b) == 0); + assert(qwen36_set_lora_tensor(reference, slot, 0, &full_a.back()) == 0); + assert(qwen36_set_lora_tensor(reference, slot, 1, &full_b.back()) == 0); + } + + auto input_ids = at::tensor({1, 2, 3}, + at::TensorOptions().device(at::kCUDA).dtype(at::kLong)).reshape({1, 3}); + auto target_mask = at::ones({1, 3}, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); + auto attention_mask = at::ones({1, 3}, + at::TensorOptions().device(at::kCUDA).dtype(at::kBool)); + + const double distributed_eval = qwen36_eval_step( + distributed, &input_ids, &target_mask, &attention_mask); + const double reference_eval = qwen36_eval_step( + reference, &input_ids, &target_mask, &attention_mask); + assert(distributed_eval > 0.0 && reference_eval > 0.0); + + const double distributed_micro = qwen36_train_micro_step( + distributed, &input_ids, &target_mask, &attention_mask, 1.0, 0); + const double reference_micro = qwen36_train_micro_step( + reference, &input_ids, &target_mask, &attention_mask, 1.0, 0); + assert(distributed_micro > 0.0 && reference_micro > 0.0); + + double max_a_grad_diff = 0.0; + double max_b_grad_diff = 0.0; + for (int64_t layer = 0; layer < 2; ++layer) { + const int64_t slot = layer * slots_per_layer; + auto* local_a_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(distributed, slot, 0)); + auto* local_b_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(distributed, slot, 1)); + auto* full_a_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(reference, slot, 0)); + auto* full_b_grad = reinterpret_cast( + qwen36_get_lora_grad_accumulator(reference, slot, 1)); + assert(local_a_grad && local_b_grad && full_a_grad && full_b_grad); + max_a_grad_diff = std::max(max_a_grad_diff, max_diff( + *local_a_grad, full_a_grad->narrow( + 0, rank * local_rank, local_rank))); + max_b_grad_diff = std::max(max_b_grad_diff, max_diff( + *local_b_grad, full_b_grad->narrow( + 1, rank * local_rank, local_rank))); + } + assert(qwen36_abort_gradient_accumulation(distributed) == 0); + assert(qwen36_abort_gradient_accumulation(reference) == 0); + + const double distributed_loss = qwen36_train_step( + distributed, &input_ids, &target_mask, &attention_mask); + const double reference_loss = qwen36_train_step( + reference, &input_ids, &target_mask, &attention_mask); + assert(distributed_loss > 0.0 && reference_loss > 0.0); + + constexpr int64_t optimizer_count = 2 * 2 * slots_per_layer; + std::vector local_m(optimizer_count), local_v(optimizer_count); + std::vector full_m(optimizer_count), full_v(optimizer_count); + assert(qwen36_export_optimizer_state( + distributed, local_m.data(), local_v.data(), optimizer_count) == + optimizer_count); + assert(qwen36_export_optimizer_state( + reference, full_m.data(), full_v.data(), optimizer_count) == + optimizer_count); + + double max_a_param_diff = 0.0; + double max_b_param_diff = 0.0; + double max_m_diff = 0.0; + double max_v_diff = 0.0; + double max_adam_error = 0.0; + for (int64_t layer = 0; layer < 2; ++layer) { + const int64_t slot = layer * slots_per_layer; + auto* local_a = reinterpret_cast( + qwen36_get_lora_a(distributed, slot)); + auto* local_b = reinterpret_cast( + qwen36_get_lora_b(distributed, slot)); + auto* full_a_param = reinterpret_cast( + qwen36_get_lora_a(reference, slot)); + auto* full_b_param = reinterpret_cast( + qwen36_get_lora_b(reference, slot)); + assert(local_a && local_b && full_a_param && full_b_param); + max_a_param_diff = std::max(max_a_param_diff, max_diff( + *local_a, full_a_param->narrow( + 0, rank * local_rank, local_rank))); + max_b_param_diff = std::max(max_b_param_diff, max_diff( + *local_b, full_b_param->narrow( + 1, rank * local_rank, local_rank))); + + auto* local_m_a = reinterpret_cast(local_m[2 * slot]); + auto* local_m_b = reinterpret_cast(local_m[2 * slot + 1]); + auto* local_v_a = reinterpret_cast(local_v[2 * slot]); + auto* local_v_b = reinterpret_cast(local_v[2 * slot + 1]); + auto* full_m_a = reinterpret_cast(full_m[2 * slot]); + auto* full_m_b = reinterpret_cast(full_m[2 * slot + 1]); + auto* full_v_a = reinterpret_cast(full_v[2 * slot]); + auto* full_v_b = reinterpret_cast(full_v[2 * slot + 1]); + assert(local_m_a && local_m_b && local_v_a && local_v_b); + assert(full_m_a && full_m_b && full_v_a && full_v_b); + max_m_diff = std::max({max_m_diff, + max_diff(*local_m_a, full_m_a->narrow( + 0, rank * local_rank, local_rank)), + max_diff(*local_m_b, full_m_b->narrow( + 1, rank * local_rank, local_rank))}); + max_v_diff = std::max({max_v_diff, + max_diff(*local_v_a, full_v_a->narrow( + 0, rank * local_rank, local_rank)), + max_diff(*local_v_b, full_v_b->narrow( + 1, rank * local_rank, local_rank))}); + + auto local_a_before = full_a[layer].narrow( + 0, rank * local_rank, local_rank).contiguous(); + auto local_b_before = full_b[layer].narrow( + 1, rank * local_rank, local_rank).contiguous(); + auto expected_a = ( + local_a_before.to(at::kFloat) - 1e-3 * + (*local_m_a / (1.0 - 0.9)) / + (((*local_v_a / (1.0 - 0.999)).sqrt()) + 1e-8)) + .to(at::kBFloat16); + auto expected_b = ( + local_b_before.to(at::kFloat) - 1e-3 * + (*local_m_b / (1.0 - 0.9)) / + (((*local_v_b / (1.0 - 0.999)).sqrt()) + 1e-8)) + .to(at::kBFloat16); + max_adam_error = std::max({max_adam_error, + max_diff(*local_a, expected_a), max_diff(*local_b, expected_b)}); + } + + std::printf( + "latent_tp_two_layer_smoke rank=%d eval_diff=%0.8e loss_diff=%0.8e a_grad_diff=%0.8e b_grad_diff=%0.8e m_diff=%0.8e v_diff=%0.8e adam_error=%0.8e a_param_diff=%0.8e b_param_diff=%0.8e\n", + rank, std::abs(distributed_eval - reference_eval), + std::abs(distributed_loss - reference_loss), max_a_grad_diff, + max_b_grad_diff, max_m_diff, max_v_diff, max_adam_error, + max_a_param_diff, max_b_param_diff); + std::fflush(stdout); + assert(std::abs(distributed_eval - reference_eval) < 5e-3); + assert(std::abs(distributed_loss - reference_loss) < 5e-3); + assert(max_a_grad_diff < 5e-4 && max_b_grad_diff < 5e-4); + assert(max_m_diff < 5e-5 && max_v_diff < 5e-8); + assert(max_adam_error == 0.0); + assert(max_a_param_diff < 1e-3 && max_b_param_diff < 1e-3); + + qwen36_free_training_context(reference); + qwen36_free_training_context(distributed); + return 0; +} diff --git a/crates/rustrain-server/src/checkpoint.rs b/crates/rustrain-server/src/checkpoint.rs index 35403ce0..2974f38e 100644 --- a/crates/rustrain-server/src/checkpoint.rs +++ b/crates/rustrain-server/src/checkpoint.rs @@ -1,13 +1,34 @@ //! Checkpoint save/load: adapter (LoRA A/B) + optimizer state (Adam m/v) + step count. -use anyhow::{Context, Result, bail}; +use anyhow::{bail, Context, Result}; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; use std::env; use std::path::{Path, PathBuf}; use tch::Tensor; -const TP_CHECKPOINT_FORMAT: &str = "rustrain-checkpoint-v3-tp"; +const TP_CHECKPOINT_FORMAT: &str = "rustrain-checkpoint-v4-tp"; +const LEGACY_TP_CHECKPOINT_FORMAT: &str = "rustrain-checkpoint-v3-tp"; + +pub fn is_legacy_tensor_parallel_checkpoint(manifest: &CheckpointManifest) -> bool { + manifest.format == LEGACY_TP_CHECKPOINT_FORMAT +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LoraTpShardLayout { + #[default] + LatentRank, + ColumnParallel, + RowParallel, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct LoraSlotIdentity { + pub index: usize, + pub layer: usize, + pub module: String, +} #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ParallelCheckpointManifest { @@ -110,6 +131,10 @@ pub struct TensorShardManifest { pub global_shape: Vec, pub local_shape: Vec, pub partition_axis: usize, + #[serde(default)] + pub layout: LoraTpShardLayout, + #[serde(default)] + pub replicated: bool, pub global_offset: Vec, pub replica_identity: String, } @@ -129,6 +154,88 @@ pub struct CheckpointManifest { pub parallel: Option, #[serde(default)] pub tensor_shards: Vec, + #[serde(default)] + pub fixed_shard_layouts: Vec, + #[serde(default)] + pub fixed_slot_identities: Vec, +} + +pub fn validate_fixed_tp_resume( + manifest: &CheckpointManifest, + expected_layouts: &[LoraTpShardLayout], + expected_identities: &[LoraSlotIdentity], +) -> Result<()> { + if is_legacy_tensor_parallel_checkpoint(manifest) { + if expected_layouts + .iter() + .any(|layout| *layout != LoraTpShardLayout::LatentRank) + { + bail!( + "legacy tensor-parallel v3 checkpoints cannot restore fixed Q/K/V/O LoRA into the projection-aware layout; use a v4 checkpoint or migrate the adapter from a merged artifact" + ); + } + return Ok(()); + } + if manifest.fixed_shard_layouts != expected_layouts { + bail!( + "fixed LoRA shard layouts do not match the current runtime slots: checkpoint={:?}, runtime={expected_layouts:?}", + manifest.fixed_shard_layouts + ); + } + if manifest.fixed_slot_identities != expected_identities { + bail!( + "fixed LoRA slot identities do not match the current runtime slots: checkpoint={:?}, runtime={expected_identities:?}", + manifest.fixed_slot_identities + ); + } + Ok(()) +} + +pub fn validate_dynamic_tp_resume( + manifest: &CheckpointManifest, + adapter_id: i64, + saved_layouts: &[LoraTpShardLayout], + expected_layouts: &[LoraTpShardLayout], +) -> Result<()> { + if is_legacy_tensor_parallel_checkpoint(manifest) { + if expected_layouts + .iter() + .any(|layout| *layout != LoraTpShardLayout::LatentRank) + { + bail!( + "legacy tensor-parallel v3 checkpoint adapter {adapter_id} contains Q/K/V/O LoRA that cannot be restored into the projection-aware layout; use a v4 checkpoint or migrate the adapter from a merged artifact" + ); + } + return Ok(()); + } + if saved_layouts != expected_layouts { + bail!( + "dynamic adapter {adapter_id} shard layouts do not match the current runtime slots: checkpoint={saved_layouts:?}, runtime={expected_layouts:?}" + ); + } + Ok(()) +} + +pub fn fixed_restore_slot_indices( + saved_a_count: usize, + saved_b_count: usize, + active_slot_indices: &[usize], + native_slot_count: usize, +) -> Result> { + if saved_a_count != saved_b_count { + bail!("checkpoint fixed LoRA A/B count mismatch: {saved_a_count}/{saved_b_count}"); + } + if saved_a_count == active_slot_indices.len() { + return Ok(active_slot_indices.to_vec()); + } + if saved_a_count == native_slot_count { + // v1/v2 checkpoints stored inactive positional placeholders. + return Ok((0..native_slot_count).collect()); + } + bail!( + "checkpoint LoRA slot count mismatch: checkpoint A/B={saved_a_count}/{saved_b_count}, active={}, native={native_slot_count}", + active_slot_indices.len() + ) } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -143,6 +250,8 @@ pub struct DynamicAdapterManifest { pub optimizer_step: u64, pub target_layers: Vec, pub target_modules: Vec, + #[serde(default)] + pub shard_layouts: Vec, pub parameter_count: usize, pub optimizer_count: usize, } @@ -218,6 +327,8 @@ pub fn save_checkpoint_with_dynamic( adam_m, adam_v, dynamic_adapters, + &[], + &[], None, ) } @@ -235,6 +346,8 @@ pub fn save_checkpoint_with_dynamic_for_topology( adam_m: &[Tensor], adam_v: &[Tensor], dynamic_adapters: &[DynamicAdapterCheckpoint], + fixed_shard_layouts: &[LoraTpShardLayout], + fixed_slot_identities: &[LoraSlotIdentity], parallel: &ParallelCheckpointManifest, ) -> Result<()> { if !parallel.is_tensor_parallel() { @@ -265,6 +378,8 @@ pub fn save_checkpoint_with_dynamic_for_topology( adam_m, adam_v, dynamic_adapters, + fixed_shard_layouts, + fixed_slot_identities, Some(parallel), ) } @@ -282,6 +397,8 @@ fn save_checkpoint_with_dynamic_at( adam_m: &[Tensor], adam_v: &[Tensor], dynamic_adapters: &[DynamicAdapterCheckpoint], + fixed_shard_layouts: &[LoraTpShardLayout], + fixed_slot_identities: &[LoraSlotIdentity], parallel: Option<&ParallelCheckpointManifest>, ) -> Result<()> { validate_tensor_counts( @@ -292,6 +409,20 @@ fn save_checkpoint_with_dynamic_at( dynamic_adapters, parallel.is_some(), )?; + if parallel.is_some() && fixed_slot_identities.len() != lora_a.len() { + bail!( + "fixed LoRA slot identity count {} does not match parameter count {}", + fixed_slot_identities.len(), + lora_a.len() + ); + } + let unique_fixed_slots = fixed_slot_identities + .iter() + .map(|identity| (identity.index, identity.layer, identity.module.as_str())) + .collect::>(); + if unique_fixed_slots.len() != fixed_slot_identities.len() { + bail!("fixed LoRA slot identities must be unique"); + } let tensor_shards = match parallel { Some(parallel) => build_tensor_shard_manifest( parallel, @@ -301,6 +432,7 @@ fn save_checkpoint_with_dynamic_at( adam_m, adam_v, dynamic_adapters, + fixed_shard_layouts, )?, None => Vec::new(), }; @@ -355,6 +487,8 @@ fn save_checkpoint_with_dynamic_at( dynamic_adapters: dynamic_manifests, parallel: parallel.cloned(), tensor_shards, + fixed_shard_layouts: fixed_shard_layouts.to_vec(), + fixed_slot_identities: fixed_slot_identities.to_vec(), }; let manifest_path = dir.join("manifest.json"); std::fs::write(&manifest_path, serde_json::to_string_pretty(&manifest)?) @@ -397,9 +531,11 @@ fn load_checkpoint_at( .with_context(|| "parse manifest.json")?; match expected_parallel { Some(expected) => { - if manifest.format != TP_CHECKPOINT_FORMAT { + if manifest.format != TP_CHECKPOINT_FORMAT + && manifest.format != LEGACY_TP_CHECKPOINT_FORMAT + { bail!( - "tensor-parallel resume requires {TP_CHECKPOINT_FORMAT}, found {}", + "tensor-parallel resume requires {TP_CHECKPOINT_FORMAT} or {LEGACY_TP_CHECKPOINT_FORMAT}, found {}", manifest.format ); } @@ -413,7 +549,9 @@ fn load_checkpoint_at( ); } } - None if manifest.format == TP_CHECKPOINT_FORMAT => { + None if manifest.format == TP_CHECKPOINT_FORMAT + || manifest.format == LEGACY_TP_CHECKPOINT_FORMAT => + { bail!("tensor-parallel checkpoint must be loaded with rank topology"); } None => {} @@ -460,6 +598,7 @@ fn load_checkpoint_at( &adam_m, &adam_v, &dynamic_adapters, + &manifest.fixed_shard_layouts, )?; validate_saved_shards(&manifest.tensor_shards, &expected_shards)?; } @@ -535,6 +674,7 @@ fn build_tensor_shard_manifest( adam_m: &[Tensor], adam_v: &[Tensor], dynamic_adapters: &[DynamicAdapterCheckpoint], + fixed_shard_layouts: &[LoraTpShardLayout], ) -> Result> { validate_tensor_counts(lora_a, lora_b, adam_m, adam_v, dynamic_adapters, true)?; let mut shards = Vec::new(); @@ -547,6 +687,7 @@ fn build_tensor_shard_manifest( lora_b, adam_m, adam_v, + fixed_shard_layouts, )?; for adapter in dynamic_adapters { append_adapter_shards( @@ -558,6 +699,7 @@ fn build_tensor_shard_manifest( &adapter.lora_b, &adapter.adam_m, &adapter.adam_v, + &adapter.manifest.shard_layouts, )?; } Ok(shards) @@ -573,7 +715,21 @@ fn append_adapter_shards( lora_b: &[Tensor], adam_m: &[Tensor], adam_v: &[Tensor], + shard_layouts: &[LoraTpShardLayout], ) -> Result<()> { + if !shard_layouts.is_empty() && shard_layouts.len() != lora_a.len() { + bail!( + "LoRA shard layout count {} does not match parameter count {}", + shard_layouts.len(), + lora_a.len() + ); + } + let layout = |index: usize| { + shard_layouts + .get(index) + .copied() + .unwrap_or(LoraTpShardLayout::LatentRank) + }; let adapter_prefix = adapter_id .map(|id| format!("dynamic_{id}_")) .unwrap_or_default(); @@ -586,6 +742,7 @@ fn append_adapter_shards( format!("{adapter_prefix}a_{index}"), "lora_a", LoraSide::A, + layout(index), tensor, )?); } @@ -598,6 +755,7 @@ fn append_adapter_shards( format!("{adapter_prefix}b_{index}"), "lora_b", LoraSide::B, + layout(index), tensor, )?); } @@ -614,6 +772,7 @@ fn append_adapter_shards( } else { LoraSide::B }, + layout(index / 2), tensor, )?); } @@ -630,6 +789,7 @@ fn append_adapter_shards( } else { LoraSide::B }, + layout(index / 2), tensor, )?); } @@ -645,6 +805,7 @@ fn tensor_shard( tensor_name: String, state: &str, side: LoraSide, + layout: LoraTpShardLayout, tensor: &Tensor, ) -> Result { let tp_size = i64::try_from(parallel.tensor_model_parallel_size) @@ -660,21 +821,60 @@ fn tensor_shard( if local_shape.len() < 2 { bail!("checkpoint tensor {file}:{tensor_name} must have at least two dimensions"); } - let partition_axis = match side { + let rank_axis = match side { LoraSide::A => local_shape.len() - 2, LoraSide::B => local_shape.len() - 1, }; let local_lora_rank = global_lora_rank / tp_size; - if local_shape[partition_axis] != local_lora_rank { - bail!( - "checkpoint tensor {file}:{tensor_name} has local rank {} on axis {partition_axis}, expected {local_lora_rank}", - local_shape[partition_axis] - ); - } let mut global_shape = local_shape.clone(); - global_shape[partition_axis] = global_lora_rank; let mut global_offset = vec![0; local_shape.len()]; - global_offset[partition_axis] = tp_rank * local_lora_rank; + let (partition_axis, replicated) = match (layout, side) { + (LoraTpShardLayout::LatentRank, _) => { + if local_shape[rank_axis] != local_lora_rank { + bail!( + "checkpoint tensor {file}:{tensor_name} has local rank {} on axis {rank_axis}, expected {local_lora_rank}", + local_shape[rank_axis] + ); + } + global_shape[rank_axis] = global_lora_rank; + global_offset[rank_axis] = tp_rank * local_lora_rank; + (rank_axis, false) + } + (LoraTpShardLayout::ColumnParallel, LoraSide::A) + | (LoraTpShardLayout::RowParallel, LoraSide::B) => { + if local_shape[rank_axis] != global_lora_rank { + bail!( + "replicated checkpoint tensor {file}:{tensor_name} has rank {} on axis {rank_axis}, expected {global_lora_rank}", + local_shape[rank_axis] + ); + } + (rank_axis, true) + } + (LoraTpShardLayout::ColumnParallel, LoraSide::B) => { + if local_shape[rank_axis] != global_lora_rank { + bail!( + "column-parallel checkpoint tensor {file}:{tensor_name} has rank {} on axis {rank_axis}, expected {global_lora_rank}", + local_shape[rank_axis] + ); + } + let axis = local_shape.len() - 2; + global_shape[axis] *= tp_size; + global_offset[axis] = tp_rank * local_shape[axis]; + (axis, false) + } + (LoraTpShardLayout::RowParallel, LoraSide::A) => { + if local_shape[rank_axis] != global_lora_rank { + bail!( + "row-parallel checkpoint tensor {file}:{tensor_name} has rank {} on axis {rank_axis}, expected {global_lora_rank}", + local_shape[rank_axis] + ); + } + let axis = local_shape.len() - 1; + global_shape[axis] *= tp_size; + global_offset[axis] = tp_rank * local_shape[axis]; + (axis, false) + } + }; Ok(TensorShardManifest { file: file.to_string(), tensor_name, @@ -684,8 +884,14 @@ fn tensor_shard( global_shape, local_shape, partition_axis, + layout, + replicated, global_offset, - replica_identity: parallel.replica_identity(), + replica_identity: if replicated { + "tp-replicated".to_string() + } else { + parallel.replica_identity() + }, }) } @@ -870,6 +1076,7 @@ mod tests { optimizer_step: 19, target_layers: vec![1, 3], target_modules: vec!["q_proj".into(), "down_proj".into()], + shard_layouts: Vec::new(), parameter_count: 2, optimizer_count: 4, }, @@ -933,6 +1140,7 @@ mod tests { }"#; let manifest: DynamicAdapterManifest = serde_json::from_str(json).unwrap(); assert_eq!(manifest.optimizer_step, 0); + assert!(manifest.shard_layouts.is_empty()); } fn tp_topology(global_rank: usize, tp_size: usize) -> ParallelCheckpointManifest { @@ -947,6 +1155,14 @@ mod tests { (vec![a], vec![b], m, v) } + fn tp_fixed_identities() -> [LoraSlotIdentity; 1] { + [LoraSlotIdentity { + index: 0, + layer: 0, + module: "in_proj_qkv".to_string(), + }] + } + fn tp_dynamic_adapter(value: f64) -> DynamicAdapterCheckpoint { let a = Tensor::full([3, 5], value, (tch::Kind::Float, tch::Device::Cpu)); let b = Tensor::full([7, 3], value + 1.0, (tch::Kind::Float, tch::Device::Cpu)); @@ -958,6 +1174,7 @@ mod tests { optimizer_step: 4, target_layers: vec![1], target_modules: vec!["q_proj".into()], + shard_layouts: Vec::new(), parameter_count: 1, optimizer_count: 2, }, @@ -987,6 +1204,8 @@ mod tests { &m, &v, &[dynamic], + &[], + &tp_fixed_identities(), &topology, ) .unwrap(); @@ -1060,6 +1279,8 @@ mod tests { &m, &v, &[], + &[], + &tp_fixed_identities(), &topology, ) .unwrap(); @@ -1088,6 +1309,8 @@ mod tests { &m, &v, &[], + &[], + &tp_fixed_identities(), &rank0, ) .unwrap(); @@ -1112,4 +1335,222 @@ mod tests { .expect("loading another rank's shard must fail"); assert!(error.to_string().contains("topology mismatch")); } + + #[test] + fn tensor_parallel_projection_layouts_record_global_tensor_geometry() { + let dir = tempfile::tempdir().unwrap(); + let topology = tp_topology(1, 2); + let column_a = Tensor::zeros([4, 8], (tch::Kind::Float, tch::Device::Cpu)); + let column_b = Tensor::zeros([8, 4], (tch::Kind::Float, tch::Device::Cpu)); + let row_a = Tensor::zeros([4, 4], (tch::Kind::Float, tch::Device::Cpu)); + let row_b = Tensor::zeros([8, 4], (tch::Kind::Float, tch::Device::Cpu)); + let lora_a = vec![column_a.shallow_clone(), row_a.shallow_clone()]; + let lora_b = vec![column_b.shallow_clone(), row_b.shallow_clone()]; + let adam_m = vec![ + column_a.zeros_like(), + column_b.zeros_like(), + row_a.zeros_like(), + row_b.zeros_like(), + ]; + let adam_v = vec![ + column_a.ones_like(), + column_b.ones_like(), + row_a.ones_like(), + row_b.ones_like(), + ]; + let layouts = [ + LoraTpShardLayout::ColumnParallel, + LoraTpShardLayout::RowParallel, + ]; + let identities = [ + LoraSlotIdentity { + index: 0, + layer: 0, + module: "q_proj".to_string(), + }, + LoraSlotIdentity { + index: 3, + layer: 0, + module: "o_proj".to_string(), + }, + ]; + + save_checkpoint_with_dynamic_for_topology( + dir.path(), + 3, + 0.5, + "Qwen/test", + 4, + 8.0, + &lora_a, + &lora_b, + &adam_m, + &adam_v, + &[], + &layouts, + &identities, + &topology, + ) + .unwrap(); + + let loaded = load_checkpoint_for_topology(dir.path(), &topology).unwrap(); + assert_eq!(loaded.manifest.fixed_shard_layouts, layouts); + assert_eq!(loaded.manifest.fixed_slot_identities, identities); + let shard = |name: &str| { + loaded + .manifest + .tensor_shards + .iter() + .find(|shard| shard.file == "adapter.safetensors" && shard.tensor_name == name) + .unwrap() + }; + + let column_a_shard = shard("a_0"); + assert_eq!(column_a_shard.layout, LoraTpShardLayout::ColumnParallel); + assert!(column_a_shard.replicated); + assert_eq!(column_a_shard.global_shape, vec![4, 8]); + assert_eq!(column_a_shard.global_offset, vec![0, 0]); + assert_eq!(column_a_shard.replica_identity, "tp-replicated"); + + let column_b_shard = shard("b_0"); + assert!(!column_b_shard.replicated); + assert_eq!(column_b_shard.partition_axis, 0); + assert_eq!(column_b_shard.global_shape, vec![16, 4]); + assert_eq!(column_b_shard.global_offset, vec![8, 0]); + + let row_a_shard = shard("a_1"); + assert_eq!(row_a_shard.layout, LoraTpShardLayout::RowParallel); + assert!(!row_a_shard.replicated); + assert_eq!(row_a_shard.partition_axis, 1); + assert_eq!(row_a_shard.global_shape, vec![4, 8]); + assert_eq!(row_a_shard.global_offset, vec![0, 4]); + + let row_b_shard = shard("b_1"); + assert!(row_b_shard.replicated); + assert_eq!(row_b_shard.global_shape, vec![8, 4]); + assert_eq!(row_b_shard.global_offset, vec![0, 0]); + assert_eq!(row_b_shard.replica_identity, "tp-replicated"); + } + + #[test] + fn tensor_parallel_loader_accepts_legacy_latent_rank_v3_manifest() { + let dir = tempfile::tempdir().unwrap(); + let topology = tp_topology(0, 2); + let (a, b, m, v) = tp_state(1.0); + save_checkpoint_with_dynamic_for_topology( + dir.path(), + 7, + 0.25, + "Qwen/test", + 4, + 8.0, + &a, + &b, + &m, + &v, + &[], + &[], + &tp_fixed_identities(), + &topology, + ) + .unwrap(); + + let manifest_path = dir.path().join("rank-00000/manifest.json"); + let mut manifest: serde_json::Value = + serde_json::from_str(&std::fs::read_to_string(&manifest_path).unwrap()).unwrap(); + let object = manifest.as_object_mut().unwrap(); + object.insert( + "format".to_string(), + serde_json::Value::String(LEGACY_TP_CHECKPOINT_FORMAT.to_string()), + ); + object.remove("fixed_shard_layouts"); + object.remove("fixed_slot_identities"); + for shard in object + .get_mut("tensor_shards") + .unwrap() + .as_array_mut() + .unwrap() + { + let shard = shard.as_object_mut().unwrap(); + shard.remove("layout"); + shard.remove("replicated"); + } + std::fs::write( + &manifest_path, + serde_json::to_string_pretty(&manifest).unwrap(), + ) + .unwrap(); + + let loaded = load_checkpoint_for_topology(dir.path(), &topology).unwrap(); + assert_eq!(loaded.manifest.format, LEGACY_TP_CHECKPOINT_FORMAT); + assert!(loaded.manifest.fixed_shard_layouts.is_empty()); + assert!(loaded + .manifest + .tensor_shards + .iter() + .all(|shard| shard.layout == LoraTpShardLayout::LatentRank && !shard.replicated)); + } + + fn resume_validation_manifest(format: &str) -> CheckpointManifest { + CheckpointManifest { + format: format.to_string(), + step: 0, + loss: 0.0, + model_path: "Qwen/test".to_string(), + lora_rank: 4, + lora_alpha: 8.0, + files: Vec::new(), + dynamic_adapters: Vec::new(), + parallel: None, + tensor_shards: Vec::new(), + fixed_shard_layouts: vec![LoraTpShardLayout::ColumnParallel], + fixed_slot_identities: vec![LoraSlotIdentity { + index: 0, + layer: 0, + module: "q_proj".to_string(), + }], + } + } + + #[test] + fn tensor_parallel_resume_rejects_fixed_slot_identity_mismatch() { + let manifest = resume_validation_manifest(TP_CHECKPOINT_FORMAT); + let expected = [LoraSlotIdentity { + index: 1, + layer: 0, + module: "k_proj".to_string(), + }]; + let error = + validate_fixed_tp_resume(&manifest, &[LoraTpShardLayout::ColumnParallel], &expected) + .unwrap_err(); + assert!(error.to_string().contains("slot identities")); + } + + #[test] + fn legacy_tensor_parallel_resume_rejects_projection_aware_layouts() { + let manifest = resume_validation_manifest(LEGACY_TP_CHECKPOINT_FORMAT); + let fixed_error = + validate_fixed_tp_resume(&manifest, &[LoraTpShardLayout::ColumnParallel], &[]) + .unwrap_err(); + assert!(fixed_error.to_string().contains("fixed Q/K/V/O")); + let dynamic_error = + validate_dynamic_tp_resume(&manifest, 17, &[], &[LoraTpShardLayout::RowParallel]) + .unwrap_err(); + assert!(dynamic_error.to_string().contains("adapter 17")); + validate_fixed_tp_resume(&manifest, &[LoraTpShardLayout::LatentRank], &[]).unwrap(); + } + + #[test] + fn fixed_restore_mapping_supports_compact_and_legacy_positional_slots() { + assert_eq!( + fixed_restore_slot_indices(2, 2, &[1, 4], 7).unwrap(), + vec![1, 4] + ); + assert_eq!( + fixed_restore_slot_indices(7, 7, &[1, 4], 7).unwrap(), + (0..7).collect::>() + ); + assert!(fixed_restore_slot_indices(2, 1, &[1, 4], 7).is_err()); + assert!(fixed_restore_slot_indices(3, 3, &[1, 4], 7).is_err()); + } } diff --git a/crates/rustrain-server/src/session.rs b/crates/rustrain-server/src/session.rs index d34fffbd..ba268017 100644 --- a/crates/rustrain-server/src/session.rs +++ b/crates/rustrain-server/src/session.rs @@ -1,6 +1,6 @@ //! Training session trait + Qwen3.6 implementation. -use anyhow::{Context, Result, anyhow, bail}; +use anyhow::{anyhow, bail, Context, Result}; use std::path::PathBuf; use std::sync::Arc; use tch::{Device, Kind, Tensor}; @@ -10,6 +10,16 @@ use crate::checkpoint; use crate::metrics::{FileMetricsSink, MetricsSink, StepMetric}; use rustrain_qwen3_6::lora::{Qwen36AdapterArtifact, Qwen36LoraConfig, Qwen36LoraTargetModule}; +fn lora_tp_shard_layout(module: Qwen36LoraTargetModule) -> checkpoint::LoraTpShardLayout { + match module { + Qwen36LoraTargetModule::QProj + | Qwen36LoraTargetModule::KProj + | Qwen36LoraTargetModule::VProj => checkpoint::LoraTpShardLayout::ColumnParallel, + Qwen36LoraTargetModule::OProj => checkpoint::LoraTpShardLayout::RowParallel, + _ => checkpoint::LoraTpShardLayout::LatentRank, + } +} + /// Session states. #[derive(Debug, Clone)] pub enum SessionState { @@ -334,13 +344,41 @@ impl TrainingSession for Qwen36Session { } let is_ep = ep_world_size > 1 && runtime_config.is_moe && tp_size == 1; let is_data_parallel = ep_world_size > 1 && !runtime_config.is_moe && tp_size == 1; + let base_tp_attention = tp_size > 1; let base_tp_mlp = tp_size > 1 && !runtime_config.is_moe; - if base_tp_mlp { + if base_tp_attention { if runtime_config.mtp_num_hidden_layers > 0 { return Err(anyhow!( - "base dense TP MLP currently requires MTP to be disabled" + "frozen base TP currently requires MTP to be disabled" )); } + if runtime_config.num_attention_heads <= 0 + || runtime_config.num_attention_heads % tp_size as i64 != 0 + || runtime_config.num_key_value_heads <= 0 + || runtime_config.num_key_value_heads % tp_size as i64 != 0 + || runtime_config.num_attention_heads % runtime_config.num_key_value_heads != 0 + { + return Err(anyhow!( + "full-attention heads (q={}, kv={}) must preserve GQA groups and be divisible by TP_SIZE={tp_size}", + runtime_config.num_attention_heads, + runtime_config.num_key_value_heads + )); + } + let rotary_dim = + (runtime_config.head_dim as f64 * runtime_config.partial_rotary_factor) as i64; + if runtime_config.head_dim <= 0 + || rotary_dim < 0 + || rotary_dim > runtime_config.head_dim + || rotary_dim % 2 != 0 + { + return Err(anyhow!( + "full-attention head_dim={} and partial_rotary_factor={} produce invalid rotary_dim={rotary_dim}", + runtime_config.head_dim, + runtime_config.partial_rotary_factor + )); + } + } + if base_tp_mlp { if runtime_config.intermediate_size <= 0 || runtime_config.intermediate_size % tp_size as i64 != 0 { @@ -406,13 +444,24 @@ impl TrainingSession for Qwen36Session { .to_kind(self.compute_kind); weights.insert(name, narrowed); } else { - let local_shard = if base_tp_mlp { - rustrain_qwen3_6::kernel::shard_dense_mlp_weight_for_tp( - &name, - &tensor, - tp_size, - ep_rank % tp_size, - )? + let local_shard = if base_tp_attention { + let attention_shard = + rustrain_qwen3_6::kernel::shard_full_attention_weight_for_tp( + &name, + &tensor, + tp_size, + ep_rank % tp_size, + )?; + if attention_shard.is_some() || !base_tp_mlp { + attention_shard + } else { + rustrain_qwen3_6::kernel::shard_dense_mlp_weight_for_tp( + &name, + &tensor, + tp_size, + ep_rank % tp_size, + )? + } } else { None }; @@ -457,6 +506,7 @@ impl TrainingSession for Qwen36Session { req.eps, lora_scaling, req.rank, + base_tp_attention, base_tp_mlp, &all_layers, &target_modules, @@ -604,30 +654,83 @@ impl TrainingSession for Qwen36Session { .as_ref() .ok_or_else(|| anyhow!("LoRA not initialized"))?; - let lora_count = ctx.lora_count(); - let mut lora_a = Vec::new(); - let mut lora_b = Vec::new(); - for i in 0..lora_count { - if let (Some(a), Some(b)) = (ctx.get_lora_a(i as i64), ctx.get_lora_b(i as i64)) { - lora_a.push(a); - lora_b.push(b); - } + let model_path = self + .model_path + .as_ref() + .ok_or_else(|| anyhow!("model path unavailable for checkpoint"))?; + let runtime_config = + rustrain_qwen3_6::config::read_qwen36_runtime_config(std::path::Path::new(model_path))?; + let parallel = checkpoint::ParallelCheckpointManifest::from_env()?; + + let fixed_config = Qwen36LoraConfig { + rank: self.lora_rank, + alpha: self.lora_alpha, + target_layers: self.lora_target_layers.clone(), + target_modules: self.lora_target_modules.clone(), + }; + let fixed_slots = rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, &fixed_config); + let lora_count = ctx.lora_count() as usize; + if fixed_slots.len() != lora_count { + bail!( + "fixed LoRA registry count {} does not match native slot count {lora_count}", + fixed_slots.len() + ); + } + let (all_adam_m, all_adam_v) = ctx.export_optimizer_state()?; + let expected_optimizer_count = lora_count.saturating_mul(2); + if all_adam_m.len() != expected_optimizer_count + || all_adam_v.len() != expected_optimizer_count + { + bail!( + "fixed optimizer state count mismatch: m={}, v={}, expected={expected_optimizer_count}", + all_adam_m.len(), + all_adam_v.len() + ); } - // Export Adam optimizer state - let (adam_m, adam_v) = ctx.export_optimizer_state()?; + // TP v4 can compact inactive native slots because the manifest records + // exact identities. Keep the positional full-slot representation for + // v1/v2, whose manifests have no fixed-slot identity metadata. + let saved_fixed_slots = fixed_slots + .iter() + .filter(|slot| parallel.tensor_model_parallel_size <= 1 || slot.active) + .collect::>(); + let saved_fixed_count = saved_fixed_slots.len(); + let mut lora_a = Vec::with_capacity(saved_fixed_count); + let mut lora_b = Vec::with_capacity(saved_fixed_count); + let mut adam_m = Vec::with_capacity(saved_fixed_count.saturating_mul(2)); + let mut adam_v = Vec::with_capacity(saved_fixed_count.saturating_mul(2)); + let mut fixed_shard_layouts = Vec::with_capacity(saved_fixed_count); + let mut fixed_slot_identities = Vec::with_capacity(saved_fixed_count); + for slot in saved_fixed_slots { + lora_a.push(ctx.get_lora_a(slot.index as i64).with_context(|| { + format!("fixed LoRA A is missing for native slot {}", slot.index) + })?); + lora_b.push(ctx.get_lora_b(slot.index as i64).with_context(|| { + format!("fixed LoRA B is missing for native slot {}", slot.index) + })?); + let optimizer_index = slot.index.saturating_mul(2); + adam_m.push(all_adam_m[optimizer_index].shallow_clone()); + adam_m.push(all_adam_m[optimizer_index + 1].shallow_clone()); + adam_v.push(all_adam_v[optimizer_index].shallow_clone()); + adam_v.push(all_adam_v[optimizer_index + 1].shallow_clone()); + fixed_shard_layouts.push(lora_tp_shard_layout(slot.module)); + fixed_slot_identities.push(checkpoint::LoraSlotIdentity { + index: slot.index, + layer: slot.layer, + module: slot.module.cpp_name().to_string(), + }); + } let mut dynamic_adapters = Vec::new(); if !self.dynamic_lora_configs.is_empty() { - let model_path = self - .model_path - .as_ref() - .ok_or_else(|| anyhow!("model path unavailable for dynamic LoRA checkpoint"))?; - let runtime_config = rustrain_qwen3_6::config::read_qwen36_runtime_config( - std::path::Path::new(model_path), - )?; for (&adapter_id, lora_config) in &self.dynamic_lora_configs { let slots = rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, lora_config); + let shard_layouts = slots + .iter() + .filter(|slot| slot.active) + .map(|slot| lora_tp_shard_layout(slot.module)) + .collect::>(); let mut dynamic_a = Vec::new(); let mut dynamic_b = Vec::new(); let mut dynamic_m = Vec::new(); @@ -738,6 +841,7 @@ impl TrainingSession for Qwen36Session { .iter() .map(|module| module.cpp_name().to_string()) .collect(), + shard_layouts, parameter_count: dynamic_a.len(), optimizer_count: dynamic_m.len(), }, @@ -749,12 +853,11 @@ impl TrainingSession for Qwen36Session { } } - let parallel = checkpoint::ParallelCheckpointManifest::from_env()?; checkpoint::save_checkpoint_with_dynamic_for_topology( std::path::Path::new(path), self.step, self.last_loss, - self.model_path.as_deref().unwrap_or(""), + model_path, self.lora_rank, self.lora_alpha, &lora_a, @@ -762,6 +865,8 @@ impl TrainingSession for Qwen36Session { &adam_m, &adam_v, &dynamic_adapters, + &fixed_shard_layouts, + &fixed_slot_identities, ¶llel, )?; @@ -771,14 +876,69 @@ impl TrainingSession for Qwen36Session { fn load_checkpoint(&mut self, path: &str) -> Result<(u64, f64)> { let parallel = checkpoint::ParallelCheckpointManifest::from_env()?; let data = checkpoint::load_checkpoint_for_topology(std::path::Path::new(path), ¶llel)?; - if !data.dynamic_adapters.is_empty() { - let model_path = self - .model_path - .as_ref() - .ok_or_else(|| anyhow!("model path unavailable for dynamic LoRA checkpoint"))?; - let runtime_config = rustrain_qwen3_6::config::read_qwen36_runtime_config( - std::path::Path::new(model_path), + let model_path = self + .model_path + .as_ref() + .ok_or_else(|| anyhow!("model path unavailable for checkpoint restore"))?; + let runtime_config = + rustrain_qwen3_6::config::read_qwen36_runtime_config(std::path::Path::new(model_path))?; + let fixed_config = Qwen36LoraConfig { + rank: self.lora_rank, + alpha: self.lora_alpha, + target_layers: self.lora_target_layers.clone(), + target_modules: self.lora_target_modules.clone(), + }; + let fixed_slots = rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, &fixed_config); + let expected_fixed_layouts = fixed_slots + .iter() + .filter(|slot| slot.active) + .map(|slot| lora_tp_shard_layout(slot.module)) + .collect::>(); + let expected_fixed_identities = fixed_slots + .iter() + .filter(|slot| slot.active) + .map(|slot| checkpoint::LoraSlotIdentity { + index: slot.index, + layer: slot.layer, + module: slot.module.cpp_name().to_string(), + }) + .collect::>(); + if parallel.tensor_model_parallel_size > 1 { + checkpoint::validate_fixed_tp_resume( + &data.manifest, + &expected_fixed_layouts, + &expected_fixed_identities, )?; + } + if parallel.tensor_model_parallel_size > 1 { + for dynamic in &data.dynamic_adapters { + let target_modules = dynamic + .manifest + .target_modules + .iter() + .map(|name| Qwen36LoraTargetModule::parse(name)) + .collect::>>()?; + let config = Qwen36LoraConfig { + rank: dynamic.manifest.rank, + alpha: dynamic.manifest.alpha, + target_layers: dynamic.manifest.target_layers.clone(), + target_modules, + }; + let expected_layouts = + rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, &config) + .iter() + .filter(|slot| slot.active) + .map(|slot| lora_tp_shard_layout(slot.module)) + .collect::>(); + checkpoint::validate_dynamic_tp_resume( + &data.manifest, + dynamic.manifest.id, + &dynamic.manifest.shard_layouts, + &expected_layouts, + )?; + } + } + if !data.dynamic_adapters.is_empty() { if !self.dynamic_lora_configs.is_empty() { bail!("cannot load dynamic LoRA checkpoint into a session with active adapters"); } @@ -908,23 +1068,77 @@ impl TrainingSession for Qwen36Session { } // Import Adam optimizer state into C++ context if let Some(ctx) = &self.ctx { - if data.lora_a.len() != data.lora_b.len() - || data.lora_a.len() != ctx.lora_count() as usize + let native_slot_count = ctx.lora_count() as usize; + if fixed_slots.len() != native_slot_count { + bail!( + "fixed LoRA registry count {} does not match native slot count {native_slot_count}", + fixed_slots.len() + ); + } + let active_slot_indices = fixed_slots + .iter() + .filter(|slot| slot.active) + .map(|slot| slot.index) + .collect::>(); + let restore_slot_indices = checkpoint::fixed_restore_slot_indices( + data.lora_a.len(), + data.lora_b.len(), + &active_slot_indices, + native_slot_count, + )?; + for ((a, b), &slot_index) in data + .lora_a + .iter() + .zip(&data.lora_b) + .zip(&restore_slot_indices) { - return Err(anyhow!( - "checkpoint LoRA slot count mismatch: checkpoint A/B={}/{}, context={}", - data.lora_a.len(), - data.lora_b.len(), - ctx.lora_count() - )); + ctx.set_lora_tensor(slot_index as i64, false, a)?; + ctx.set_lora_tensor(slot_index as i64, true, b)?; } - for (index, (a, b)) in data.lora_a.iter().zip(&data.lora_b).enumerate() { - ctx.set_lora_tensor(index as i64, false, a)?; - ctx.set_lora_tensor(index as i64, true, b)?; + if data.adam_m.is_empty() != data.adam_v.is_empty() { + bail!( + "checkpoint fixed optimizer m/v count mismatch: {}/{}", + data.adam_m.len(), + data.adam_v.len() + ); } - if !data.adam_m.is_empty() && !data.adam_v.is_empty() { - ctx.import_optimizer_state(&data.adam_m, &data.adam_v)?; - tracing::info!(imported = data.adam_m.len(), "optimizer state imported"); + if !data.adam_m.is_empty() { + let expected_saved_optimizer_count = restore_slot_indices.len().saturating_mul(2); + if data.adam_m.len() != expected_saved_optimizer_count + || data.adam_v.len() != expected_saved_optimizer_count + { + bail!( + "checkpoint fixed optimizer count mismatch: m={}, v={}, expected={expected_saved_optimizer_count}", + data.adam_m.len(), + data.adam_v.len() + ); + } + let (mut all_adam_m, mut all_adam_v) = ctx.export_optimizer_state()?; + let expected_native_optimizer_count = native_slot_count.saturating_mul(2); + if all_adam_m.len() != expected_native_optimizer_count + || all_adam_v.len() != expected_native_optimizer_count + { + bail!( + "native fixed optimizer count mismatch: m={}, v={}, expected={expected_native_optimizer_count}", + all_adam_m.len(), + all_adam_v.len() + ); + } + for (saved_slot, &native_slot) in restore_slot_indices.iter().enumerate() { + let saved = saved_slot.saturating_mul(2); + let native = native_slot.saturating_mul(2); + all_adam_m[native] = data.adam_m[saved].shallow_clone(); + all_adam_m[native + 1] = data.adam_m[saved + 1].shallow_clone(); + all_adam_v[native] = data.adam_v[saved].shallow_clone(); + all_adam_v[native + 1] = data.adam_v[saved + 1].shallow_clone(); + } + let imported = ctx.import_optimizer_state(&all_adam_m, &all_adam_v)?; + if imported != expected_native_optimizer_count as i64 { + bail!( + "native fixed optimizer import restored {imported} tensors, expected {expected_native_optimizer_count}" + ); + } + tracing::info!(imported, "optimizer state imported"); } let native_step = i64::try_from(data.manifest.step) .context("checkpoint step exceeds the native optimizer range")?; diff --git a/docs/plans/qwen-lora-megatron-progress.md b/docs/plans/qwen-lora-megatron-progress.md index a9db0caa..1149c393 100644 --- a/docs/plans/qwen-lora-megatron-progress.md +++ b/docs/plans/qwen-lora-megatron-progress.md @@ -12,9 +12,9 @@ timestamp: 2026-07-17T00:00:00Z # Current State -Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP parity against a full-expert reference, variable-split EP A2A with fixed-LoRA data sharding, dense replicated-DP smoke with per-tenant token weighting, TP-only latent-rank-sharded LoRA smoke, frozen dense SwiGLU MLP base-weight TP, dynamic batch logical-step update, ABI11 FP32 gradient storage/aggregation, standard Adam bias correction, per-tenant optimizer-step restore, selected-tenant isolation, same-topology rank-aware checkpointing, and 5D topology mapping. +Verified: Qwen3.5/3.6 native forward/backward slices, GDN CUDA path, grouped MoE fallback, EP parity against a full-expert reference, variable-split EP A2A with fixed-LoRA data sharding, dense replicated-DP smoke with per-tenant token weighting, TP-only latent-rank-sharded LoRA smoke, frozen full-attention and dense SwiGLU MLP base-weight TP, projection-aware Q/K/V/O fixed and selected-dynamic LoRA TP, dynamic batch logical-step update, ABI11 FP32 gradient storage/aggregation, standard Adam bias correction, per-tenant optimizer-step restore, selected-tenant isolation, checkpoint v4 projection layouts and fixed-slot identities, and 5D topology mapping. -Not yet verified or implemented: base-weight TP for attention/GDN/MoE/embedding/LM-head, MLP-targeted LoRA under base TP, multi-axis TP+DP/EP, PP/CP, DeepEP/TE prebuilt integration, server-side source-sharded dynamic dispatch, heterogeneous dynamic adapter signatures, dynamic+MTP objective normalization, cross-topology checkpoint resharding, and matched Megatron throughput. Native direct dynamic source metadata under sharded A2A is implemented and full-reference smoke-tested. +Not yet verified or implemented: base-weight TP for GDN/MoE/embedding/LM-head, MLP-targeted LoRA under base TP, fused QKV/gate-up and sequence parallel, multi-axis TP+DP/EP, PP/CP, DeepEP/TE/FLA prebuilt integration, server-side source-sharded dynamic dispatch, heterogeneous dynamic adapter signatures, dynamic+MTP objective normalization, cross-topology checkpoint resharding, old-v3 attention checkpoint migration, and matched Megatron throughput. Native direct dynamic source metadata under sharded A2A is implemented and full-reference smoke-tested. # Durable Milestones @@ -40,6 +40,9 @@ Not yet verified or implemented: base-weight TP for attention/GDN/MoE/embedding/ - H20 `123.57.26.97:28004`: `native_ep_bench.cpp` fresh ABI0 benchmark (`seq=128, hidden=256, experts=8, intermediate=256, warmup=2, iters=10`) passed legacy and sharded A2A with `rank_statuses=0,0`. Legacy median was about `5.47 ms` / `46.4k processed tokens/s` (`23.2k unique tokens/s`), while sharded A2A was about `6.72 ms` / `37.8k processed and unique tokens/s`. This is a synthetic native baseline, not Megatron-LM parity. - H20 `123.57.26.97:28004`: fresh ABI1 dynamic sharded native smoke passed on both ranks (`rank_statuses=0,0`) against a full-expert reference with complementary source masks. Dynamic grouped-expert parameter/m/v maxima were `1.53e-5` / `4.88e-5` / `5.75e-8`; it also exercised a third tenant with zero global target tokens, clocks `[2,2,0]`, and explicit rejection of ordinary `train_step`. The server path still broadcasts replicated source batches, and dynamic+MTP is explicitly rejected until its two objective denominators are separated. - H20 `123.57.26.97:28004`: fresh ABI13 dense base-MLP TP2 smoke passed on both ranks. CLI and server use one CPU sharding helper for gate/up rows and matching down columns; the per-context native flag validates local shapes without process-global env state. C++ reduces the row-parallel output and all-reduces column-parallel input dgrad. Eval/train loss differed from the replicated full-weight reference by `1.84e-5`; FP32 LoRA gradient-accumulator maxima were `3.05e-5` / `4.58e-5`, and the largest post-Adam LoRA slice difference was `9.31e-9`. The smoke also rejects MLP LoRA targets, prevents incompatible TP state transitions, and covers eval-to-train cache invalidation. +- ABI14 working tree: frozen full attention shards Q/K/V output heads and O input columns, while fixed and selected dynamic Q/K/V/O LoRA use projection-aware layouts. Replicated-side gradients are summed once at the optimizer boundary. Latent-rank LoRA now applies copy-to-TP-region before its local A/B path so a later sharded branch sums input dgrad before it reaches preceding replicated layers. Checkpoint v4 records the layout, replicated tensor geometry, and exact fixed native slot identity; v3 remains loadable only for latent-rank tensors and attention migration is explicitly rejected. +- H20 target: ABI14 4Q/2KV-head full-attention TP2 smoke passed Q/K/V/O fixed and selected-dynamic full-reference oracles on both ranks. Fixed eval/loss differed by `4.40e-4`; Q/K/V-B and O-A gradient maxima were `1.83e-4` / `1.14e-5` / `3.05e-4` / `1.83e-4`, FP32 m/v maxima were `1.53e-5` / `6.63e-9`, and the local standard-Adam formula error was below `3.73e-9`. Selected-dynamic loss differed by `6.82e-5` and the largest Q/K/V/O parameter difference was `3.05e-5`. +- H20 target: ABI14 two-layer GDN latent-rank TP2 full-reference smoke passed on both ranks. Eval/loss differed by `5.46e-5`; A/B gradient maxima were `9.54e-7` / `4.77e-7`, FP32 m/v maxima were `4.77e-8` / `3.71e-14`, and the local standard-Adam formula error was zero. ABI14 general single-GPU native smoke and dense-MLP TP2 regression also passed. # Decisions During Execution @@ -47,10 +50,12 @@ Not yet verified or implemented: base-weight TP for attention/GDN/MoE/embedding/ - Publish multi-LoRA `n_max` directly from rank 0 with `ncclBroadcast`; filesystem rendezvous can reuse stale files across process restarts and give ranks different chunk schedules. - Do not enable Qwen TP/PP/CP by merely relaxing runtime validation. - Treat dense base-MLP TP as one accepted slice, not full model TP. Until projection-aware LoRA collectives exist, reject gate/up/down LoRA targets instead of applying the replicated-projection reduction rule to disjoint output shards. +- For frozen full-attention TP, use Q/K/V output-head sharding and O input-column sharding. Q/K/V replicate LoRA A and shard B; O shards A and replicates B. Sum only the replicated-side gradient at the optimizer boundary so the activation collective is not duplicated. +- Do not reinterpret v3 attention LoRA tensors as projection-aware shards. Their latent-rank geometry is incompatible, so require v4 for Q/K/V/O resume and keep v3 compatibility for latent-rank-only modules. - Treat Exa/Jina dependency search failures as missing evidence, not as proof that a package is compatible. # Verification -Passed: `cargo check -p rustrain-qwen3-6 -p rustrain-server -p rustrain-ipc` (with the repository host venv), `cargo test -p rustrain-qwen3-6 --lib` (5), `cargo test -p rustrain-server --lib` (6), Qwen integration tests, remote ABI8 smoke, remote ABI9 selected-tenant native smoke, remote ABI10 single-rank native smoke, remote ABI10 two-rank TP native smoke, remote ABI11 single/TP2/DP2/EP2 native smoke with numerical Adam and parity oracles, ABI1 dynamic sharded full-reference smoke, and ABI13 dense base-MLP TP2 parity smoke. +Passed: `cargo check -p rustrain-qwen3-6 -p rustrain-server -p rustrain-ipc` (with the repository host venv), `cargo test -p rustrain-qwen3-6 --lib`, `cargo test -p rustrain-server checkpoint::tests --lib` (11), Qwen integration tests, remote ABI8 smoke, remote ABI9 selected-tenant native smoke, remote ABI10 single-rank native smoke, remote ABI10 two-rank TP native smoke, remote ABI11 single/TP2/DP2/EP2 native smoke with numerical Adam and parity oracles, ABI1 dynamic sharded full-reference smoke, ABI13 dense base-MLP TP2 parity smoke, and ABI14 two-layer latent TP2, GQA full-attention TP2, general-native, and dense-MLP regression smokes. -Not run: full-model base TP beyond dense MLP, multi-axis TP+DP/EP, PP/CP, server-side source-sharded dynamic dispatch, heterogeneous dynamic adapter signatures, dynamic+MTP, cross-topology resharding, and matched Megatron performance benchmark. The target host lacks importable Megatron/Transformer Engine/DeepEP/flash-attn prebuilt packages, so no dependency installation or JIT workaround was used. +Not run: full-model base TP beyond full attention plus dense MLP, multi-axis TP+DP/EP, PP/CP, server-side source-sharded dynamic dispatch, heterogeneous dynamic adapter signatures, dynamic+MTP, cross-topology resharding, and matched Megatron performance benchmark. The target host lacks importable Megatron/Transformer Engine/FLA/DeepEP/flash-attn prebuilt packages, so no dependency installation or JIT workaround was used. diff --git a/docs/qwen35-qwen36-megatron-audit.md b/docs/qwen35-qwen36-megatron-audit.md index 41793b7d..cbdba779 100644 --- a/docs/qwen35-qwen36-megatron-audit.md +++ b/docs/qwen35-qwen36-megatron-audit.md @@ -4,10 +4,10 @@ ## 结论 -- 模型语义：Qwen3.5 dense、Qwen3.6 dense/MoE 的 native forward/backward 路径已经覆盖 hybrid full attention、GDN、MoE、MTP 和 LoRA 目标模块；已有配置解析、集成测试及 H20 native smoke 证据。 +- 模型语义：Qwen3.5 dense、Qwen3.6 dense/MoE 的 native forward/backward 路径已经覆盖 hybrid full attention、GDN、MoE、MTP 和 LoRA 目标模块；已有配置解析、合成 oracle、集成测试及 H20 native smoke 证据，但尚未完成真实 35B/3.6 权重的长时间训练验证。 - 已实现并可验证的分布式子集：MoE expert parallel，以及 replicated LoRA 的 data parallel；梯度累积和 dynamic multi-LoRA 已有 logical-step 边界。DP 动态租户按 adapter token count 加权，sharded A2A native 路径会保留 source flattened row 来恢复租户，并按全局租户 token count 归一化。 - 性能：MoE grouped dispatch 相对逐 expert matmul 的已有 microbenchmark 为约 3.70x（E=32, N=4096, H=2048, I=768，结果误差为 0）；这不是端到端训练吞吐或 Megatron 对比。 -- 已实现 LoRA latent-rank TP-only，以及 frozen dense SwiGLU MLP 的 gate/up row shard、down column shard 和输出 all-reduce。attention/GDN/MoE/vocab 仍复制，MLP LoRA 在 base TP 下暂拒绝；PP/CP 和 TP 与 EP/DP 的组合仍未实现。 +- 已实现 LoRA latent-rank TP-only、frozen full-attention TP（Q/K/V ColumnParallel、O RowParallel），以及 frozen dense SwiGLU MLP 的 gate/up row shard、down column shard 和输出 all-reduce。Q/K/V/O fixed 与 selected dynamic LoRA 使用 projection-aware shard 和梯度 reduction；GDN/MoE/embedding/LM-head 仍复制，MLP LoRA 在 base TP 下暂拒绝；PP/CP 和 TP 与 EP/DP 的组合仍未实现。 - 因此当前实现不能宣称“Megatron-LM 级别”。它是一个计算集中在 C++ 的 LoRA/EP/DP 子集，离 Megatron 的完整并行和通信重叠仍有实质差距。 ## 当前能力矩阵 @@ -23,16 +23,18 @@ | microbatch accumulation | 已实现子集 | non-final microbatch 只 backward，final microbatch 才 optimizer；FP32 accumulator 存储/聚合，autograd leaf backward 仍为 BF16 | | replicated data parallel | 已实现 | logical-step 边界同步 replicated LoRA；EP expert 参数不走该 reduction | | expert parallel | 已实现子集 | 默认 routed-output all-reduce；gated variable-split A2A 已验证 fixed-LoRA 和 native dynamic-LoRA data sharding；GPU-only split planning、异步 overlap 和 DeepEP backend 未实现 | -| tensor parallel | dense MLP 子集 | latent rank 分片和独立 TP communicator 已验证；dense gate/up/down base 权重按 intermediate 维切分并有 TP2 full-reference smoke；attention/GDN/MoE/LM-head 仍不切分，MLP LoRA/MTP 暂拒绝 | +| tensor parallel | full attention + dense MLP 子集 | full attention 的 Q/K/V 输出头分片、O 输入列分片和 projection-aware fixed/dynamic LoRA 已通过 TP2 full-reference smoke；dense gate/up/down base 权重按 intermediate 维切分；GDN/MoE/embedding/LM-head 仍不切分，MLP LoRA/MTP 暂拒绝 | | pipeline parallel | 未实现于 Qwen native | 没有 stage 切分、microbatch scheduler 或 activation send/recv | | context parallel | 未实现于 Qwen native | 没有 ring attention、跨 rank KV/索引合并 | -| distributed checkpoint | 已实现子集 | same-topology TP rank-sharded v3 已验证；跨 topology reshard 和 PP/CP 未实现 | +| distributed checkpoint | 已实现子集 | v4 记录 projection layout、replicated tensor geometry 和 fixed slot identity，并验证 same-topology rank shard；v3 仅兼容 latent-rank checkpoint，旧 attention LoRA 因形状不可迁移而明确拒绝；跨 topology reshard 和 PP/CP 未实现 | ## 与 Megatron-LM 的关键差距 ### 并行语义 -Megatron 的通用 MLP 使用 ColumnParallel fc1、local gated activation 和 RowParallel fc2，并进一步提供 fused fc1/activation、sequence parallel、通信重叠和 sharded-state 支持。当前 Qwen native 已补上算法等价的 separate gate/up row shard 与 down column shard，但仍是两次独立 GEMM，且 attention/GDN/MoE/vocab 不分片。Megatron-LM 本仓库也没有直接的 Qwen3.5/3.6 模型实现，精确模型入口依赖外部 bridge，因此这里比较的是成熟的通用并行基础设施，不是同模型端到端实现。 +Megatron 的通用 MLP 使用 ColumnParallel fc1、local gated activation 和 RowParallel fc2，并进一步提供 fused fc1/activation、sequence parallel、通信重叠和 sharded-state 支持。当前 Qwen native 已补上算法等价的 separate gate/up row shard 与 down column shard，但仍是两次独立 GEMM。 + +本地 Megatron-LM 的 `experimental/lite/megatron/lite/model/qwen3_5` 已包含直接的 Qwen3.5 实现和 LoRA adapter，而不是只能依赖外部 bridge：`primitive/modules/gqa.py` 使用 fused QKV ColumnParallel 和 O RowParallel；`gated_delta_net.py` 对输入/输出投影和 local heads 做 TP，并接入 sequence parallel、context parallel 与 FLA 高性能路径；`lora.py` 的 `LinearLoRA` 还能拆分 latent rank/output 并配套 gather、reduce-scatter 或 input-gradient all-reduce。没有发现显式的 Qwen3.6 model registration。当前 Qwen native 的 full-attention TP 数学布局已经对齐，但 Q/K/V 仍是三次独立 GEMM，LoRA 的 replicated 一侧也比 Megatron Lite 更保守；GDN、MoE、embedding/LM-head 和多轴并行仍是主要差距。 PP 会把层分到不同 stage 并使用 1F1B 等调度；CP 会在序列和 attention state 上做跨 rank 通信。当前 Qwen native `TrainingContext` 仍在每个进程执行完整层栈，因此仅增加 PP/CP 配置不能得到正确语义。 @@ -44,22 +46,24 @@ PP 会把层分到不同 stage 并使用 1F1B 等调度；CP 会在序列和 att ### 性能工程 -当前粗粒度 C++ FFI、grouped MoE 和 activation checkpoint/offload 是有效优化，但尚无 Megatron/Transformer Engine 级别的端到端数据：没有完整模型在同一 GPU、序列长度、microbatch、精度和通信配置下的 tokens/s、显存、扩展效率对照，也没有 FP8/FP4 参数与 fused attention/DeepEP 的 Qwen 路径。 +当前粗粒度 C++ FFI、grouped MoE 和 activation checkpoint/offload 是有效优化，但 full attention 仍通过 ATen linear/SDPA，QKV 未融合，TP collective 同步执行，GDN 和 vocab 路径仍复制。尚无 Megatron/Transformer Engine 级别的端到端数据：没有完整模型在同一 GPU、序列长度、microbatch、精度和通信配置下的 tokens/s、显存、扩展效率对照，也没有 FP8/FP4 参数与 fused attention/DeepEP 的 Qwen 路径。 本次 native benchmark 没有证明 gated A2A 的端到端 step-time 优势：在该小型 workload 上 sharded A2A 的中位 step 反而比 legacy 高约 `23%`。它没有实现 DeepEP 的 fused permutation、GPU-only split planning 或通信计算 overlap，也没有覆盖 H=`2048`/E=`256` 的完整 Qwen3.6 workload。legacy 模式复制输入 batch，因此必须同时报告唯一样本吞吐，不能只看所有 rank 的 processed tokens/s。 -目标 H20 的 ABI1 环境有 PyTorch 2.12.1、Triton 和 NumPy，但没有 Megatron、Transformer Engine、DeepEP、flash-attn、Apex 或缓存的兼容 prebuilt wheel。本地 Megatron 的 Qwen3.5 35B-A3B 入口强制 TE/flash-attn，且是 full-parameter SFT，不提供 trainable LoRA wrapper；其 `moe_perf` 也固定 TE grouped MLP 和 H100 条件。因此当前不能诚实地产出 matched Megatron-LoRA benchmark，且本工作没有通过 JIT 或自构建依赖绕过该限制。 +目标 H20 的 ABI1 环境有 PyTorch 2.12.1、Triton 和 NumPy，但没有 Megatron、Transformer Engine、FLA、DeepEP、flash-attn、Apex 或缓存的兼容 prebuilt wheel。虽然本地 Megatron Lite 源码包含 Qwen3.5 LoRA adapter，目标机仍不具备其高性能依赖。因此当前不能诚实地产出 matched Megatron-LoRA benchmark，且本工作没有通过 JIT 或自构建依赖绕过该限制。 ## 验证边界 -已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试，以及 H20 ABI11 的单卡、TP2、EP2 和 DP2 native smoke。DP2 的 weighted m/grouped/v/Adam BF16 delta oracle 分别达到 `2.43e-8`、`2.27e-8`、`7.33e-8` 和 `0`；legacy EP2 与 full-expert reference 的 loss、LoRA、Adam state 和标准 Adam 首步 oracle 在两 rank 均为零差异。Replicated A2A 与 fixed-LoRA sharded A2A 也均通过两 rank full-expert reference；sharded token counts `[1,3]` 的加权 loss 与 global reference 相差约 `9.6e-7`，m/v 最大差 `1.22e-5` / `3.92e-9`，Adam oracle 差为 `0`。H20 ABI1 dynamic sharded full-reference smoke 在两 rank 返回 `0`：dynamic grouped-expert 参数最大差 `1.53e-5`，m/v 最大差 `4.88e-5` / `5.75e-8`。新增 ABI13 dense base-MLP TP2 smoke 对 gate/up/down 使用半尺寸本地权重，验证 row-parallel forward sum 与 column-parallel input dgrad sum；eval/train loss 与完整权重参考相差 `1.84e-5`，并直接比较 FP32 LoRA accumulator 和更新切片。没有完成完整大模型长时间训练、跨节点通信、完整 base TP、PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖且有直接 oracle 的子集，而不是所有并行配置。 +已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试，以及 H20 ABI11 的单卡、TP2、EP2 和 DP2 native smoke。DP2 的 weighted m/grouped/v/Adam BF16 delta oracle 分别达到 `2.43e-8`、`2.27e-8`、`7.33e-8` 和 `0`；legacy EP2 与 full-expert reference 的 loss、LoRA、Adam state 和标准 Adam 首步 oracle 在两 rank 均为零差异。Replicated A2A 与 fixed-LoRA sharded A2A 也均通过两 rank full-expert reference；sharded token counts `[1,3]` 的加权 loss 与 global reference 相差约 `9.6e-7`，m/v 最大差 `1.22e-5` / `3.92e-9`，Adam oracle 差为 `0`。H20 ABI1 dynamic sharded full-reference smoke 在两 rank 返回 `0`：dynamic grouped-expert 参数最大差 `1.53e-5`，m/v 最大差 `4.88e-5` / `5.75e-8`。ABI13 dense base-MLP TP2 smoke 对 gate/up/down 使用半尺寸本地权重，eval/train loss 与完整权重参考相差 `1.84e-5`。 + +ABI14 的 4Q/2KV-head GQA full-attention TP2 smoke 覆盖 Q/K/V/O fixed 和 selected dynamic LoRA。fixed eval/loss 与完整参考最大差 `4.40e-4`；Q/K/V-B 与 O-A 梯度最大差分别为 `1.83e-4`、`1.14e-5`、`3.05e-4` 和 `1.83e-4`，FP32 Adam m/v 最大差 `1.53e-5` / `6.63e-9`，本地标准 Adam 公式误差小于 `3.73e-9`。selected dynamic loss 最大差 `6.82e-5`，Q/K/V/O 参数最大差 `3.05e-5`。另一个两层 GDN smoke 直接验证 latent-rank TP 的 input-dgrad backward all-reduce：A/B 梯度最大差 `9.54e-7` / `4.77e-7`，m/v 最大差 `4.77e-8` / `3.71e-14`，本地 Adam 误差为 `0`。BF16 参数直接对照最多跨约两个量化 bin，因此验收以 FP32 梯度、m/v 和本地 Adam 公式为主。没有完成完整大模型长时间训练、跨节点通信、GDN/MoE/vocab base TP、PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖且有直接 oracle 的子集，而不是所有并行配置。 ## 继续达到 Megatron 级别所需的最小工作包 -1. 建立 5D TP/PP/DP/EP/CP topology，并让 launcher、NCCL process groups 和 checkpoint 使用同一 rank 映射。 -2. 补齐 Qwen full/GDN attention、MoE、LM-head/CE 的 TP shard，并为 dense MLP 增加 projection-aware LoRA、fused gate/up FC1 和 MTP 支持；为 PP 实现 stage forward/backward 与 1F1B scheduler；为 CP 实现 ring attention/state exchange。 +1. 把现有 5D TP/PP/DP/EP/CP topology contract 落成可组合 runtime process groups、launcher 和 checkpoint rank mapping；当前 native 仍拒绝多轴组合。 +2. 补齐 Qwen GDN attention、MoE、embedding/LM-head/CE 的 TP shard，并为 dense MLP 增加 projection-aware LoRA、fused gate/up FC1 和 MTP 支持；为 full attention 融合 QKV/SDPA 并增加 sequence parallel；为 PP 实现 stage forward/backward 与 1F1B scheduler；为 CP 实现 ring attention/state exchange。 3. 将 EP dispatch/combine 替换为 fused/异步路径，并测量通信与计算重叠。 -4. 为 LoRA 增加 FP32 accumulation、每 adapter optimizer step、可恢复的 accumulation 状态和 rank-sharded checkpoint。 +4. 为 checkpoint 增加跨 topology reshard、可恢复的 pending accumulation state，并为旧 v3 attention checkpoint 提供离线迁移工具。 5. 在固定硬件和 workload 上，与 Megatron-LM 记录 tokens/s、step time、峰值显存、通信占比和 loss 曲线。 ## Native EP Benchmark Artifact From 69c4153b3d6ee2cc32ea0820bd07bb13b64de7bc Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 14:06:09 +0800 Subject: [PATCH 026/156] feat: add qwen gdn tensor parallelism --- .../kernels/qwen3_6_kernels.cpp | 113 ++- crates/rustrain-qwen3-6/src/kernel.rs | 202 ++++- crates/rustrain-qwen3-6/src/session.rs | 49 +- .../rustrain-qwen3-6/tests/native_smoke.cpp | 2 +- .../tests/native_tp_gdn_bench.cpp | 345 +++++++++ .../tests/native_tp_gdn_smoke.cpp | 707 ++++++++++++++++++ .../tests/native_tp_latent_smoke.cpp | 4 +- crates/rustrain-server/src/checkpoint.rs | 237 +++++- crates/rustrain-server/src/session.rs | 66 +- docs/agent/linear-attention.md | 31 + docs/qwen35-qwen36-megatron-audit.md | 16 +- scripts/run_qwen36_native_gdn_tp.sh | 192 +++++ 12 files changed, 1914 insertions(+), 50 deletions(-) create mode 100644 crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp create mode 100644 crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp create mode 100755 scripts/run_qwen36_native_gdn_tp.sh diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 51583d1d..756bfad6 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -1820,6 +1820,8 @@ static at::Tensor forward_single_layer( // Linear attention auto in_proj_qkv = *w[2], in_proj_z = *w[3], in_proj_a = *w[4], in_proj_b = *w[5]; auto a_log = *w[6], dt_bias = *w[7], conv1d_w = *w[8], norm_w = *w[9], out_proj = *w[10]; + TORCH_CHECK(!base_tp_attention_enabled(ctx) || use_batched, + "base linear-attention TP requires the activation-level LoRA path"); if (use_batched) { attn_output = linear_attention_batched( ctx, attn_input, layer_idx, @@ -2020,9 +2022,8 @@ struct TrainingContext { // input-sharded by the Rust weight loader. The local row contribution // is reduced over the TP communicator before the residual add. bool base_tp_mlp = false; - // Frozen full-attention TP: Q/K/V own disjoint head bundles and O owns - // the matching input columns. GDN layers remain replicated until their - // state/head bundle partition has a dedicated implementation. + // Frozen attention TP: full attention and GDN own disjoint head bundles; + // their output projections own the matching input columns. bool base_tp_attention = false; // Set when a legacy NCCL setter supplies an incompatible mixed topology. // Training entry points reject the context before touching parameters. @@ -2231,14 +2232,16 @@ static LoraTpLayout lora_tp_layout( layer_idx < 0 || layer_idx >= ctx->num_layers) return LoraTpLayout::LatentRank; const auto& cfg = ctx->layer_configs[layer_idx]; - if (cfg.layer_type != 0) return LoraTpLayout::LatentRank; auto table = lora_projection_table(cfg); TORCH_CHECK(pair_idx >= 0 && pair_idx < table.count, "invalid LoRA pair for TP layout"); const std::string name(table.entries[pair_idx].name); - if (name == "q_proj" || name == "k_proj" || name == "v_proj") + if (name == "q_proj" || name == "k_proj" || name == "v_proj" || + name == "in_proj_qkv" || name == "in_proj_z" || + name == "in_proj_a" || name == "in_proj_b") return LoraTpLayout::ColumnParallel; - if (name == "o_proj") return LoraTpLayout::RowParallel; + if (name == "o_proj" || name == "out_proj") + return LoraTpLayout::RowParallel; return LoraTpLayout::LatentRank; } @@ -3038,8 +3041,8 @@ at::Tensor compute_attn_only( } // Legacy path: weight-level LoRA - TORCH_CHECK(!ctx->base_tp_attention || cfg.layer_type != 0, - "base full-attention TP requires the activation-level LoRA path"); + TORCH_CHECK(!ctx->base_tp_attention, + "base attention TP requires the activation-level LoRA path"); int64_t lora_count = lora_pair_count(cfg); int64_t la_offset = ctx->lora_layer_offset[layer_idx]; bool has_lora = (la_offset + lora_count) <= (int64_t)ctx->lora_a.size(); @@ -3408,17 +3411,26 @@ static at::Tensor linear_attention_batched( ) { // Full reimplementation of linear_attention non-chunked path with // activation-level LoRA delta on QKV, Z, and out_proj. - auto device = hidden.device(); int64_t batch = hidden.size(0), seq = hidden.size(1); + auto projection_input = tp_copy_base_attention_input(ctx, hidden); + if (base_tp_attention_enabled(ctx)) { + TORCH_CHECK(num_k_heads > 0 && num_v_heads > 0 && + num_v_heads % num_k_heads == 0 && + num_k_heads % ctx->tp_world_size == 0 && + num_v_heads % ctx->tp_world_size == 0, + "linear-attention heads must preserve value-head groups and be divisible by TP_SIZE"); + num_k_heads /= ctx->tp_world_size; + num_v_heads /= ctx->tp_world_size; + } int64_t q_size = num_k_heads * key_dim; int64_t v_size = num_v_heads * val_dim; int64_t qkv_dim = q_size * 2 + v_size; // QKV projection + LoRA delta - auto qkv = at::matmul(hidden, in_proj_qkv.t()); + auto qkv = at::matmul(projection_input, in_proj_qkv.t()); auto it_qkv = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 0)); if (it_qkv != ctx->lora_batch_cache.end()) { - qkv = qkv + lora_activation_delta(ctx, hidden, it_qkv->second.a_stack, + qkv = qkv + lora_activation_delta(ctx, projection_input, it_qkv->second.a_stack, it_qkv->second.b_stack, it_qkv->second.scaling, it_qkv->second.layout); } @@ -3487,24 +3499,24 @@ static at::Tensor linear_attention_batched( v_f[0][0][0][0].item(), v_f[0][0][0][1].item(), v_f[0][0][0][2].item()); } - auto a = at::matmul(hidden, in_proj_a.t()); - auto b = at::matmul(hidden, in_proj_b.t()); + auto a = at::matmul(projection_input, in_proj_a.t()); + auto b = at::matmul(projection_input, in_proj_b.t()); auto it_a = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 2)); if (it_a != ctx->lora_batch_cache.end()) { - a = a + lora_activation_delta(ctx, hidden, it_a->second.a_stack, + a = a + lora_activation_delta(ctx, projection_input, it_a->second.a_stack, it_a->second.b_stack, it_a->second.scaling, it_a->second.layout); } auto it_b = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 3)); if (it_b != ctx->lora_batch_cache.end()) { - b = b + lora_activation_delta(ctx, hidden, it_b->second.a_stack, + b = b + lora_activation_delta(ctx, projection_input, it_b->second.a_stack, it_b->second.b_stack, it_b->second.scaling, it_b->second.layout); } // Z projection + LoRA delta - auto z = at::matmul(hidden, in_proj_z.t()); + auto z = at::matmul(projection_input, in_proj_z.t()); auto it_z = ctx->lora_batch_cache.find(lora_cache_key(layer_idx, 1)); if (it_z != ctx->lora_batch_cache.end()) { - z = z + lora_activation_delta(ctx, hidden, it_z->second.a_stack, + z = z + lora_activation_delta(ctx, projection_input, it_z->second.a_stack, it_z->second.b_stack, it_z->second.scaling, it_z->second.layout); } z = z.reshape({batch, seq, num_v_heads, head_v_dim}); @@ -3621,7 +3633,7 @@ static at::Tensor linear_attention_batched( it_op->second.b_stack, it_op->second.scaling, it_op->second.layout); } - return result; + return tp_allreduce_base_attention(ctx, result); } // ────────────────────────────────────────────────────────────────────── @@ -4578,7 +4590,7 @@ static at::Tensor mtp_compute_loss( extern "C" { __attribute__((visibility("default"))) int64_t qwen36_kernel_abi_version() { - return 14; + return 15; } static constexpr int32_t QWEN36_CONTEXT_BASE_TP_ATTENTION = 1 << 0; @@ -4667,7 +4679,7 @@ static void* qwen36_create_training_context_impl( if (ctx->base_tp_attention) { TORCH_CHECK(ctx->tp_world_size > 1, - "base full-attention TP requires TP_SIZE>1"); + "base attention TP requires TP_SIZE>1"); int64_t weight_offset = 0; for (int64_t layer = 0; layer < num_layers; ++layer) { const auto& cfg = ctx->layer_configs[layer]; @@ -4706,6 +4718,67 @@ static void* qwen36_create_training_context_impl( "base full-attention TP received inconsistent local weight shapes at layer ", layer, ": q=", q->sizes(), " k=", k->sizes(), " v=", v->sizes(), " o=", o->sizes()); + } else { + TORCH_CHECK(cfg.num_k_heads > 0 && cfg.num_v_heads > 0 && + cfg.key_dim == 128 && cfg.val_dim == 128 && + cfg.conv_kernel > 0, + "invalid linear-attention configuration at layer ", layer); + TORCH_CHECK(cfg.num_v_heads % cfg.num_k_heads == 0 && + cfg.num_k_heads % ctx->tp_world_size == 0 && + cfg.num_v_heads % ctx->tp_world_size == 0, + "linear-attention heads must preserve value-head groups and be divisible by TP_SIZE at layer ", + layer); + auto* qkv = ctx->weight_ptrs[weight_offset + 2]; + auto* z = ctx->weight_ptrs[weight_offset + 3]; + auto* a = ctx->weight_ptrs[weight_offset + 4]; + auto* b = ctx->weight_ptrs[weight_offset + 5]; + auto* a_log = ctx->weight_ptrs[weight_offset + 6]; + auto* dt_bias = ctx->weight_ptrs[weight_offset + 7]; + auto* conv = ctx->weight_ptrs[weight_offset + 8]; + auto* norm = ctx->weight_ptrs[weight_offset + 9]; + auto* out = ctx->weight_ptrs[weight_offset + 10]; + TORCH_CHECK(qkv && z && a && b && a_log && dt_bias && + conv && norm && out, + "base linear-attention TP received null weights at layer ", layer); + const int64_t local_k_heads = + cfg.num_k_heads / ctx->tp_world_size; + const int64_t local_v_heads = + cfg.num_v_heads / ctx->tp_world_size; + const int64_t local_q = local_k_heads * cfg.key_dim; + const int64_t local_v = local_v_heads * cfg.val_dim; + const int64_t local_qkv = local_q * 2 + local_v; + TORCH_CHECK(qkv->dim() == 2 && + qkv->size(0) == local_qkv, + "base linear-attention TP QKV shape mismatch at layer ", + layer, ": ", qkv->sizes(), " expected rows=", local_qkv); + const int64_t hidden_size = qkv->size(1); + TORCH_CHECK(z->dim() == 2 && z->size(0) == local_v && + z->size(1) == hidden_size && + a->dim() == 2 && a->size(0) == local_v_heads && + a->size(1) == hidden_size && + b->dim() == 2 && b->sizes() == a->sizes(), + "base linear-attention TP Z/A/B shapes mismatch at layer ", layer, + ": z=", z->sizes(), " a=", a->sizes(), " b=", b->sizes()); + TORCH_CHECK(a_log->dim() == 1 && + a_log->size(0) == local_v_heads && + dt_bias->dim() == 1 && + dt_bias->size(0) == local_v_heads, + "base linear-attention TP A_log/dt_bias shapes mismatch at layer ", + layer, ": A_log=", a_log->sizes(), + " dt_bias=", dt_bias->sizes()); + TORCH_CHECK(conv->dim() == 3 && + conv->size(0) == local_qkv && conv->size(1) == 1 && + conv->size(2) == cfg.conv_kernel, + "base linear-attention TP depthwise-conv shape mismatch at layer ", + layer, ": ", conv->sizes()); + TORCH_CHECK(norm->dim() == 1 && norm->size(0) == cfg.val_dim, + "base linear-attention TP norm must remain replicated at layer ", + layer, ": ", norm->sizes()); + TORCH_CHECK(out->dim() == 2 && out->size(0) == hidden_size && + out->size(1) == local_v, + "base linear-attention TP output shape mismatch at layer ", + layer, ": ", out->sizes(), " expected [", hidden_size, + ", ", local_v, "]"); } weight_offset += weight_count_for_layer(cfg); } diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs index 9bf5a9e9..7c000d45 100644 --- a/crates/rustrain-qwen3-6/src/kernel.rs +++ b/crates/rustrain-qwen3-6/src/kernel.rs @@ -197,7 +197,7 @@ unsafe fn load_kernels() -> Option { }}; } let abi_version: FnKernelAbiVersion = sym!("qwen36_kernel_abi_version"); - if abi_version() != 14 { + if abi_version() != 15 { return None; } Some(KernelHandles { @@ -326,6 +326,97 @@ pub fn shard_full_attention_weight_for_tp( )) } +/// Return the rank-local frozen GDN shard for TP. Q/K/V use the model's flat +/// `[Q_all | K_all | V_all]` layout, so QKV and depthwise-conv tensors must be +/// sliced segment by segment before being packed into the local flat layout. +#[allow(clippy::too_many_arguments)] +pub fn shard_linear_attention_weight_for_tp( + name: &str, + tensor: &Tensor, + tp_size: usize, + tp_rank: usize, + num_k_heads: i64, + key_dim: i64, + num_v_heads: i64, + val_dim: i64, +) -> Result> { + let is_qkv = name.ends_with(".linear_attn.in_proj_qkv.weight"); + let is_conv = name.ends_with(".linear_attn.conv1d.weight"); + let is_z = name.ends_with(".linear_attn.in_proj_z.weight"); + let is_ab = name.ends_with(".linear_attn.in_proj_a.weight") + || name.ends_with(".linear_attn.in_proj_b.weight"); + let is_head = name.ends_with(".linear_attn.A_log") || name.ends_with(".linear_attn.dt_bias"); + let is_out = name.ends_with(".linear_attn.out_proj.weight"); + if !(is_qkv || is_conv || is_z || is_ab || is_head || is_out) { + return Ok(None); + } + if tp_size <= 1 || tp_rank >= tp_size { + bail!("invalid linear-attention TP shard: tp_rank={tp_rank}, tp_size={tp_size}"); + } + let tp_size_i64 = tp_size as i64; + if num_k_heads <= 0 + || num_v_heads <= 0 + || key_dim <= 0 + || val_dim <= 0 + || num_v_heads % num_k_heads != 0 + || num_k_heads % tp_size_i64 != 0 + || num_v_heads % tp_size_i64 != 0 + { + bail!( + "linear-attention heads/dimensions must preserve value-head groups and be divisible by TP_SIZE={tp_size}: k_heads={num_k_heads}, v_heads={num_v_heads}, key_dim={key_dim}, val_dim={val_dim}" + ); + } + + let q_total = num_k_heads * key_dim; + let v_total = num_v_heads * val_dim; + let qkv_total = q_total * 2 + v_total; + let local_q = q_total / tp_size_i64; + let local_v = v_total / tp_size_i64; + let local_heads = num_v_heads / tp_size_i64; + let rank = tp_rank as i64; + let shape = tensor.size(); + + let require_axis = |axis: usize, expected: i64| -> Result<()> { + let actual = shape.get(axis).copied().ok_or_else(|| { + anyhow::anyhow!("base TP linear-attention weight {name} has no dimension {axis}") + })?; + if actual != expected { + bail!( + "base TP linear-attention weight {name} dimension {axis}={actual}, expected {expected}" + ); + } + Ok(()) + }; + + if is_qkv || is_conv { + require_axis(0, qkv_total)?; + if is_qkv && shape.len() != 2 { + bail!("base TP linear-attention QKV weight {name} must be rank 2"); + } + if is_conv && (shape.len() != 3 || shape[1] != 1) { + bail!("base TP linear-attention depthwise conv weight {name} must have shape [channels, 1, kernel]"); + } + let q = tensor.narrow(0, rank * local_q, local_q); + let k = tensor.narrow(0, q_total + rank * local_q, local_q); + let v = tensor.narrow(0, q_total * 2 + rank * local_v, local_v); + return Ok(Some(Tensor::cat(&[&q, &k, &v], 0).contiguous())); + } + if is_z { + require_axis(0, v_total)?; + return Ok(Some(tensor.narrow(0, rank * local_v, local_v).contiguous())); + } + if is_ab || is_head { + require_axis(0, num_v_heads)?; + return Ok(Some( + tensor + .narrow(0, rank * local_heads, local_heads) + .contiguous(), + )); + } + require_axis(1, v_total)?; + Ok(Some(tensor.narrow(1, rank * local_v, local_v).contiguous())) +} + pub fn build_weight_ptrs( weights: &std::collections::BTreeMap, config: &crate::config::Qwen36RuntimeConfig, @@ -1191,7 +1282,10 @@ impl Drop for CppTrainingContext { #[cfg(test)] mod tests { - use super::{shard_dense_mlp_weight_for_tp, shard_full_attention_weight_for_tp}; + use super::{ + shard_dense_mlp_weight_for_tp, shard_full_attention_weight_for_tp, + shard_linear_attention_weight_for_tp, + }; use tch::{Kind, Tensor}; #[test] @@ -1266,4 +1360,108 @@ mod tests { ) .is_err()); } + + #[test] + fn linear_attention_tp_preserves_flat_qkv_and_conv_segments() { + let qkv = Tensor::arange(48, (Kind::Float, tch::Device::Cpu)).reshape([12, 4]); + let conv = Tensor::arange(24, (Kind::Float, tch::Device::Cpu)).reshape([12, 1, 2]); + let qkv_rank_one = shard_linear_attention_weight_for_tp( + "model.layers.0.linear_attn.in_proj_qkv.weight", + &qkv, + 2, + 1, + 2, + 2, + 4, + 1, + ) + .unwrap() + .unwrap(); + let conv_rank_one = shard_linear_attention_weight_for_tp( + "model.layers.0.linear_attn.conv1d.weight", + &conv, + 2, + 1, + 2, + 2, + 4, + 1, + ) + .unwrap() + .unwrap(); + assert_eq!(qkv_rank_one.size(), [6, 4]); + assert_eq!(conv_rank_one.size(), [6, 1, 2]); + // Global rows are Q=[0..4], K=[4..8], V=[8..12]. Rank 1 owns the + // second half of each segment and repacks them as local [Q|K|V]. + assert_eq!(qkv_rank_one.double_value(&[0, 0]), 8.0); + assert_eq!(qkv_rank_one.double_value(&[2, 0]), 24.0); + assert_eq!(qkv_rank_one.double_value(&[4, 0]), 40.0); + assert_eq!(conv_rank_one.double_value(&[0, 0, 0]), 4.0); + assert_eq!(conv_rank_one.double_value(&[2, 0, 0]), 12.0); + assert_eq!(conv_rank_one.double_value(&[4, 0, 0]), 20.0); + } + + #[test] + fn linear_attention_tp_shards_value_head_tensors_and_output_columns() { + let z = Tensor::arange(32, (Kind::Float, tch::Device::Cpu)).reshape([8, 4]); + let a = Tensor::arange(16, (Kind::Float, tch::Device::Cpu)).reshape([4, 4]); + let a_log = Tensor::arange(4, (Kind::Float, tch::Device::Cpu)); + let out = Tensor::arange(32, (Kind::Float, tch::Device::Cpu)).reshape([4, 8]); + let shard = |name: &str, tensor: &Tensor| { + shard_linear_attention_weight_for_tp(name, tensor, 2, 1, 2, 2, 4, 2) + .unwrap() + .unwrap() + }; + let z = shard("model.layers.0.linear_attn.in_proj_z.weight", &z); + let a = shard("model.layers.0.linear_attn.in_proj_a.weight", &a); + let a_log = shard("model.layers.0.linear_attn.A_log", &a_log); + let out = shard("model.layers.0.linear_attn.out_proj.weight", &out); + assert_eq!(z.size(), [4, 4]); + assert_eq!(a.size(), [2, 4]); + assert_eq!(a_log.size(), [2]); + assert_eq!(out.size(), [4, 4]); + assert_eq!(z.double_value(&[0, 0]), 16.0); + assert_eq!(a.double_value(&[0, 0]), 8.0); + assert_eq!(a_log.double_value(&[0]), 2.0); + assert_eq!(out.double_value(&[0, 0]), 4.0); + } + + #[test] + fn linear_attention_tp_rejects_invalid_groups_and_shapes() { + let tensor = Tensor::zeros([12, 4], (Kind::Float, tch::Device::Cpu)); + assert!(shard_linear_attention_weight_for_tp( + "model.layers.0.linear_attn.norm.weight", + &tensor, + 2, + 0, + 2, + 2, + 4, + 1, + ) + .unwrap() + .is_none()); + assert!(shard_linear_attention_weight_for_tp( + "model.layers.0.linear_attn.in_proj_qkv.weight", + &tensor, + 2, + 0, + 3, + 2, + 4, + 1, + ) + .is_err()); + assert!(shard_linear_attention_weight_for_tp( + "model.layers.0.linear_attn.in_proj_qkv.weight", + &Tensor::zeros([11, 4], (Kind::Float, tch::Device::Cpu)), + 2, + 0, + 2, + 2, + 4, + 1, + ) + .is_err()); + } } diff --git a/crates/rustrain-qwen3-6/src/session.rs b/crates/rustrain-qwen3-6/src/session.rs index 8a9767f5..50d725c9 100644 --- a/crates/rustrain-qwen3-6/src/session.rs +++ b/crates/rustrain-qwen3-6/src/session.rs @@ -328,9 +328,10 @@ fn train_impl( ); } - // Full attention follows Megatron's Q/K/V ColumnParallel and O - // RowParallel layout. Dense MLP additionally shards gate/up rows and down - // columns. GDN, MoE experts, embeddings, and the LM head remain replicated. + // Full attention and GDN follow head-aligned ColumnParallel input + // projections and RowParallel output projections. Dense MLP additionally + // shards gate/up rows and down columns. MoE experts, embeddings, and the + // LM head remain replicated. let base_tp_attention = tp_size > 1; let base_tp_mlp = tp_size > 1 && !runtime_config.is_moe; if base_tp_attention { @@ -362,6 +363,30 @@ fn train_impl( runtime_config.partial_rotary_factor ); } + if runtime_config + .layer_types + .iter() + .any(|layer| *layer == LayerType::LinearAttention) + { + if runtime_config.linear_num_key_heads <= 0 + || runtime_config.linear_num_value_heads <= 0 + || runtime_config.linear_num_value_heads % runtime_config.linear_num_key_heads != 0 + || runtime_config.linear_num_key_heads % tp_size as i64 != 0 + || runtime_config.linear_num_value_heads % tp_size as i64 != 0 + || runtime_config.linear_key_head_dim != 128 + || runtime_config.linear_value_head_dim != 128 + || runtime_config.linear_conv_kernel_dim <= 0 + { + bail!( + "linear-attention TP requires k/v heads divisible by TP_SIZE with preserved value-head groups, 128-wide key/value heads, and a positive conv kernel: k_heads={}, v_heads={}, key_dim={}, value_dim={}, conv_kernel={}, tp={tp_size}", + runtime_config.linear_num_key_heads, + runtime_config.linear_num_value_heads, + runtime_config.linear_key_head_dim, + runtime_config.linear_value_head_dim, + runtime_config.linear_conv_kernel_dim, + ); + } + } } if base_tp_mlp { if runtime_config.intermediate_size <= 0 @@ -440,12 +465,26 @@ fn train_impl( } else { for (name, tensor) in &weights { let local_shard = if base_tp_attention { - let attention_shard = crate::kernel::shard_full_attention_weight_for_tp( + let full_attention_shard = crate::kernel::shard_full_attention_weight_for_tp( name, tensor, tp_size, rank % tp_size, )?; + let attention_shard = if full_attention_shard.is_some() { + full_attention_shard + } else { + crate::kernel::shard_linear_attention_weight_for_tp( + name, + tensor, + tp_size, + rank % tp_size, + runtime_config.linear_num_key_heads, + runtime_config.linear_key_head_dim, + runtime_config.linear_num_value_heads, + runtime_config.linear_value_head_dim, + )? + }; if attention_shard.is_some() || !base_tp_mlp { attention_shard } else { @@ -471,7 +510,7 @@ fn train_impl( tp_size, tp_rank = rank % tp_size, base_tp_mlp, - "frozen base TP enabled: full-attention head shards and optional dense MLP shards" + "frozen base TP enabled: full-attention/GDN head shards and optional dense MLP shards" ); } } diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index ed57a659..36114967 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -366,7 +366,7 @@ static int run_dynamic_dp_smoke( } int main() { - assert(qwen36_kernel_abi_version() == 14); + assert(qwen36_kernel_abi_version() == 15); const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); const int process_rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0"); const int local_rank = std::atoi(std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); diff --git a/crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp b/crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp new file mode 100644 index 00000000..5d03651e --- /dev/null +++ b/crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp @@ -0,0 +1,345 @@ +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct LayerConfig { + int64_t layer_type, num_heads, num_kv_heads, head_dim; + int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel; + double partial_rotary_factor, rope_theta, rms_eps; + int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count; + int64_t intermediate_size; + int32_t norm_topk_prob; + void* nccl_comm; + void* nccl_stream; +}; + +extern "C" int64_t qwen36_kernel_abi_version(); +extern "C" void qwen36_set_cuda_device(int32_t); +extern "C" void* qwen36_create_training_context( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*); +extern "C" void* qwen36_create_training_context_ex( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*, int32_t); +extern "C" int32_t qwen36_init_nccl(void*); +extern "C" double qwen36_train_step(void*, void*, void*, void*); +extern "C" void qwen36_free_training_context(void*); + +namespace { + +constexpr int64_t kAbiVersion = 15; +constexpr int32_t kBaseTpAttention = 1 << 0; + +static int env_int(const char* name, int fallback) { + const char* value = std::getenv(name); + if (!value || value[0] == '\0') return fallback; + const int parsed = std::atoi(value); + assert(parsed > 0); + return parsed; +} + +static int env_int_or(const char* primary, const char* secondary, int fallback) { + return std::getenv(primary) ? env_int(primary, fallback) + : env_int(secondary, fallback); +} + +static at::Tensor seeded_randn( + std::initializer_list shape, double scale, int64_t seed +) { + at::manual_seed(seed); + return (at::randn(shape, + at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)) * scale) + .to(at::kBFloat16); +} + +static at::Tensor unit_weight(std::initializer_list shape) { + return at::ones( + shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)); +} + +static std::vector pointers(std::vector& tensors) { + std::vector result; + result.reserve(tensors.size()); + for (auto& tensor : tensors) result.push_back(&tensor); + return result; +} + +static void append_gdn_layer( + std::vector& weights, int64_t layer, + int64_t hidden, int64_t intermediate, + int64_t global_k_heads, int64_t global_v_heads, + int64_t key_dim, int64_t value_dim, int64_t conv_kernel, + int tp_world +) { + const int64_t local_k_heads = global_k_heads / tp_world; + const int64_t local_v_heads = global_v_heads / tp_world; + const int64_t q_size = local_k_heads * key_dim; + const int64_t v_size = local_v_heads * value_dim; + const int64_t qkv_size = 2 * q_size + v_size; + const int64_t seed = 1000 + layer * 100; + + weights.push_back(unit_weight({hidden})); + weights.push_back(unit_weight({hidden})); + weights.push_back(seeded_randn({qkv_size, hidden}, 0.0020, seed + 2)); + weights.push_back(seeded_randn({v_size, hidden}, 0.0020, seed + 3)); + weights.push_back(seeded_randn({local_v_heads, hidden}, 0.0020, seed + 4)); + weights.push_back(seeded_randn({local_v_heads, hidden}, 0.0020, seed + 5)); + // A realistic negative time-step bias keeps the recurrent decay near one. + // A zero bias would produce decay ~= 0.5 and make reverse state recovery + // exponentially ill-conditioned over this synthetic long sequence. + weights.push_back(at::zeros({local_v_heads}, + at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + weights.push_back(at::full({local_v_heads}, -4.0, + at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + weights.push_back(seeded_randn( + {qkv_size, 1, conv_kernel}, 0.0020, seed + 8)); + weights.push_back(unit_weight({value_dim})); + weights.push_back(seeded_randn({hidden, v_size}, 0.0020, seed + 10)); + weights.push_back(seeded_randn( + {intermediate, hidden}, 0.0020, seed + 11)); + weights.push_back(seeded_randn( + {intermediate, hidden}, 0.0020, seed + 12)); + weights.push_back(seeded_randn( + {hidden, intermediate}, 0.0020, seed + 13)); +} + +static double percentile(std::vector values, double quantile) { + assert(!values.empty()); + std::sort(values.begin(), values.end()); + const double position = quantile * static_cast(values.size() - 1); + const size_t lower = static_cast(position); + const size_t upper = std::min(lower + 1, values.size() - 1); + const double fraction = position - static_cast(lower); + return values[lower] * (1.0 - fraction) + values[upper] * fraction; +} + +static double gib(size_t bytes) { + return static_cast(bytes) / (1024.0 * 1024.0 * 1024.0); +} + +static size_t used_since(size_t initial_free, size_t observed_free) { + return initial_free > observed_free ? initial_free - observed_free : 0; +} + +} // namespace + +int main() { + const char* mode_env = std::getenv("BENCH_MODE"); + const std::string mode = mode_env ? mode_env : "single"; + assert(mode == "single" || mode == "tp2"); + const bool use_tp = mode == "tp2"; + const int expected_world = use_tp ? 2 : 1; + const int rank = std::getenv("RANK") ? std::atoi(std::getenv("RANK")) : 0; + const int world = std::getenv("WORLD_SIZE") + ? std::atoi(std::getenv("WORLD_SIZE")) : 1; + const int local_rank = std::getenv("LOCAL_RANK") + ? std::atoi(std::getenv("LOCAL_RANK")) : rank; + assert(world == expected_world && rank >= 0 && rank < world); + assert(qwen36_kernel_abi_version() == kAbiVersion); + qwen36_set_cuda_device(local_rank); + assert(cudaFree(nullptr) == cudaSuccess); + + const int batch = env_int("BENCH_BATCH", 2); + const int seq = env_int("BENCH_SEQ", 512); + const int hidden = env_int("BENCH_HIDDEN", 2048); + const int key_heads = env_int("BENCH_K_HEADS", 16); + const int value_heads = env_int("BENCH_V_HEADS", 32); + const int key_dim = env_int_or("BENCH_KEY_DIM", "BENCH_HEAD_DIM", 128); + const int value_dim = env_int_or("BENCH_VALUE_DIM", "BENCH_HEAD_DIM", 128); + const int conv_kernel = env_int("BENCH_CONV", 4); + const int layers = env_int("BENCH_LAYERS", 3); + const int intermediate = env_int("BENCH_INTERMEDIATE", 2048); + const int vocab = env_int("BENCH_VOCAB", 4096); + const int lora_rank = env_int("BENCH_LORA_RANK", 8); + const int warmup = env_int("BENCH_WARMUP", 5); + const int iters = env_int("BENCH_ITERS", 30); + assert(seq >= 2 && hidden > 0 && key_dim > 0 && value_dim > 0); + assert(value_heads % key_heads == 0); + assert(key_heads % expected_world == 0 && value_heads % expected_world == 0); + assert(lora_rank % expected_world == 0); + + size_t free_start = 0, total_bytes = 0; + assert(cudaMemGetInfo(&free_start, &total_bytes) == cudaSuccess); + size_t min_observed_free = free_start; + + std::vector weights; + weights.reserve(layers * 14); + for (int64_t layer = 0; layer < layers; ++layer) { + append_gdn_layer(weights, layer, hidden, intermediate, + key_heads, value_heads, key_dim, value_dim, conv_kernel, + expected_world); + } + for (auto& weight : weights) weight.set_requires_grad(false); + auto weight_ptrs = pointers(weights); + + auto embed = seeded_randn({vocab, hidden}, 0.0020, 31); + auto final_norm = unit_weight({hidden}); + auto lm_head = seeded_randn({vocab, hidden}, 0.0020, 37); + embed.set_requires_grad(false); + final_norm.set_requires_grad(false); + lm_head.set_requires_grad(false); + + std::vector configs(layers); + for (auto& config : configs) { + config.layer_type = 1; + config.num_k_heads = key_heads; + config.key_dim = key_dim; + config.num_v_heads = value_heads; + config.val_dim = value_dim; + config.conv_kernel = conv_kernel; + config.rms_eps = 1e-5; + config.intermediate_size = intermediate; + } + std::vector target_layers(layers); + std::iota(target_layers.begin(), target_layers.end(), 0); + constexpr const char* targets = + "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj"; + + std::vector host_ids(batch * seq); + for (int b = 0; b < batch; ++b) { + for (int s = 0; s < seq; ++s) + host_ids[b * seq + s] = 1 + (b * 131 + s * 17) % (vocab - 1); + } + auto input_ids = at::from_blob(host_ids.data(), {batch, seq}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)).clone().to(at::kCUDA); + auto target_mask = at::ones({batch, seq}, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); + auto attention_mask = at::ones({batch, seq}, + at::TensorOptions().device(at::kCUDA).dtype(at::kBool)); + + size_t free_before_context = 0; + assert(cudaMemGetInfo(&free_before_context, &total_bytes) == cudaSuccess); + min_observed_free = std::min(min_observed_free, free_before_context); + setenv("TP_SIZE", use_tp ? "2" : "1", 1); + unsetenv("RUSTRAIN_DATA_PARALLEL"); + void* context = use_tp + ? qwen36_create_training_context_ex( + weight_ptrs.data(), weight_ptrs.size(), &embed, &final_norm, &lm_head, + configs.data(), layers, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + target_layers.data(), layers, targets, kBaseTpAttention) + : qwen36_create_training_context( + weight_ptrs.data(), weight_ptrs.size(), &embed, &final_norm, &lm_head, + configs.data(), layers, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, + target_layers.data(), layers, targets); + assert(context); + if (use_tp) assert(qwen36_init_nccl(context) == 0); + + size_t free_after_context = 0; + assert(cudaMemGetInfo(&free_after_context, &total_bytes) == cudaSuccess); + min_observed_free = std::min(min_observed_free, free_after_context); + double last_loss = 0.0; + for (int i = 0; i < warmup; ++i) { + last_loss = qwen36_train_step( + context, &input_ids, &target_mask, &attention_mask); + assert(std::isfinite(last_loss)); + } + assert(cudaDeviceSynchronize() == cudaSuccess); + + size_t free_after_warmup = 0; + assert(cudaMemGetInfo(&free_after_warmup, &total_bytes) == cudaSuccess); + min_observed_free = std::min(min_observed_free, free_after_warmup); + + cudaEvent_t start = nullptr; + cudaEvent_t stop = nullptr; + assert(cudaEventCreate(&start) == cudaSuccess); + assert(cudaEventCreate(&stop) == cudaSuccess); + std::vector times; + times.reserve(iters); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + for (int i = 0; i < iters; ++i) { + assert(cudaEventRecord(start, stream) == cudaSuccess); + last_loss = qwen36_train_step( + context, &input_ids, &target_mask, &attention_mask); + assert(cudaEventRecord(stop, stream) == cudaSuccess); + assert(cudaEventSynchronize(stop) == cudaSuccess); + assert(std::isfinite(last_loss)); + float elapsed_ms = 0.0f; + assert(cudaEventElapsedTime(&elapsed_ms, start, stop) == cudaSuccess); + times.push_back(elapsed_ms); + size_t free_now = 0; + assert(cudaMemGetInfo(&free_now, &total_bytes) == cudaSuccess); + min_observed_free = std::min(min_observed_free, free_now); + } + assert(cudaEventDestroy(start) == cudaSuccess); + assert(cudaEventDestroy(stop) == cudaSuccess); + + const double mean = std::accumulate(times.begin(), times.end(), 0.0) / + static_cast(times.size()); + double variance = 0.0; + for (const double value : times) variance += (value - mean) * (value - mean); + variance /= static_cast(times.size()); + const double p50 = percentile(times, 0.50); + const double p90 = percentile(times, 0.90); + const double model_tokens = static_cast(batch) * (seq - 1); + const double layer_tokens = model_tokens * layers; + const double model_tokens_per_second = model_tokens / (p50 / 1000.0); + const double layer_tokens_per_second = layer_tokens / (p50 / 1000.0); + + cudaDeviceProp properties{}; + assert(cudaGetDeviceProperties(&properties, local_rank) == cudaSuccess); + std::ostringstream output; + output << std::fixed << std::setprecision(6) + << "native_tp_gdn_bench {" + << "\"mode\":\"" << mode << "\"," + << "\"rank\":" << rank << ",\"world\":" << world << "," + << "\"gpu\":\"" << properties.name << "\"," + << "\"abi\":" << kAbiVersion << "," + << "\"batch\":" << batch << ",\"seq\":" << seq << "," + << "\"hidden\":" << hidden << ",\"layers\":" << layers << "," + << "\"key_heads\":" << key_heads << "," + << "\"value_heads\":" << value_heads << "," + << "\"key_dim\":" << key_dim << "," + << "\"value_dim\":" << value_dim << "," + << "\"n_rep\":" << (value_heads / key_heads) << "," + << "\"conv\":" << conv_kernel << "," + << "\"intermediate\":" << intermediate << "," + << "\"vocab\":" << vocab << ",\"lora_rank\":" << lora_rank << "," + << "\"warmup\":" << warmup << ",\"iters\":" << iters << "," + << "\"last_loss\":" << last_loss << "," + << "\"step_ms_mean\":" << mean << "," + << "\"step_ms_p50\":" << p50 << "," + << "\"step_ms_p90\":" << p90 << "," + << "\"step_ms_std\":" << std::sqrt(variance) << "," + << "\"model_tokens\":" << model_tokens << "," + << "\"gdn_layer_tokens\":" << layer_tokens << "," + << "\"model_tokens_per_sec\":" << model_tokens_per_second << "," + << "\"gdn_layer_tokens_per_sec\":" << layer_tokens_per_second << "," + << "\"device_total_gib\":" << gib(total_bytes) << "," + << "\"free_start_gib\":" << gib(free_start) << "," + << "\"free_before_context_gib\":" << gib(free_before_context) << "," + << "\"free_after_context_gib\":" << gib(free_after_context) << "," + << "\"free_after_warmup_gib\":" << gib(free_after_warmup) << "," + << "\"context_resident_delta_gib\":" + << gib(used_since(free_before_context, free_after_context)) << "," + << "\"max_observed_resident_gib\":" + << gib(used_since(free_start, min_observed_free)) << "," + << "\"samples_ms\":["; + for (size_t i = 0; i < times.size(); ++i) { + if (i) output << ','; + output << times[i]; + } + output << "]}\n"; + const std::string line = output.str(); + std::fwrite(line.data(), 1, line.size(), stdout); + std::fflush(stdout); + + qwen36_free_training_context(context); + return 0; +} diff --git a/crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp new file mode 100644 index 00000000..aabd341b --- /dev/null +++ b/crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp @@ -0,0 +1,707 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct LayerConfig { + int64_t layer_type, num_heads, num_kv_heads, head_dim; + int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel; + double partial_rotary_factor, rope_theta, rms_eps; + int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count; + int64_t intermediate_size; + int32_t norm_topk_prob; + void* nccl_comm; + void* nccl_stream; +}; + +extern "C" int64_t qwen36_kernel_abi_version(); +extern "C" void qwen36_set_cuda_device(int32_t); +extern "C" void* qwen36_create_training_context( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*); +extern "C" void* qwen36_create_training_context_ex( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*, int32_t); +extern "C" int32_t qwen36_init_nccl(void*); +extern "C" int64_t qwen36_add_lora( + void*, int64_t, double, const int64_t*, int64_t, const char*); +extern "C" void* qwen36_get_adapter_lora_tensor( + void*, int64_t, int64_t, const char*, int32_t); +extern "C" int32_t qwen36_set_adapter_lora_tensor( + void*, int64_t, int64_t, const char*, int32_t, void*); +extern "C" void* qwen36_get_adapter_optimizer_tensor( + void*, int64_t, int64_t, const char*, int32_t, int32_t); +extern "C" void* qwen36_get_lora_a(void*, int64_t); +extern "C" void* qwen36_get_lora_b(void*, int64_t); +extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*); +extern "C" void* qwen36_get_lora_grad_accumulator(void*, int64_t, int32_t); +extern "C" int32_t qwen36_abort_gradient_accumulation(void*); +extern "C" int64_t qwen36_export_optimizer_state( + void*, void**, void**, int64_t); +extern "C" double qwen36_eval_step(void*, void*, void*, void*); +extern "C" double qwen36_train_step(void*, void*, void*, void*); +extern "C" double qwen36_train_micro_step( + void*, void*, void*, void*, double, int32_t); +extern "C" double qwen36_train_multi_lora_selected( + void*, void*, void*, void*, const int64_t*, int32_t, int32_t); +extern "C" void qwen36_free_training_context(void*); + +namespace { + +constexpr int64_t kAbiVersion = 15; +constexpr int32_t kBaseTpAttention = 1 << 0; +constexpr int64_t kLayers = 2; +constexpr int64_t kHidden = 32; +constexpr int64_t kIntermediate = 48; +constexpr int64_t kKeyHeads = 4; +constexpr int64_t kValueHeads = 8; +constexpr int64_t kKeyDim = 128; +constexpr int64_t kValueDim = 128; +constexpr int64_t kConvKernel = 4; +constexpr int64_t kVocab = 64; +constexpr int64_t kLoraRank = 4; +constexpr int64_t kPairsPerLayer = 8; +constexpr int64_t kQSize = kKeyHeads * kKeyDim; +constexpr int64_t kVSize = kValueHeads * kValueDim; +constexpr int64_t kQkvSize = 2 * kQSize + kVSize; + +constexpr std::array kGdnModules = { + "in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b", "out_proj"}; + +static int required_env_int(const char* name) { + const char* value = std::getenv(name); + assert(value && value[0] != '\0'); + return std::atoi(value); +} + +static at::Tensor fingerprint( + std::initializer_list shape, double scale, int64_t offset +) { + int64_t count = 1; + for (const int64_t dim : shape) count *= dim; + auto values = at::arange( + count, at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); + return ((values.add(offset).remainder(97) - 48.0) * scale) + .reshape(shape).to(at::kBFloat16); +} + +static at::Tensor unit_weight(std::initializer_list shape) { + return at::ones( + shape, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)); +} + +static void append_full_gdn_layer( + std::vector& weights, int64_t layer +) { + const int64_t base = 1000 * layer; + weights.push_back(unit_weight({kHidden})); + weights.push_back(unit_weight({kHidden})); + + auto q = fingerprint({kQSize, kHidden}, 0.00035, base + 11); + auto k = fingerprint({kQSize, kHidden}, 0.00041, base + 211); + auto v = fingerprint({kVSize, kHidden}, 0.00029, base + 421); + weights.push_back(at::cat({q, k, v}, 0).contiguous()); + weights.push_back(fingerprint({kVSize, kHidden}, 0.00031, base + 631)); + weights.push_back(fingerprint({kValueHeads, kHidden}, 0.00043, base + 719)); + weights.push_back(fingerprint({kValueHeads, kHidden}, 0.00047, base + 811)); + weights.push_back(fingerprint({kValueHeads}, 0.0008, base + 907)); + weights.push_back(fingerprint({kValueHeads}, 0.0007, base + 953)); + + auto q_conv = fingerprint( + {kQSize, 1, kConvKernel}, 0.0009, base + 101); + auto k_conv = fingerprint( + {kQSize, 1, kConvKernel}, 0.0011, base + 307); + auto v_conv = fingerprint( + {kVSize, 1, kConvKernel}, 0.0007, base + 509); + weights.push_back(at::cat({q_conv, k_conv, v_conv}, 0).contiguous()); + weights.push_back(unit_weight({kValueDim})); + weights.push_back(fingerprint({kHidden, kVSize}, 0.00033, base + 613)); + + weights.push_back(fingerprint( + {kIntermediate, kHidden}, 0.00045, base + 701)); + weights.push_back(fingerprint( + {kIntermediate, kHidden}, 0.00039, base + 797)); + weights.push_back(fingerprint( + {kHidden, kIntermediate}, 0.00037, base + 887)); +} + +static at::Tensor shard_flat_qkv(const at::Tensor& full, int rank, int dim) { + const int64_t local_q = kQSize / 2; + const int64_t local_v = kVSize / 2; + auto q = full.narrow(dim, rank * local_q, local_q); + auto k = full.narrow(dim, kQSize + rank * local_q, local_q); + auto v = full.narrow(dim, 2 * kQSize + rank * local_v, local_v); + return at::cat({q, k, v}, dim).contiguous(); +} + +static std::vector make_local_weights( + const std::vector& full, int rank +) { + std::vector local; + local.reserve(full.size()); + for (int64_t layer = 0; layer < kLayers; ++layer) { + const int64_t offset = layer * 14; + local.push_back(full[offset + 0]); + local.push_back(full[offset + 1]); + local.push_back(shard_flat_qkv(full[offset + 2], rank, 0)); + local.push_back(full[offset + 3] + .narrow(0, rank * (kVSize / 2), kVSize / 2).contiguous()); + local.push_back(full[offset + 4] + .narrow(0, rank * (kValueHeads / 2), kValueHeads / 2).contiguous()); + local.push_back(full[offset + 5] + .narrow(0, rank * (kValueHeads / 2), kValueHeads / 2).contiguous()); + local.push_back(full[offset + 6] + .narrow(0, rank * (kValueHeads / 2), kValueHeads / 2).contiguous()); + local.push_back(full[offset + 7] + .narrow(0, rank * (kValueHeads / 2), kValueHeads / 2).contiguous()); + local.push_back(shard_flat_qkv(full[offset + 8], rank, 0)); + local.push_back(full[offset + 9]); + local.push_back(full[offset + 10] + .narrow(1, rank * (kVSize / 2), kVSize / 2).contiguous()); + local.insert(local.end(), full.begin() + offset + 11, + full.begin() + offset + 14); + } + return local; +} + +static std::vector pointers(std::vector& tensors) { + std::vector result; + result.reserve(tensors.size()); + for (auto& tensor : tensors) result.push_back(&tensor); + return result; +} + +static LayerConfig gdn_config() { + LayerConfig config{}; + config.layer_type = 1; + config.num_k_heads = kKeyHeads; + config.key_dim = kKeyDim; + config.num_v_heads = kValueHeads; + config.val_dim = kValueDim; + config.conv_kernel = kConvKernel; + config.rms_eps = 1e-5; + config.intermediate_size = kIntermediate; + return config; +} + +struct Batch { + at::Tensor input_ids; + at::Tensor target_mask; + at::Tensor attention_mask; +}; + +static Batch make_batch(int64_t batch, int64_t seq, int64_t offset) { + std::vector host(batch * seq); + for (int64_t b = 0; b < batch; ++b) { + for (int64_t s = 0; s < seq; ++s) { + host[b * seq + s] = 1 + (offset + b * 17 + s * 5) % (kVocab - 1); + } + } + auto ids = at::from_blob(host.data(), {batch, seq}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)).clone().to(at::kCUDA); + auto target = at::ones({batch, seq}, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)); + auto attention = at::ones({batch, seq}, + at::TensorOptions().device(at::kCUDA).dtype(at::kBool)); + return {std::move(ids), std::move(target), std::move(attention)}; +} + +static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) { + assert(lhs.sizes() == rhs.sizes()); + return (lhs.to(at::kFloat) - rhs.to(at::kFloat)) + .abs().max().item(); +} + +static double relative_l2(const at::Tensor& lhs, const at::Tensor& rhs) { + assert(lhs.sizes() == rhs.sizes()); + auto delta = lhs.to(at::kFloat) - rhs.to(at::kFloat); + const double denom = std::max(rhs.to(at::kFloat).norm().item(), 1e-12); + return delta.norm().item() / denom; +} + +static bool is_column_parallel(const char* module) { + return std::strcmp(module, "out_proj") != 0; +} + +static int64_t full_output_size(const char* module) { + if (std::strcmp(module, "in_proj_qkv") == 0) return kQkvSize; + if (std::strcmp(module, "in_proj_z") == 0) return kVSize; + if (std::strcmp(module, "in_proj_a") == 0 || + std::strcmp(module, "in_proj_b") == 0) return kValueHeads; + assert(std::strcmp(module, "out_proj") == 0); + return kHidden; +} + +static int64_t full_input_size(const char* module) { + return std::strcmp(module, "out_proj") == 0 ? kVSize : kHidden; +} + +static at::Tensor shard_projection_output( + const at::Tensor& full, const char* module, int rank, int dim +) { + if (std::strcmp(module, "in_proj_qkv") == 0) + return shard_flat_qkv(full, rank, dim); + const int64_t local = full.size(dim) / 2; + return full.narrow(dim, rank * local, local).contiguous(); +} + +static at::Tensor expected_local_factor( + const at::Tensor& full, const char* module, bool is_b, int rank +) { + if (is_column_parallel(module)) { + return is_b ? shard_projection_output(full, module, rank, 0) + : full; + } + return is_b ? full + : full.narrow(1, rank * (kVSize / 2), kVSize / 2).contiguous(); +} + +struct LoraFixture { + int64_t layer; + int64_t pair; + const char* module; + at::Tensor full_a; + at::Tensor full_b; + at::Tensor local_a; + at::Tensor local_b; +}; + +static std::vector make_lora_fixtures(int rank, int64_t seed_base) { + std::vector fixtures; + fixtures.reserve(kLayers * kGdnModules.size()); + for (int64_t layer = 0; layer < kLayers; ++layer) { + for (int64_t pair = 0; pair < static_cast(kGdnModules.size()); ++pair) { + const char* module = kGdnModules[pair]; + const int64_t base = seed_base + layer * 100 + pair * 13; + auto full_a = fingerprint( + {kLoraRank, full_input_size(module)}, 0.0007, base + 1); + auto full_b = fingerprint( + {full_output_size(module), kLoraRank}, 0.0006, base + 7); + auto local_a = expected_local_factor(full_a, module, false, rank); + auto local_b = expected_local_factor(full_b, module, true, rank); + fixtures.push_back({layer, pair, module, std::move(full_a), + std::move(full_b), std::move(local_a), std::move(local_b)}); + } + } + return fixtures; +} + +static void assert_weight_contract( + const std::vector& full, + const std::vector& local, int rank +) { + assert(full.size() == 28 && local.size() == full.size()); + for (int64_t layer = 0; layer < kLayers; ++layer) { + const int64_t offset = layer * 14; + assert(local[offset + 2].sizes() == + at::IntArrayRef({kQkvSize / 2, kHidden})); + assert(local[offset + 8].sizes() == + at::IntArrayRef({kQkvSize / 2, 1, kConvKernel})); + assert(local[offset + 3].sizes() == + at::IntArrayRef({kVSize / 2, kHidden})); + assert(local[offset + 4].sizes() == + at::IntArrayRef({kValueHeads / 2, kHidden})); + assert(local[offset + 9].sizes() == at::IntArrayRef({kValueDim})); + assert(local[offset + 10].sizes() == + at::IntArrayRef({kHidden, kVSize / 2})); + assert(max_diff(local[offset + 2], + shard_flat_qkv(full[offset + 2], rank, 0)) == 0.0); + assert(max_diff(local[offset + 8], + shard_flat_qkv(full[offset + 8], rank, 0)) == 0.0); + } +} + +struct ContextPair { + void* distributed = nullptr; + void* reference = nullptr; +}; + +static ContextPair create_context_pair( + std::vector& local_weights, + std::vector& full_weights, + at::Tensor& embed, at::Tensor& final_norm, at::Tensor& lm_head, + LayerConfig* configs, const char* fixed_targets, + bool run_negative_guard +) { + const int64_t target_layers[kLayers] = {0, 1}; + auto local_ptrs = pointers(local_weights); + setenv("TP_SIZE", "2", 1); + unsetenv("RUSTRAIN_DATA_PARALLEL"); + + if (run_negative_guard) { + auto invalid_weights = local_weights; + invalid_weights[2] = invalid_weights[2] + .narrow(0, 0, invalid_weights[2].size(0) - 1).contiguous(); + auto invalid_ptrs = pointers(invalid_weights); + void* invalid = qwen36_create_training_context_ex( + invalid_ptrs.data(), invalid_ptrs.size(), &embed, &final_norm, &lm_head, + configs, kLayers, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, kVocab, 1e-5, kLoraRank, + target_layers, kLayers, fixed_targets, kBaseTpAttention); + assert(!invalid && "invalid local flat-QKV shape must be rejected"); + } + + void* distributed = qwen36_create_training_context_ex( + local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head, + configs, kLayers, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, kVocab, 1e-5, kLoraRank, + target_layers, kLayers, fixed_targets, kBaseTpAttention); + assert(distributed && qwen36_init_nccl(distributed) == 0); + + setenv("TP_SIZE", "1", 1); + auto full_ptrs = pointers(full_weights); + void* reference = qwen36_create_training_context( + full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head, + configs, kLayers, static_cast(at::kBFloat16), + 1.0, 1e-3, 0.9, 0.999, 1e-8, kVocab, 1e-5, kLoraRank, + target_layers, kLayers, fixed_targets); + assert(reference); + return {distributed, reference}; +} + +static void set_fixed_lora( + ContextPair contexts, std::vector& fixtures +) { + for (auto& fixture : fixtures) { + const int64_t slot = fixture.layer * kPairsPerLayer + fixture.pair; + assert(qwen36_set_lora_tensor( + contexts.distributed, slot, 0, &fixture.local_a) == 0); + assert(qwen36_set_lora_tensor( + contexts.distributed, slot, 1, &fixture.local_b) == 0); + assert(qwen36_set_lora_tensor( + contexts.reference, slot, 0, &fixture.full_a) == 0); + assert(qwen36_set_lora_tensor( + contexts.reference, slot, 1, &fixture.full_b) == 0); + } +} + +struct ErrorSummary { + double a = 0.0; + double b = 0.0; + double relative = 0.0; + double m = 0.0; + double v = 0.0; + double param = 0.0; + double adam = 0.0; +}; + +static void check_fixed_path( + ContextPair contexts, std::vector& fixtures, + Batch& short_batch, Batch& batch, int rank +) { + const double short_distributed = qwen36_eval_step(contexts.distributed, + &short_batch.input_ids, &short_batch.target_mask, + &short_batch.attention_mask); + const double short_reference = qwen36_eval_step(contexts.reference, + &short_batch.input_ids, &short_batch.target_mask, + &short_batch.attention_mask); + const double eval_distributed = qwen36_eval_step(contexts.distributed, + &batch.input_ids, &batch.target_mask, &batch.attention_mask); + const double eval_reference = qwen36_eval_step(contexts.reference, + &batch.input_ids, &batch.target_mask, &batch.attention_mask); + assert(std::isfinite(short_distributed) && std::isfinite(short_reference)); + assert(std::isfinite(eval_distributed) && std::isfinite(eval_reference)); + + const double micro_distributed = qwen36_train_micro_step( + contexts.distributed, &batch.input_ids, &batch.target_mask, + &batch.attention_mask, 1.0, 0); + const double micro_reference = qwen36_train_micro_step( + contexts.reference, &batch.input_ids, &batch.target_mask, + &batch.attention_mask, 1.0, 0); + assert(std::isfinite(micro_distributed) && std::isfinite(micro_reference)); + + ErrorSummary errors; + for (auto& fixture : fixtures) { + const int64_t slot = fixture.layer * kPairsPerLayer + fixture.pair; + auto* local_a = reinterpret_cast( + qwen36_get_lora_grad_accumulator(contexts.distributed, slot, 0)); + auto* local_b = reinterpret_cast( + qwen36_get_lora_grad_accumulator(contexts.distributed, slot, 1)); + auto* full_a = reinterpret_cast( + qwen36_get_lora_grad_accumulator(contexts.reference, slot, 0)); + auto* full_b = reinterpret_cast( + qwen36_get_lora_grad_accumulator(contexts.reference, slot, 1)); + assert(local_a && local_b && full_a && full_b); + assert(local_a->scalar_type() == at::kFloat && + local_b->scalar_type() == at::kFloat); + auto expected_a = expected_local_factor(*full_a, fixture.module, false, rank); + auto expected_b = expected_local_factor(*full_b, fixture.module, true, rank); + const double a_diff = max_diff(*local_a, expected_a); + const double b_diff = max_diff(*local_b, expected_b); + const double a_relative = relative_l2(*local_a, expected_a); + const double b_relative = relative_l2(*local_b, expected_b); + std::printf( + "native_tp_gdn_grad rank=%d layer=%ld module=%s " + "a_diff=%0.8e b_diff=%0.8e a_relative_l2=%0.8e " + "b_relative_l2=%0.8e a_reference_norm=%0.8e " + "b_reference_norm=%0.8e\n", + rank, fixture.layer, fixture.module, a_diff, b_diff, + a_relative, b_relative, + expected_a.to(at::kFloat).norm().item(), + expected_b.to(at::kFloat).norm().item()); + // Before the final optimizer boundary, replicated factors have only + // their rank-local contribution. Compare shard-local factors here; + // final Adam state checks below cover the synchronized replicas. + if (is_column_parallel(fixture.module)) { + errors.b = std::max(errors.b, b_diff); + errors.relative = std::max(errors.relative, b_relative); + } else { + errors.a = std::max(errors.a, a_diff); + errors.relative = std::max(errors.relative, a_relative); + } + } + assert(qwen36_abort_gradient_accumulation(contexts.distributed) == 0); + assert(qwen36_abort_gradient_accumulation(contexts.reference) == 0); + + const double loss_distributed = qwen36_train_step(contexts.distributed, + &batch.input_ids, &batch.target_mask, &batch.attention_mask); + const double loss_reference = qwen36_train_step(contexts.reference, + &batch.input_ids, &batch.target_mask, &batch.attention_mask); + assert(std::isfinite(loss_distributed) && std::isfinite(loss_reference)); + + constexpr int64_t optimizer_count = 2 * kLayers * kPairsPerLayer; + std::vector local_m(optimizer_count), local_v(optimizer_count); + std::vector full_m(optimizer_count), full_v(optimizer_count); + assert(qwen36_export_optimizer_state(contexts.distributed, + local_m.data(), local_v.data(), optimizer_count) == optimizer_count); + assert(qwen36_export_optimizer_state(contexts.reference, + full_m.data(), full_v.data(), optimizer_count) == optimizer_count); + + for (auto& fixture : fixtures) { + const int64_t slot = fixture.layer * kPairsPerLayer + fixture.pair; + auto* updated_a = reinterpret_cast( + qwen36_get_lora_a(contexts.distributed, slot)); + auto* updated_b = reinterpret_cast( + qwen36_get_lora_b(contexts.distributed, slot)); + auto* reference_a = reinterpret_cast( + qwen36_get_lora_a(contexts.reference, slot)); + auto* reference_b = reinterpret_cast( + qwen36_get_lora_b(contexts.reference, slot)); + assert(updated_a && updated_b && reference_a && reference_b); + + auto* m_a = reinterpret_cast(local_m[2 * slot]); + auto* m_b = reinterpret_cast(local_m[2 * slot + 1]); + auto* v_a = reinterpret_cast(local_v[2 * slot]); + auto* v_b = reinterpret_cast(local_v[2 * slot + 1]); + auto* ref_m_a = reinterpret_cast(full_m[2 * slot]); + auto* ref_m_b = reinterpret_cast(full_m[2 * slot + 1]); + auto* ref_v_a = reinterpret_cast(full_v[2 * slot]); + auto* ref_v_b = reinterpret_cast(full_v[2 * slot + 1]); + assert(m_a && m_b && v_a && v_b && ref_m_a && ref_m_b && + ref_v_a && ref_v_b); + + errors.m = std::max({errors.m, + max_diff(*m_a, expected_local_factor( + *ref_m_a, fixture.module, false, rank)), + max_diff(*m_b, expected_local_factor( + *ref_m_b, fixture.module, true, rank))}); + errors.v = std::max({errors.v, + max_diff(*v_a, expected_local_factor( + *ref_v_a, fixture.module, false, rank)), + max_diff(*v_b, expected_local_factor( + *ref_v_b, fixture.module, true, rank))}); + errors.param = std::max({errors.param, + max_diff(*updated_a, expected_local_factor( + *reference_a, fixture.module, false, rank)), + max_diff(*updated_b, expected_local_factor( + *reference_b, fixture.module, true, rank))}); + + auto expected_a = (fixture.local_a.to(at::kFloat) - 1e-3 * + (*m_a / (1.0 - 0.9)) / + (((*v_a / (1.0 - 0.999)).sqrt()) + 1e-8)).to(at::kBFloat16); + auto expected_b = (fixture.local_b.to(at::kFloat) - 1e-3 * + (*m_b / (1.0 - 0.9)) / + (((*v_b / (1.0 - 0.999)).sqrt()) + 1e-8)).to(at::kBFloat16); + errors.adam = std::max({errors.adam, + max_diff(*updated_a, expected_a), max_diff(*updated_b, expected_b)}); + } + + const double short_diff = std::abs(short_distributed - short_reference); + const double eval_diff = std::abs(eval_distributed - eval_reference); + const double micro_diff = std::abs(micro_distributed - micro_reference); + const double loss_diff = std::abs(loss_distributed - loss_reference); + std::printf( + "native_tp_gdn_fixed rank=%d short_eval_diff=%0.8e eval_diff=%0.8e " + "micro_diff=%0.8e loss_diff=%0.8e a_grad_diff=%0.8e " + "b_grad_diff=%0.8e grad_relative_l2=%0.8e m_diff=%0.8e " + "v_diff=%0.8e param_diff=%0.8e adam_error=%0.8e\n", + rank, short_diff, eval_diff, micro_diff, loss_diff, errors.a, errors.b, + errors.relative, errors.m, errors.v, errors.param, errors.adam); + std::fflush(stdout); + + assert(short_diff < 5e-3 && eval_diff < 5e-3); + assert(micro_diff < 5e-3 && loss_diff < 5e-3); + assert(errors.a < 5e-4 && errors.b < 5e-4); + assert(errors.relative < 2e-2); + assert(errors.m < 5e-5 && errors.v < 5e-8); + assert(errors.param <= 2e-3); + assert(errors.adam < 1e-8); +} + +static void set_dynamic_lora( + void* context, int64_t adapter_id, + std::vector& fixtures, bool local +) { + for (auto& fixture : fixtures) { + auto& a = local ? fixture.local_a : fixture.full_a; + auto& b = local ? fixture.local_b : fixture.full_b; + assert(qwen36_set_adapter_lora_tensor(context, adapter_id, + fixture.layer, fixture.module, 0, &a) == 0); + assert(qwen36_set_adapter_lora_tensor(context, adapter_id, + fixture.layer, fixture.module, 1, &b) == 0); + } +} + +static void check_dynamic_path( + ContextPair contexts, std::vector& fixtures, + Batch& batch, int rank +) { + const int64_t target_layers[kLayers] = {0, 1}; + constexpr const char* targets = + "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj"; + const int64_t distributed_id = qwen36_add_lora(contexts.distributed, + kLoraRank, kLoraRank, target_layers, kLayers, targets); + const int64_t reference_id = qwen36_add_lora(contexts.reference, + kLoraRank, kLoraRank, target_layers, kLayers, targets); + assert(distributed_id > 0 && reference_id > 0); + set_dynamic_lora(contexts.distributed, distributed_id, fixtures, true); + set_dynamic_lora(contexts.reference, reference_id, fixtures, false); + + const double distributed_loss = qwen36_train_multi_lora_selected( + contexts.distributed, &batch.input_ids, &batch.target_mask, + &batch.attention_mask, &distributed_id, 1, kLoraRank); + const double reference_loss = qwen36_train_multi_lora_selected( + contexts.reference, &batch.input_ids, &batch.target_mask, + &batch.attention_mask, &reference_id, 1, kLoraRank); + assert(std::isfinite(distributed_loss) && std::isfinite(reference_loss)); + + ErrorSummary errors; + for (auto& fixture : fixtures) { + auto* local_a = reinterpret_cast( + qwen36_get_adapter_lora_tensor(contexts.distributed, + distributed_id, fixture.layer, fixture.module, 0)); + auto* local_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor(contexts.distributed, + distributed_id, fixture.layer, fixture.module, 1)); + auto* full_a = reinterpret_cast( + qwen36_get_adapter_lora_tensor(contexts.reference, + reference_id, fixture.layer, fixture.module, 0)); + auto* full_b = reinterpret_cast( + qwen36_get_adapter_lora_tensor(contexts.reference, + reference_id, fixture.layer, fixture.module, 1)); + assert(local_a && local_b && full_a && full_b); + errors.param = std::max({errors.param, + max_diff(*local_a, expected_local_factor( + *full_a, fixture.module, false, rank)), + max_diff(*local_b, expected_local_factor( + *full_b, fixture.module, true, rank))}); + + auto state = [&](void* context, int64_t adapter, bool is_b, bool is_v) { + auto* tensor = reinterpret_cast( + qwen36_get_adapter_optimizer_tensor(context, adapter, + fixture.layer, fixture.module, is_b, is_v)); + assert(tensor && tensor->scalar_type() == at::kFloat); + return tensor; + }; + auto* m_a = state(contexts.distributed, distributed_id, false, false); + auto* m_b = state(contexts.distributed, distributed_id, true, false); + auto* v_a = state(contexts.distributed, distributed_id, false, true); + auto* v_b = state(contexts.distributed, distributed_id, true, true); + auto* full_m_a = state(contexts.reference, reference_id, false, false); + auto* full_m_b = state(contexts.reference, reference_id, true, false); + auto* full_v_a = state(contexts.reference, reference_id, false, true); + auto* full_v_b = state(contexts.reference, reference_id, true, true); + errors.m = std::max({errors.m, + max_diff(*m_a, expected_local_factor( + *full_m_a, fixture.module, false, rank)), + max_diff(*m_b, expected_local_factor( + *full_m_b, fixture.module, true, rank))}); + errors.v = std::max({errors.v, + max_diff(*v_a, expected_local_factor( + *full_v_a, fixture.module, false, rank)), + max_diff(*v_b, expected_local_factor( + *full_v_b, fixture.module, true, rank))}); + + auto expected_a = (fixture.local_a.to(at::kFloat) - 1e-3 * + (*m_a / (1.0 - 0.9)) / + (((*v_a / (1.0 - 0.999)).sqrt()) + 1e-8)).to(at::kBFloat16); + auto expected_b = (fixture.local_b.to(at::kFloat) - 1e-3 * + (*m_b / (1.0 - 0.9)) / + (((*v_b / (1.0 - 0.999)).sqrt()) + 1e-8)).to(at::kBFloat16); + errors.adam = std::max({errors.adam, + max_diff(*local_a, expected_a), max_diff(*local_b, expected_b)}); + } + + const double loss_diff = std::abs(distributed_loss - reference_loss); + std::printf( + "native_tp_gdn_dynamic rank=%d loss_diff=%0.8e m_diff=%0.8e " + "v_diff=%0.8e param_diff=%0.8e adam_error=%0.8e\n", + rank, loss_diff, errors.m, errors.v, errors.param, errors.adam); + std::fflush(stdout); + assert(loss_diff < 5e-3); + assert(errors.m < 5e-5 && errors.v < 5e-8); + assert(errors.param <= 2e-3); + assert(errors.adam < 1e-8); +} + +} // namespace + +int main() { + const int rank = required_env_int("RANK"); + const int world = required_env_int("WORLD_SIZE"); + const int local_rank = std::getenv("LOCAL_RANK") + ? std::atoi(std::getenv("LOCAL_RANK")) : rank; + assert(world == 2 && rank >= 0 && rank < world); + assert(qwen36_kernel_abi_version() == kAbiVersion); + qwen36_set_cuda_device(local_rank); + + std::vector full_weights; + full_weights.reserve(kLayers * 14); + for (int64_t layer = 0; layer < kLayers; ++layer) + append_full_gdn_layer(full_weights, layer); + for (auto& weight : full_weights) weight.set_requires_grad(false); + auto local_weights = make_local_weights(full_weights, rank); + for (auto& weight : local_weights) weight.set_requires_grad(false); + assert_weight_contract(full_weights, local_weights, rank); + + auto embed = fingerprint({kVocab, kHidden}, 0.0013, 17); + auto final_norm = unit_weight({kHidden}); + auto lm_head = fingerprint({kVocab, kHidden}, 0.0011, 59); + embed.set_requires_grad(false); + final_norm.set_requires_grad(false); + lm_head.set_requires_grad(false); + LayerConfig configs[kLayers] = {gdn_config(), gdn_config()}; + auto short_batch = make_batch(2, 3, 3); + auto batch = make_batch(2, 9, 11); + auto dynamic_batch = make_batch(1, 9, 23); + + constexpr const char* all_targets = + "in_proj_qkv,in_proj_z,in_proj_a,in_proj_b,out_proj"; + auto fixed_contexts = create_context_pair(local_weights, full_weights, + embed, final_norm, lm_head, configs, all_targets, true); + auto fixed_fixtures = make_lora_fixtures(rank, 2000); + set_fixed_lora(fixed_contexts, fixed_fixtures); + check_fixed_path( + fixed_contexts, fixed_fixtures, short_batch, batch, rank); + qwen36_free_training_context(fixed_contexts.reference); + qwen36_free_training_context(fixed_contexts.distributed); + + auto dynamic_contexts = create_context_pair(local_weights, full_weights, + embed, final_norm, lm_head, configs, "in_proj_qkv", false); + auto dynamic_fixtures = make_lora_fixtures(rank, 4000); + check_dynamic_path( + dynamic_contexts, dynamic_fixtures, dynamic_batch, rank); + qwen36_free_training_context(dynamic_contexts.reference); + qwen36_free_training_context(dynamic_contexts.distributed); + return 0; +} diff --git a/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp index f675963c..07157a9a 100644 --- a/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp @@ -129,7 +129,9 @@ int main() { weight_ptrs.data(), weight_ptrs.size(), &embed, &final_norm, &lm_head, configs, 2, static_cast(at::kBFloat16), 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, - target_layers, 2, "in_proj_qkv", 1); + // This regression exercises latent-rank-only TP with replicated GDN + // base weights. ABI15 bit 0 now explicitly enables base GDN head TP. + target_layers, 2, "in_proj_qkv", 0); assert(distributed && qwen36_init_nccl(distributed) == 0); setenv("TP_SIZE", "1", 1); diff --git a/crates/rustrain-server/src/checkpoint.rs b/crates/rustrain-server/src/checkpoint.rs index 2974f38e..31725e0e 100644 --- a/crates/rustrain-server/src/checkpoint.rs +++ b/crates/rustrain-server/src/checkpoint.rs @@ -20,6 +20,14 @@ pub enum LoraTpShardLayout { #[default] LatentRank, ColumnParallel, + /// Column-parallel projection whose global rows use flat + /// `[Q_all | K_all | V_all]` storage while each rank stores packed + /// `[Q_local | K_local | V_local]` rows. + FlatQkvColumnParallel { + q_rows: i64, + k_rows: i64, + v_rows: i64, + }, RowParallel, } @@ -136,9 +144,20 @@ pub struct TensorShardManifest { #[serde(default)] pub replicated: bool, pub global_offset: Vec, + /// Non-contiguous mappings from rank-local storage into the original + /// global tensor. Empty for ordinary contiguous shards. + #[serde(default)] + pub segments: Vec, pub replica_identity: String, } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct TensorShardSegmentManifest { + pub local_offset: i64, + pub global_offset: i64, + pub length: i64, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CheckpointManifest { pub format: String, @@ -171,7 +190,7 @@ pub fn validate_fixed_tp_resume( .any(|layout| *layout != LoraTpShardLayout::LatentRank) { bail!( - "legacy tensor-parallel v3 checkpoints cannot restore fixed Q/K/V/O LoRA into the projection-aware layout; use a v4 checkpoint or migrate the adapter from a merged artifact" + "legacy tensor-parallel v3 checkpoints cannot restore fixed projection-aware attention LoRA; use a v4 checkpoint or migrate the adapter from a merged artifact" ); } return Ok(()); @@ -203,7 +222,7 @@ pub fn validate_dynamic_tp_resume( .any(|layout| *layout != LoraTpShardLayout::LatentRank) { bail!( - "legacy tensor-parallel v3 checkpoint adapter {adapter_id} contains Q/K/V/O LoRA that cannot be restored into the projection-aware layout; use a v4 checkpoint or migrate the adapter from a merged artifact" + "legacy tensor-parallel v3 checkpoint adapter {adapter_id} contains projection-aware attention LoRA that cannot be restored; use a v4 checkpoint or migrate the adapter from a merged artifact" ); } return Ok(()); @@ -841,6 +860,7 @@ fn tensor_shard( (rank_axis, false) } (LoraTpShardLayout::ColumnParallel, LoraSide::A) + | (LoraTpShardLayout::FlatQkvColumnParallel { .. }, LoraSide::A) | (LoraTpShardLayout::RowParallel, LoraSide::B) => { if local_shape[rank_axis] != global_lora_rank { bail!( @@ -862,6 +882,45 @@ fn tensor_shard( global_offset[axis] = tp_rank * local_shape[axis]; (axis, false) } + ( + LoraTpShardLayout::FlatQkvColumnParallel { + q_rows, + k_rows, + v_rows, + }, + LoraSide::B, + ) => { + if local_shape[rank_axis] != global_lora_rank { + bail!( + "flat-QKV column-parallel checkpoint tensor {file}:{tensor_name} has rank {} on axis {rank_axis}, expected {global_lora_rank}", + local_shape[rank_axis] + ); + } + if q_rows <= 0 + || k_rows <= 0 + || v_rows <= 0 + || q_rows % tp_size != 0 + || k_rows % tp_size != 0 + || v_rows % tp_size != 0 + { + bail!( + "flat-QKV global row segments [{q_rows}, {k_rows}, {v_rows}] must be positive and divisible by TP size {tp_size}" + ); + } + let axis = local_shape.len() - 2; + let local_q = q_rows / tp_size; + let local_k = k_rows / tp_size; + let local_v = v_rows / tp_size; + if local_shape[axis] != local_q + local_k + local_v { + bail!( + "flat-QKV checkpoint tensor {file}:{tensor_name} has {} local rows, expected {} from Q/K/V segments", + local_shape[axis], + local_q + local_k + local_v + ); + } + global_shape[axis] = q_rows + k_rows + v_rows; + (axis, false) + } (LoraTpShardLayout::RowParallel, LoraSide::A) => { if local_shape[rank_axis] != global_lora_rank { bail!( @@ -875,6 +934,38 @@ fn tensor_shard( (axis, false) } }; + let segments = match (layout, side) { + ( + LoraTpShardLayout::FlatQkvColumnParallel { + q_rows, + k_rows, + v_rows, + }, + LoraSide::B, + ) => { + let local_q = q_rows / tp_size; + let local_k = k_rows / tp_size; + let local_v = v_rows / tp_size; + vec![ + TensorShardSegmentManifest { + local_offset: 0, + global_offset: tp_rank * local_q, + length: local_q, + }, + TensorShardSegmentManifest { + local_offset: local_q, + global_offset: q_rows + tp_rank * local_k, + length: local_k, + }, + TensorShardSegmentManifest { + local_offset: local_q + local_k, + global_offset: q_rows + k_rows + tp_rank * local_v, + length: local_v, + }, + ] + } + _ => Vec::new(), + }; Ok(TensorShardManifest { file: file.to_string(), tensor_name, @@ -887,6 +978,7 @@ fn tensor_shard( layout, replicated, global_offset, + segments, replica_identity: if replicated { "tp-replicated".to_string() } else { @@ -1342,24 +1434,43 @@ mod tests { let topology = tp_topology(1, 2); let column_a = Tensor::zeros([4, 8], (tch::Kind::Float, tch::Device::Cpu)); let column_b = Tensor::zeros([8, 4], (tch::Kind::Float, tch::Device::Cpu)); + let flat_qkv_a = Tensor::zeros([4, 8], (tch::Kind::Float, tch::Device::Cpu)); + let flat_qkv_b = Tensor::zeros([12, 4], (tch::Kind::Float, tch::Device::Cpu)); let row_a = Tensor::zeros([4, 4], (tch::Kind::Float, tch::Device::Cpu)); let row_b = Tensor::zeros([8, 4], (tch::Kind::Float, tch::Device::Cpu)); - let lora_a = vec![column_a.shallow_clone(), row_a.shallow_clone()]; - let lora_b = vec![column_b.shallow_clone(), row_b.shallow_clone()]; + let lora_a = vec![ + column_a.shallow_clone(), + flat_qkv_a.shallow_clone(), + row_a.shallow_clone(), + ]; + let lora_b = vec![ + column_b.shallow_clone(), + flat_qkv_b.shallow_clone(), + row_b.shallow_clone(), + ]; let adam_m = vec![ column_a.zeros_like(), column_b.zeros_like(), + flat_qkv_a.zeros_like(), + flat_qkv_b.zeros_like(), row_a.zeros_like(), row_b.zeros_like(), ]; let adam_v = vec![ column_a.ones_like(), column_b.ones_like(), + flat_qkv_a.ones_like(), + flat_qkv_b.ones_like(), row_a.ones_like(), row_b.ones_like(), ]; let layouts = [ LoraTpShardLayout::ColumnParallel, + LoraTpShardLayout::FlatQkvColumnParallel { + q_rows: 4, + k_rows: 4, + v_rows: 16, + }, LoraTpShardLayout::RowParallel, ]; let identities = [ @@ -1368,6 +1479,11 @@ mod tests { layer: 0, module: "q_proj".to_string(), }, + LoraSlotIdentity { + index: 4, + layer: 1, + module: "in_proj_qkv".to_string(), + }, LoraSlotIdentity { index: 3, layer: 0, @@ -1418,20 +1534,125 @@ mod tests { assert_eq!(column_b_shard.global_shape, vec![16, 4]); assert_eq!(column_b_shard.global_offset, vec![8, 0]); - let row_a_shard = shard("a_1"); + let flat_qkv_a_shard = shard("a_1"); + assert_eq!( + flat_qkv_a_shard.layout, + LoraTpShardLayout::FlatQkvColumnParallel { + q_rows: 4, + k_rows: 4, + v_rows: 16, + } + ); + assert!(flat_qkv_a_shard.replicated); + assert_eq!(flat_qkv_a_shard.global_shape, vec![4, 8]); + assert_eq!(flat_qkv_a_shard.global_offset, vec![0, 0]); + + let flat_qkv_b_shard = shard("b_1"); + assert_eq!( + flat_qkv_b_shard.layout, + LoraTpShardLayout::FlatQkvColumnParallel { + q_rows: 4, + k_rows: 4, + v_rows: 16, + } + ); + assert!(!flat_qkv_b_shard.replicated); + assert_eq!(flat_qkv_b_shard.partition_axis, 0); + assert_eq!(flat_qkv_b_shard.global_shape, vec![24, 4]); + assert_eq!(flat_qkv_b_shard.global_offset, vec![0, 0]); + assert_eq!( + flat_qkv_b_shard.segments, + vec![ + TensorShardSegmentManifest { + local_offset: 0, + global_offset: 2, + length: 2, + }, + TensorShardSegmentManifest { + local_offset: 2, + global_offset: 6, + length: 2, + }, + TensorShardSegmentManifest { + local_offset: 4, + global_offset: 16, + length: 8, + }, + ] + ); + + let row_a_shard = shard("a_2"); assert_eq!(row_a_shard.layout, LoraTpShardLayout::RowParallel); assert!(!row_a_shard.replicated); assert_eq!(row_a_shard.partition_axis, 1); assert_eq!(row_a_shard.global_shape, vec![4, 8]); assert_eq!(row_a_shard.global_offset, vec![0, 4]); - let row_b_shard = shard("b_1"); + let row_b_shard = shard("b_2"); assert!(row_b_shard.replicated); assert_eq!(row_b_shard.global_shape, vec![8, 4]); assert_eq!(row_b_shard.global_offset, vec![0, 0]); assert_eq!(row_b_shard.replica_identity, "tp-replicated"); } + #[test] + fn flat_qkv_segments_reconstruct_global_lora_and_optimizer_rows() { + let layout = LoraTpShardLayout::FlatQkvColumnParallel { + q_rows: 4, + k_rows: 4, + v_rows: 16, + }; + let mut reconstructed = [vec![f64::NAN; 24], vec![f64::NAN; 24], vec![f64::NAN; 24]]; + + for rank in 0..2 { + let topology = tp_topology(rank, 2); + let global_rows = (0..24).map(f64::from).collect::>(); + let local_rows = [ + &global_rows[rank * 2..rank * 2 + 2], + &global_rows[4 + rank * 2..4 + rank * 2 + 2], + &global_rows[8 + rank * 8..8 + rank * 8 + 8], + ] + .concat(); + + for (state_index, (state, delta)) in + [("lora_b", 0.0), ("adam_m", 100.0), ("adam_v", 200.0)] + .into_iter() + .enumerate() + { + let values = local_rows + .iter() + .map(|value| value + delta) + .collect::>(); + let tensor = Tensor::from_slice(&values).reshape([12, 1]).repeat([1, 4]); + let shard = tensor_shard( + &topology, + None, + 4, + "state.safetensors", + state.to_string(), + state, + LoraSide::B, + layout, + &tensor, + ) + .unwrap(); + for segment in shard.segments { + for offset in 0..segment.length { + reconstructed[state_index][(segment.global_offset + offset) as usize] = + tensor.double_value(&[segment.local_offset + offset, 0]); + } + } + } + } + + for (state_index, delta) in [0.0, 100.0, 200.0].into_iter().enumerate() { + let expected = (0..24) + .map(|row| f64::from(row) + delta) + .collect::>(); + assert_eq!(reconstructed[state_index], expected); + } + } + #[test] fn tensor_parallel_loader_accepts_legacy_latent_rank_v3_manifest() { let dir = tempfile::tempdir().unwrap(); @@ -1532,7 +1753,9 @@ mod tests { let fixed_error = validate_fixed_tp_resume(&manifest, &[LoraTpShardLayout::ColumnParallel], &[]) .unwrap_err(); - assert!(fixed_error.to_string().contains("fixed Q/K/V/O")); + assert!(fixed_error + .to_string() + .contains("fixed projection-aware attention LoRA")); let dynamic_error = validate_dynamic_tp_resume(&manifest, 17, &[], &[LoraTpShardLayout::RowParallel]) .unwrap_err(); diff --git a/crates/rustrain-server/src/session.rs b/crates/rustrain-server/src/session.rs index ba268017..8ea9568c 100644 --- a/crates/rustrain-server/src/session.rs +++ b/crates/rustrain-server/src/session.rs @@ -10,12 +10,25 @@ use crate::checkpoint; use crate::metrics::{FileMetricsSink, MetricsSink, StepMetric}; use rustrain_qwen3_6::lora::{Qwen36AdapterArtifact, Qwen36LoraConfig, Qwen36LoraTargetModule}; -fn lora_tp_shard_layout(module: Qwen36LoraTargetModule) -> checkpoint::LoraTpShardLayout { +fn lora_tp_shard_layout( + module: Qwen36LoraTargetModule, + config: &rustrain_qwen3_6::config::Qwen36RuntimeConfig, +) -> checkpoint::LoraTpShardLayout { match module { Qwen36LoraTargetModule::QProj | Qwen36LoraTargetModule::KProj - | Qwen36LoraTargetModule::VProj => checkpoint::LoraTpShardLayout::ColumnParallel, - Qwen36LoraTargetModule::OProj => checkpoint::LoraTpShardLayout::RowParallel, + | Qwen36LoraTargetModule::VProj + | Qwen36LoraTargetModule::InProjZ + | Qwen36LoraTargetModule::InProjA + | Qwen36LoraTargetModule::InProjB => checkpoint::LoraTpShardLayout::ColumnParallel, + Qwen36LoraTargetModule::InProjQkv => checkpoint::LoraTpShardLayout::FlatQkvColumnParallel { + q_rows: config.linear_num_key_heads * config.linear_key_head_dim, + k_rows: config.linear_num_key_heads * config.linear_key_head_dim, + v_rows: config.linear_num_value_heads * config.linear_value_head_dim, + }, + Qwen36LoraTargetModule::OProj | Qwen36LoraTargetModule::OutProj => { + checkpoint::LoraTpShardLayout::RowParallel + } _ => checkpoint::LoraTpShardLayout::LatentRank, } } @@ -352,6 +365,29 @@ impl TrainingSession for Qwen36Session { "frozen base TP currently requires MTP to be disabled" )); } + if runtime_config + .layer_types + .iter() + .any(|layer| *layer == rustrain_qwen3_6::config::LayerType::LinearAttention) + && (runtime_config.linear_num_key_heads <= 0 + || runtime_config.linear_num_value_heads <= 0 + || runtime_config.linear_num_value_heads % runtime_config.linear_num_key_heads + != 0 + || runtime_config.linear_num_key_heads % tp_size as i64 != 0 + || runtime_config.linear_num_value_heads % tp_size as i64 != 0 + || runtime_config.linear_key_head_dim != 128 + || runtime_config.linear_value_head_dim != 128 + || runtime_config.linear_conv_kernel_dim <= 0) + { + return Err(anyhow!( + "linear-attention TP requires k/v heads divisible by TP_SIZE with preserved value-head groups, 128-wide key/value heads, and a positive conv kernel: k_heads={}, v_heads={}, key_dim={}, value_dim={}, conv_kernel={}, tp={tp_size}", + runtime_config.linear_num_key_heads, + runtime_config.linear_num_value_heads, + runtime_config.linear_key_head_dim, + runtime_config.linear_value_head_dim, + runtime_config.linear_conv_kernel_dim, + )); + } if runtime_config.num_attention_heads <= 0 || runtime_config.num_attention_heads % tp_size as i64 != 0 || runtime_config.num_key_value_heads <= 0 @@ -445,13 +481,27 @@ impl TrainingSession for Qwen36Session { weights.insert(name, narrowed); } else { let local_shard = if base_tp_attention { - let attention_shard = + let full_attention_shard = rustrain_qwen3_6::kernel::shard_full_attention_weight_for_tp( &name, &tensor, tp_size, ep_rank % tp_size, )?; + let attention_shard = if full_attention_shard.is_some() { + full_attention_shard + } else { + rustrain_qwen3_6::kernel::shard_linear_attention_weight_for_tp( + &name, + &tensor, + tp_size, + ep_rank % tp_size, + runtime_config.linear_num_key_heads, + runtime_config.linear_key_head_dim, + runtime_config.linear_num_value_heads, + runtime_config.linear_value_head_dim, + )? + }; if attention_shard.is_some() || !base_tp_mlp { attention_shard } else { @@ -714,7 +764,7 @@ impl TrainingSession for Qwen36Session { adam_m.push(all_adam_m[optimizer_index + 1].shallow_clone()); adam_v.push(all_adam_v[optimizer_index].shallow_clone()); adam_v.push(all_adam_v[optimizer_index + 1].shallow_clone()); - fixed_shard_layouts.push(lora_tp_shard_layout(slot.module)); + fixed_shard_layouts.push(lora_tp_shard_layout(slot.module, &runtime_config)); fixed_slot_identities.push(checkpoint::LoraSlotIdentity { index: slot.index, layer: slot.layer, @@ -729,7 +779,7 @@ impl TrainingSession for Qwen36Session { let shard_layouts = slots .iter() .filter(|slot| slot.active) - .map(|slot| lora_tp_shard_layout(slot.module)) + .map(|slot| lora_tp_shard_layout(slot.module, &runtime_config)) .collect::>(); let mut dynamic_a = Vec::new(); let mut dynamic_b = Vec::new(); @@ -892,7 +942,7 @@ impl TrainingSession for Qwen36Session { let expected_fixed_layouts = fixed_slots .iter() .filter(|slot| slot.active) - .map(|slot| lora_tp_shard_layout(slot.module)) + .map(|slot| lora_tp_shard_layout(slot.module, &runtime_config)) .collect::>(); let expected_fixed_identities = fixed_slots .iter() @@ -928,7 +978,7 @@ impl TrainingSession for Qwen36Session { rustrain_qwen3_6::lora::native_lora_slots(&runtime_config, &config) .iter() .filter(|slot| slot.active) - .map(|slot| lora_tp_shard_layout(slot.module)) + .map(|slot| lora_tp_shard_layout(slot.module, &runtime_config)) .collect::>(); checkpoint::validate_dynamic_tp_resume( &data.manifest, diff --git a/docs/agent/linear-attention.md b/docs/agent/linear-attention.md index f5589545..fe289a11 100644 --- a/docs/agent/linear-attention.md +++ b/docs/agent/linear-attention.md @@ -40,6 +40,37 @@ S = S + k_t ⊗ delta # state update out_t = S @ q_t # output ``` +## Tensor Parallel Layout + +GDN TP partitions K/V head groups while preserving `n_rep = V_heads / K_heads`: + +- `in_proj_qkv.weight` and `conv1d.weight`: slice Q, K, and V segments independently, then repack each rank as `[Q_local | K_local | V_local]`. +- `in_proj_z.weight`: shard output rows by local V heads. +- `in_proj_a.weight`, `in_proj_b.weight`, `A_log`, `dt_bias`: shard by local V-head rows. +- `norm.weight`: replicate `[value_head_dim]`. +- `out_proj.weight`: shard input columns matching local V heads, then all-reduce the local output. +- Column-parallel LoRA replicates A and shards B; row-parallel LoRA shards A and replicates B. Replicated-factor gradients are summed once at the optimizer boundary. + +The v4 flat-QKV checkpoint manifest records three local-to-global row segments for Q, K, and V. Merge or reshard must apply those segments instead of treating the packed rank-local rows as one contiguous global slice. + +One `TpCopyToRegion` must wrap the shared input before the QKV/Z/A/B forks so backward sums their input-gradient contributions once. Do not all-reduce each fork separately. + +## Backward Stability + +The current fused CUDA backward reconstructs earlier states by dividing the decayed state by `g_exp`. This is fast and verified with realistic negative `dt_bias`, where decay stays near one, but it is ill-conditioned for synthetic long sequences with decay near `0.5`. A stable production replacement should checkpoint recurrent state by chunks and replay each chunk during backward instead of repeatedly inverting the decay. + +## Native GDN TP Verification + +Use a Python environment with ABI-compatible prebuilt PyTorch, CUDA, and NCCL, then run: + +```bash +PYTHON=/path/to/python scripts/run_qwen36_native_gdn_tp.sh smoke +PYTHON=/path/to/python scripts/run_qwen36_native_gdn_tp.sh bench-single +PYTHON=/path/to/python scripts/run_qwen36_native_gdn_tp.sh bench-tp2 +``` + +The script builds only the repository kernel and native harnesses. It discovers and links the prebuilt dependency files from the selected Python environment; missing headers or libraries are reported instead of building third-party dependencies. + ## L2 Normalization - Transformers: `x * rsqrt(sum(x²) + eps)` (eps=1e-6, in denominator) diff --git a/docs/qwen35-qwen36-megatron-audit.md b/docs/qwen35-qwen36-megatron-audit.md index cbdba779..97cd43bd 100644 --- a/docs/qwen35-qwen36-megatron-audit.md +++ b/docs/qwen35-qwen36-megatron-audit.md @@ -7,7 +7,7 @@ - 模型语义：Qwen3.5 dense、Qwen3.6 dense/MoE 的 native forward/backward 路径已经覆盖 hybrid full attention、GDN、MoE、MTP 和 LoRA 目标模块；已有配置解析、合成 oracle、集成测试及 H20 native smoke 证据，但尚未完成真实 35B/3.6 权重的长时间训练验证。 - 已实现并可验证的分布式子集：MoE expert parallel，以及 replicated LoRA 的 data parallel；梯度累积和 dynamic multi-LoRA 已有 logical-step 边界。DP 动态租户按 adapter token count 加权，sharded A2A native 路径会保留 source flattened row 来恢复租户，并按全局租户 token count 归一化。 - 性能：MoE grouped dispatch 相对逐 expert matmul 的已有 microbenchmark 为约 3.70x（E=32, N=4096, H=2048, I=768，结果误差为 0）；这不是端到端训练吞吐或 Megatron 对比。 -- 已实现 LoRA latent-rank TP-only、frozen full-attention TP（Q/K/V ColumnParallel、O RowParallel），以及 frozen dense SwiGLU MLP 的 gate/up row shard、down column shard 和输出 all-reduce。Q/K/V/O fixed 与 selected dynamic LoRA 使用 projection-aware shard 和梯度 reduction；GDN/MoE/embedding/LM-head 仍复制，MLP LoRA 在 base TP 下暂拒绝；PP/CP 和 TP 与 EP/DP 的组合仍未实现。 +- 已实现 LoRA latent-rank TP-only、frozen full-attention TP、frozen GDN TP，以及 frozen dense SwiGLU MLP TP。GDN 按 K/V head group 切分复合 QKV、depthwise conv、Z/A/B、A_log/dt_bias 和 out-proj input columns；fixed 与 selected dynamic LoRA 使用 projection-aware shard 和梯度 reduction。MoE/embedding/LM-head 仍复制，MLP LoRA 在 base TP 下暂拒绝；PP/CP 和 TP 与 EP/DP 的组合仍未实现。 - 因此当前实现不能宣称“Megatron-LM 级别”。它是一个计算集中在 C++ 的 LoRA/EP/DP 子集，离 Megatron 的完整并行和通信重叠仍有实质差距。 ## 当前能力矩阵 @@ -23,7 +23,7 @@ | microbatch accumulation | 已实现子集 | non-final microbatch 只 backward，final microbatch 才 optimizer；FP32 accumulator 存储/聚合，autograd leaf backward 仍为 BF16 | | replicated data parallel | 已实现 | logical-step 边界同步 replicated LoRA；EP expert 参数不走该 reduction | | expert parallel | 已实现子集 | 默认 routed-output all-reduce；gated variable-split A2A 已验证 fixed-LoRA 和 native dynamic-LoRA data sharding；GPU-only split planning、异步 overlap 和 DeepEP backend 未实现 | -| tensor parallel | full attention + dense MLP 子集 | full attention 的 Q/K/V 输出头分片、O 输入列分片和 projection-aware fixed/dynamic LoRA 已通过 TP2 full-reference smoke；dense gate/up/down base 权重按 intermediate 维切分；GDN/MoE/embedding/LM-head 仍不切分，MLP LoRA/MTP 暂拒绝 | +| tensor parallel | full attention + GDN + dense MLP 子集 | full attention 与 GDN 均使用 head-aligned ColumnParallel input projection 和 RowParallel output projection；GDN 的 flat QKV/conv 按 `[Q_local|K_local|V_local]` 重排。两种 attention 的 fixed/dynamic LoRA 均通过 TP2 full-reference smoke；dense gate/up/down base 权重按 intermediate 维切分；MoE/embedding/LM-head 仍不切分，MLP LoRA/MTP 暂拒绝 | | pipeline parallel | 未实现于 Qwen native | 没有 stage 切分、microbatch scheduler 或 activation send/recv | | context parallel | 未实现于 Qwen native | 没有 ring attention、跨 rank KV/索引合并 | | distributed checkpoint | 已实现子集 | v4 记录 projection layout、replicated tensor geometry 和 fixed slot identity，并验证 same-topology rank shard；v3 仅兼容 latent-rank checkpoint，旧 attention LoRA 因形状不可迁移而明确拒绝；跨 topology reshard 和 PP/CP 未实现 | @@ -34,7 +34,7 @@ Megatron 的通用 MLP 使用 ColumnParallel fc1、local gated activation 和 RowParallel fc2，并进一步提供 fused fc1/activation、sequence parallel、通信重叠和 sharded-state 支持。当前 Qwen native 已补上算法等价的 separate gate/up row shard 与 down column shard，但仍是两次独立 GEMM。 -本地 Megatron-LM 的 `experimental/lite/megatron/lite/model/qwen3_5` 已包含直接的 Qwen3.5 实现和 LoRA adapter，而不是只能依赖外部 bridge：`primitive/modules/gqa.py` 使用 fused QKV ColumnParallel 和 O RowParallel；`gated_delta_net.py` 对输入/输出投影和 local heads 做 TP，并接入 sequence parallel、context parallel 与 FLA 高性能路径；`lora.py` 的 `LinearLoRA` 还能拆分 latent rank/output 并配套 gather、reduce-scatter 或 input-gradient all-reduce。没有发现显式的 Qwen3.6 model registration。当前 Qwen native 的 full-attention TP 数学布局已经对齐，但 Q/K/V 仍是三次独立 GEMM，LoRA 的 replicated 一侧也比 Megatron Lite 更保守；GDN、MoE、embedding/LM-head 和多轴并行仍是主要差距。 +本地 Megatron-LM 的 `experimental/lite/megatron/lite/model/qwen3_5` 已包含直接的 Qwen3.5 实现和 LoRA adapter，而不是只能依赖外部 bridge：`primitive/modules/gqa.py` 使用 fused QKV ColumnParallel 和 O RowParallel；`gated_delta_net.py` 对输入/输出投影和 local heads 做 TP，并接入 sequence parallel、context parallel 与 FLA 高性能路径；`lora.py` 的 `LinearLoRA` 还能拆分 latent rank/output 并配套 gather、reduce-scatter 或 input-gradient all-reduce。没有发现显式的 Qwen3.6 model registration。当前 Qwen native 的 full-attention/GDN TP 数学布局已经对齐，但 full attention 的 Q/K/V 仍是三次独立 GEMM，GDN 没有 sequence/context parallel 或 FLA chunk kernel，LoRA 的 replicated 一侧也比 Megatron Lite 更保守；MoE、embedding/LM-head 和多轴并行仍是主要差距。 PP 会把层分到不同 stage 并使用 1F1B 等调度；CP 会在序列和 attention state 上做跨 rank 通信。当前 Qwen native `TrainingContext` 仍在每个进程执行完整层栈，因此仅增加 PP/CP 配置不能得到正确语义。 @@ -46,7 +46,9 @@ PP 会把层分到不同 stage 并使用 1F1B 等调度；CP 会在序列和 att ### 性能工程 -当前粗粒度 C++ FFI、grouped MoE 和 activation checkpoint/offload 是有效优化，但 full attention 仍通过 ATen linear/SDPA，QKV 未融合，TP collective 同步执行，GDN 和 vocab 路径仍复制。尚无 Megatron/Transformer Engine 级别的端到端数据：没有完整模型在同一 GPU、序列长度、microbatch、精度和通信配置下的 tokens/s、显存、扩展效率对照，也没有 FP8/FP4 参数与 fused attention/DeepEP 的 Qwen 路径。 +当前粗粒度 C++ FFI、grouped MoE、GDN head TP 和 activation checkpoint/offload 是有效优化，但 full attention 仍通过 ATen linear/SDPA，QKV 未融合，TP collective 同步执行，vocab 路径仍复制。GDN fused backward 通过反除 decay 重建历史 state；真实 time-step bias 下可运行，但极小 decay 的长序列会数值不稳，仍需 state checkpoint/chunk backward 替代。尚无 Megatron/Transformer Engine 级别的端到端数据：没有完整模型在同一 GPU、序列长度、microbatch、精度和通信配置下的 tokens/s、显存、扩展效率对照，也没有 FP8/FP4 参数与 fused attention/DeepEP 的 Qwen 路径。 + +ABI15 synthetic GDN TP benchmark 使用 3 层、S=512、H=2048、16 K heads、32 V heads 和 LoRA rank 8。在 B=2 时 single/TP2 p50 为 `81.10/约 76.06 ms`，只有约 `1.07x`；在 B=8 时为 `303.25/约 158.04 ms`，达到约 `1.92x`，每卡 observed resident 从 `4.59 GiB` 降到 `3.54 GiB`。`nsys` 的 B=2 TP2 trace 中 fused delta-rule backward 占 GPU kernel time `80.3%`、forward 占 `7.4%`，NCCL all-reduce 合计低于 `0.4%`。小 batch 下 single 和 TP2 都只需相近的 persistent-block wave 数，因此不能期待 head TP 自动加速；dynamic multi-LoRA batching 提高 BH 后才接近线性 scaling。这是 synthetic native 结果，不是完整模型或 matched Megatron 对比。 本次 native benchmark 没有证明 gated A2A 的端到端 step-time 优势：在该小型 workload 上 sharded A2A 的中位 step 反而比 legacy 高约 `23%`。它没有实现 DeepEP 的 fused permutation、GPU-only split planning 或通信计算 overlap，也没有覆盖 H=`2048`/E=`256` 的完整 Qwen3.6 workload。legacy 模式复制输入 batch，因此必须同时报告唯一样本吞吐，不能只看所有 rank 的 processed tokens/s。 @@ -56,12 +58,14 @@ PP 会把层分到不同 stage 并使用 1F1B 等调度；CP 会在序列和 att 已运行的验证包括 Rust 编译检查、core 单测、Qwen3.6 配置/集成测试，以及 H20 ABI11 的单卡、TP2、EP2 和 DP2 native smoke。DP2 的 weighted m/grouped/v/Adam BF16 delta oracle 分别达到 `2.43e-8`、`2.27e-8`、`7.33e-8` 和 `0`；legacy EP2 与 full-expert reference 的 loss、LoRA、Adam state 和标准 Adam 首步 oracle 在两 rank 均为零差异。Replicated A2A 与 fixed-LoRA sharded A2A 也均通过两 rank full-expert reference；sharded token counts `[1,3]` 的加权 loss 与 global reference 相差约 `9.6e-7`，m/v 最大差 `1.22e-5` / `3.92e-9`，Adam oracle 差为 `0`。H20 ABI1 dynamic sharded full-reference smoke 在两 rank 返回 `0`：dynamic grouped-expert 参数最大差 `1.53e-5`，m/v 最大差 `4.88e-5` / `5.75e-8`。ABI13 dense base-MLP TP2 smoke 对 gate/up/down 使用半尺寸本地权重，eval/train loss 与完整权重参考相差 `1.84e-5`。 -ABI14 的 4Q/2KV-head GQA full-attention TP2 smoke 覆盖 Q/K/V/O fixed 和 selected dynamic LoRA。fixed eval/loss 与完整参考最大差 `4.40e-4`；Q/K/V-B 与 O-A 梯度最大差分别为 `1.83e-4`、`1.14e-5`、`3.05e-4` 和 `1.83e-4`，FP32 Adam m/v 最大差 `1.53e-5` / `6.63e-9`，本地标准 Adam 公式误差小于 `3.73e-9`。selected dynamic loss 最大差 `6.82e-5`，Q/K/V/O 参数最大差 `3.05e-5`。另一个两层 GDN smoke 直接验证 latent-rank TP 的 input-dgrad backward all-reduce：A/B 梯度最大差 `9.54e-7` / `4.77e-7`，m/v 最大差 `4.77e-8` / `3.71e-14`，本地 Adam 误差为 `0`。BF16 参数直接对照最多跨约两个量化 bin，因此验收以 FP32 梯度、m/v 和本地 Adam 公式为主。没有完成完整大模型长时间训练、跨节点通信、GDN/MoE/vocab base TP、PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖且有直接 oracle 的子集，而不是所有并行配置。 +ABI14 的 4Q/2KV-head GQA full-attention TP2 smoke 覆盖 Q/K/V/O fixed 和 selected dynamic LoRA。fixed eval/loss 与完整参考最大差 `4.40e-4`；Q/K/V-B 与 O-A 梯度最大差分别为 `1.83e-4`、`1.14e-5`、`3.05e-4` 和 `1.83e-4`，FP32 Adam m/v 最大差 `1.53e-5` / `6.63e-9`，本地标准 Adam 公式误差小于 `3.73e-9`。selected dynamic loss 最大差 `6.82e-5`，Q/K/V/O 参数最大差 `3.05e-5`。 + +ABI15 的两层 GDN base-TP smoke 覆盖复合 QKV/conv head shard、Z/A/B/A_log/dt_bias、replicated norm、out-proj input columns，以及五种 GDN projection 的 fixed/selected dynamic LoRA。fixed loss 最大差 `1.30e-4`，rank-local factor 梯度最大差 `2.44e-4`，FP32 Adam m/v 最大差 `3.43e-6` / `2.49e-10`；dynamic loss 最大差 `1.16e-3`，m/v 最大差 `6.10e-6` / `1.22e-9`，标准 Adam 公式误差均为 `0`。latent-rank-only GDN TP2 回归也通过，loss 差 `5.46e-5`。BF16 参数直接对照最多跨约两个量化 bin，因此验收以 FP32 梯度、m/v 和本地 Adam 公式为主。没有完成完整大模型长时间训练、跨节点通信、MoE/vocab base TP、PP/CP smoke 或与 Megatron-LM 的同条件 benchmark。因此“正确”应理解为已覆盖且有直接 oracle 的子集，而不是所有并行配置。 ## 继续达到 Megatron 级别所需的最小工作包 1. 把现有 5D TP/PP/DP/EP/CP topology contract 落成可组合 runtime process groups、launcher 和 checkpoint rank mapping；当前 native 仍拒绝多轴组合。 -2. 补齐 Qwen GDN attention、MoE、embedding/LM-head/CE 的 TP shard，并为 dense MLP 增加 projection-aware LoRA、fused gate/up FC1 和 MTP 支持；为 full attention 融合 QKV/SDPA 并增加 sequence parallel；为 PP 实现 stage forward/backward 与 1F1B scheduler；为 CP 实现 ring attention/state exchange。 +2. 补齐 MoE、embedding/LM-head/CE 的 TP shard，并为 dense MLP 增加 projection-aware LoRA、fused gate/up FC1 和 MTP 支持；为 GDN 增加稳定的 chunk/state-checkpoint backward 与 sequence/context parallel，为 full attention 融合 QKV/SDPA 并增加 sequence parallel；为 PP 实现 stage forward/backward 与 1F1B scheduler；为 CP 实现 ring attention/state exchange。 3. 将 EP dispatch/combine 替换为 fused/异步路径，并测量通信与计算重叠。 4. 为 checkpoint 增加跨 topology reshard、可恢复的 pending accumulation state，并为旧 v3 attention checkpoint 提供离线迁移工具。 5. 在固定硬件和 workload 上，与 Megatron-LM 记录 tokens/s、step time、峰值显存、通信占比和 loss 曲线。 diff --git a/scripts/run_qwen36_native_gdn_tp.sh b/scripts/run_qwen36_native_gdn_tp.sh new file mode 100755 index 00000000..d3ae0ebd --- /dev/null +++ b/scripts/run_qwen36_native_gdn_tp.sh @@ -0,0 +1,192 @@ +#!/usr/bin/env bash +set -euo pipefail + +usage() { + echo "usage: $0 {smoke|bench-single|bench-tp2}" >&2 + exit 2 +} + +mode="${1:-}" +case "$mode" in + smoke|bench-single|bench-tp2) ;; + *) usage ;; +esac + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$repo_root" + +python_bin="${PYTHON:-python3}" +readarray -t torch_config < <( + "$python_bin" - <<'PY' +import pathlib +import torch + +root = pathlib.Path(torch.__file__).resolve().parent +print(root / "include") +print(root / "lib") +print(int(torch._C._GLIBCXX_USE_CXX11_ABI)) +print(root.parent) +print(torch.__version__) +PY +) +torch_include="${TORCH_INCLUDE_PATH:-${torch_config[0]}}" +torch_lib="${TORCH_LIB_PATH:-${torch_config[1]}}" +cxx11_abi="${GLIBCXX_USE_CXX11_ABI:-${torch_config[2]}}" +site_packages="${torch_config[3]}" +torch_version="${torch_config[4]}" + +cuda_home="${CUDA_HOME:-/usr/local/cuda}" +cuda_include="${CUDA_INCLUDE_PATH:-$cuda_home/include}" +nccl_include="${NCCL_INCLUDE_PATH:-$site_packages/nvidia/nccl/include}" +nccl_lib="${NCCL_LIB_PATH:-$site_packages/nvidia/nccl/lib}" + +for required in \ + "$torch_include/ATen/ATen.h" \ + "$torch_lib/libtorch.so" \ + "$cuda_include/cuda_runtime.h" \ + "$nccl_include/nccl.h"; do + if [[ ! -e "$required" ]]; then + echo "required prebuilt dependency file not found: $required" >&2 + exit 1 + fi +done + +export LIBTORCH_USE_PYTORCH=1 +export LIBTORCH_BYPASS_VERSION_CHECK=1 +export TORCH_INCLUDE_PATH="$torch_include" +export TORCH_LIB_PATH="$torch_lib" +export CUDA_INCLUDE_PATH="$cuda_include" +export NCCL_INCLUDE_PATH="$nccl_include" +export NCCL_LIB_PATH="$nccl_lib" +export GLIBCXX_USE_CXX11_ABI="$cxx11_abi" + +if [[ -e "$nccl_lib/libnccl.so" ]]; then + nccl_link="-lnccl" + nccl_file="$nccl_lib/libnccl.so" +elif [[ -e "$nccl_lib/libnccl.so.2" ]]; then + nccl_link="-l:libnccl.so.2" + nccl_file="$nccl_lib/libnccl.so.2" +else + echo "prebuilt NCCL library not found under $nccl_lib" >&2 + exit 1 +fi + +nvcc="$cuda_home/bin/nvcc" +if [[ ! -x "$nvcc" ]]; then + echo "CUDA compiler not found: $nvcc" >&2 + exit 1 +fi + +for tool in g++ sha256sum stat; do + if ! command -v "$tool" >/dev/null; then + echo "required build tool not found: $tool" >&2 + exit 1 + fi +done + +fingerprint="$({ + printf '%s\n' \ + "python=$python_bin" \ + "torch_version=$torch_version" \ + "torch_include=$torch_include" \ + "torch_lib=$torch_lib" \ + "cxx11_abi=$cxx11_abi" \ + "cuda_home=$cuda_home" \ + "cuda_include=$cuda_include" \ + "nccl_include=$nccl_include" \ + "nccl_lib=$nccl_lib" \ + "nccl_link=$nccl_link" + g++ --version + "$nvcc" --version + sha256sum "$0" + stat -Lc '%n:%s:%Y' \ + "$torch_lib/libtorch.so" \ + "$torch_lib/libtorch_cuda.so" \ + "$torch_lib/libc10.so" \ + "$nccl_file" \ + "$cuda_home/lib64/libcudart.so" +} | sha256sum | cut -c1-20)" +native_root="${NATIVE_BUILD_DIR:-target/native-qwen36-gdn}" +native_dir="$native_root/$fingerprint" +mkdir -p "$native_dir" + +kernel_lib="$native_dir/libqwen36_kernels.so" +kernel_sources=( + crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp + crates/rustrain-qwen3-6/kernels/delta_rule.cu + crates/rustrain-qwen3-6/kernels/delta_rule.cuh + crates/rustrain-qwen3-6/kernels/fused_kernels.cu +) +rebuild_kernel=false +if [[ ! -e "$kernel_lib" ]]; then + rebuild_kernel=true +else + for source in "${kernel_sources[@]}"; do + if [[ "$source" -nt "$kernel_lib" ]]; then + rebuild_kernel=true + break + fi + done +fi + +if [[ "$rebuild_kernel" == true ]]; then + cuda_objects=() + for source in delta_rule fused_kernels; do + object="$native_dir/$source.o" + "$nvcc" -c "crates/rustrain-qwen3-6/kernels/$source.cu" -o "$object" \ + -O2 -std=c++17 "-D_GLIBCXX_USE_CXX11_ABI=$cxx11_abi" \ + "-I$torch_include" "-I$cuda_include" -Xcompiler -fPIC + cuda_objects+=("$object") + done + + g++ -shared -fPIC -std=c++17 -O2 \ + "-D_GLIBCXX_USE_CXX11_ABI=$cxx11_abi" \ + crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp \ + "${cuda_objects[@]}" -o "$kernel_lib" \ + "-I$torch_include" "-I$cuda_include" "-I$nccl_include" \ + "-L$torch_lib" "-L$nccl_lib" "-L$cuda_home/lib64" \ + "-Wl,-rpath,$torch_lib" "-Wl,-rpath,$nccl_lib" \ + "-Wl,-rpath,$cuda_home/lib64" -Wl,--no-as-needed \ + -ltorch -ltorch_cuda -ltorch_cpu -lc10 -lc10_cuda -lcudart "$nccl_link" +fi +kernel_dir="$native_dir" + +common_flags=( + -std=c++17 -O2 "-D_GLIBCXX_USE_CXX11_ABI=$cxx11_abi" + "-I$torch_include" "-I$cuda_include" "-I$nccl_include" + "-L$kernel_dir" "-L$torch_lib" "-L$nccl_lib" "-L$cuda_home/lib64" + "-Wl,-rpath,$kernel_dir" "-Wl,-rpath,$torch_lib" + "-Wl,-rpath,$nccl_lib" "-Wl,-rpath,$cuda_home/lib64" + -Wl,--no-as-needed -Wl,--allow-shlib-undefined + -lqwen36_kernels -ltorch -ltorch_cuda -ltorch_cpu -lc10 -lc10_cuda + -lcudart "$nccl_link" +) + +smoke_bin="$native_dir/native_tp_gdn_smoke" +bench_bin="$native_dir/native_tp_gdn_bench" +smoke_source=crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp +bench_source=crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp +if [[ ! -e "$smoke_bin" || "$smoke_source" -nt "$smoke_bin" || "$kernel_lib" -nt "$smoke_bin" ]]; then + g++ "$smoke_source" -o "$smoke_bin" "${common_flags[@]}" +fi +if [[ ! -e "$bench_bin" || "$bench_source" -nt "$bench_bin" || "$kernel_lib" -nt "$bench_bin" ]]; then + g++ "$bench_source" -o "$bench_bin" "${common_flags[@]}" +fi + +cu13_lib="$site_packages/nvidia/cu13/lib" +export LD_LIBRARY_PATH="$kernel_dir:$torch_lib:$nccl_lib:$cuda_home/lib64:$cu13_lib:${LD_LIBRARY_PATH:-}" +export RUSTRAIN_NCCL_RUN_ID="${RUSTRAIN_NCCL_RUN_ID:-qwen36-gdn-$$}" + +case "$mode" in + smoke) + TP_SIZE=2 "$python_bin" -m torch.distributed.run --standalone \ + --nnodes=1 --nproc-per-node=2 --no-python "$smoke_bin" + ;; + bench-single) + BENCH_MODE=single WORLD_SIZE=1 RANK=0 LOCAL_RANK=0 "$bench_bin" + ;; + bench-tp2) + BENCH_MODE=tp2 TP_SIZE=2 "$python_bin" -m torch.distributed.run \ + --standalone --nnodes=1 --nproc-per-node=2 --no-python "$bench_bin" + ;; +esac From 99aecef5e225b1d7f25a0d3bf0deb994dd052b67 Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 15:19:17 +0800 Subject: [PATCH 027/156] feat: add qwen tensor data parallelism --- Cargo.lock | 1 + crates/rustrain-ipc/src/command.rs | 58 ++ crates/rustrain-parallel/src/topology.rs | 17 + .../kernels/qwen3_6_kernels.cpp | 561 +++++++++++++----- crates/rustrain-qwen3-6/src/kernel.rs | 49 +- crates/rustrain-qwen3-6/src/session.rs | 125 ++-- .../rustrain-qwen3-6/tests/native_smoke.cpp | 2 +- .../tests/native_tp_dp_smoke.cpp | 534 +++++++++++++++++ .../tests/native_tp_gdn_bench.cpp | 2 +- .../tests/native_tp_gdn_smoke.cpp | 2 +- .../tests/native_tp_latent_smoke.cpp | 2 +- crates/rustrain-server/Cargo.toml | 1 + crates/rustrain-server/src/api.rs | 123 +++- crates/rustrain-server/src/ep.rs | 330 +++++++++-- crates/rustrain-server/src/session.rs | 109 +++- scripts/run_qwen36_native_gdn_tp.sh | 38 +- src/main.rs | 4 + 17 files changed, 1691 insertions(+), 267 deletions(-) create mode 100644 crates/rustrain-qwen3-6/tests/native_tp_dp_smoke.cpp diff --git a/Cargo.lock b/Cargo.lock index 52308319..838be979 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2498,6 +2498,7 @@ dependencies = [ "rustrain-checkpoint", "rustrain-core", "rustrain-ipc", + "rustrain-parallel", "rustrain-qwen3-6", "rustrain-train", "safetensors 0.8.0", diff --git a/crates/rustrain-ipc/src/command.rs b/crates/rustrain-ipc/src/command.rs index 4c77edc7..0f47f7b3 100644 --- a/crates/rustrain-ipc/src/command.rs +++ b/crates/rustrain-ipc/src/command.rs @@ -57,6 +57,8 @@ pub enum EpCommand { input_ids: Vec, target_mask: Vec, attention_mask: Vec, + #[serde(default = "default_batch_size")] + batch_size: usize, seq_len: usize, }, TrainMultiLora { @@ -64,6 +66,8 @@ pub enum EpCommand { input_ids: Vec, target_mask: Vec, attention_mask: Vec, + #[serde(default = "default_batch_size")] + batch_size: usize, seq_len: usize, n_total: i32, lora_rank: i32, @@ -88,6 +92,10 @@ pub enum EpCommand { Shutdown, } +const fn default_batch_size() -> usize { + 1 +} + /// Results that workers return to the HTTP server. #[derive(Debug, Clone, Serialize, Deserialize)] pub enum EpResult { @@ -114,3 +122,53 @@ impl EpResult { EpResult::Error(msg.into()) } } + +#[cfg(test)] +mod tests { + use super::EpCommand; + + #[test] + fn train_step_serde_preserves_batch_and_sequence_shape() { + let command = EpCommand::TrainStep { + session_id: "session".into(), + input_ids: vec![1, 2, 3, 4, 5, 6], + target_mask: vec![1; 6], + attention_mask: vec![1; 6], + batch_size: 2, + seq_len: 3, + }; + + let json = serde_json::to_string(&command).unwrap(); + let decoded: EpCommand = serde_json::from_str(&json).unwrap(); + match decoded { + EpCommand::TrainStep { + batch_size, + seq_len, + .. + } => { + assert_eq!(batch_size, 2); + assert_eq!(seq_len, 3); + } + other => panic!("unexpected command: {other:?}"), + } + } + + #[test] + fn legacy_train_step_command_defaults_to_one_batch_row() { + let json = r#"{ + "TrainStep": { + "session_id": "session", + "input_ids": [1, 2, 3], + "target_mask": [1, 1, 1], + "attention_mask": [1, 1, 1], + "seq_len": 3 + } + }"#; + + let decoded: EpCommand = serde_json::from_str(json).unwrap(); + match decoded { + EpCommand::TrainStep { batch_size, .. } => assert_eq!(batch_size, 1), + other => panic!("unexpected command: {other:?}"), + } + } +} diff --git a/crates/rustrain-parallel/src/topology.rs b/crates/rustrain-parallel/src/topology.rs index b0e59137..ed1215e5 100644 --- a/crates/rustrain-parallel/src/topology.rs +++ b/crates/rustrain-parallel/src/topology.rs @@ -480,6 +480,23 @@ mod tests { assert_eq!(topology.pipeline_group(rank).unwrap(), vec![3, 9, 15, 21]); } + #[test] + fn tp2_dp2_builds_orthogonal_groups_for_every_rank() { + let topology = ParallelTopology::new(2, 1, 2, 1, 1).unwrap(); + let expected = [ + (0, 0, vec![0, 1], vec![0, 2]), + (1, 0, vec![0, 1], vec![1, 3]), + (0, 1, vec![2, 3], vec![0, 2]), + (1, 1, vec![2, 3], vec![1, 3]), + ]; + for (rank, (tp_rank, dp_rank, tp_group, dp_group)) in expected.into_iter().enumerate() { + assert_eq!(topology.tensor_rank(rank).unwrap(), tp_rank); + assert_eq!(topology.data_rank(rank).unwrap(), dp_rank); + assert_eq!(topology.tensor_group(rank).unwrap(), tp_group); + assert_eq!(topology.data_group(rank).unwrap(), dp_group); + } + } + #[test] fn five_dimensional_default_order_matches_orthogonal_mapping() { let topology = ParallelTopology::new(2, 2, 2, 2, 2).unwrap(); diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 756bfad6..191015e7 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -30,6 +30,7 @@ #include #include #include +#include #include #include #include @@ -2032,6 +2033,149 @@ struct TrainingContext { // ────────────────────────────────────────────────────────────────────── }; +struct AdapterRegistryHash { + uint64_t first = 1469598103934665603ULL; + uint64_t second = 1099511628211ULL; + + void add_u64(uint64_t value) { + first ^= value; + first *= 1099511628211ULL; + second ^= value + 0x9e3779b97f4a7c15ULL + (second << 6) + + (second >> 2); + second *= 0xbf58476d1ce4e5b9ULL; + } + + void add_string(const std::string& value) { + add_u64(value.size()); + for (const unsigned char byte : value) add_u64(byte); + } +}; + +static void hash_tensor_layout( + AdapterRegistryHash& hash, const at::Tensor& tensor +) { + hash.add_u64(tensor.defined()); + if (!tensor.defined()) return; + hash.add_u64(static_cast(tensor.scalar_type())); + hash.add_u64(tensor.requires_grad()); + hash.add_u64(tensor.dim()); + for (const auto size : tensor.sizes()) hash.add_u64(size); +} + +static void hash_adapter_layout( + AdapterRegistryHash& hash, + const TrainingContext::LoRAAdapter& adapter +) { + hash.add_u64(adapter.id); + hash.add_u64(adapter.rank); + hash.add_u64(adapter.optimizer_step); + uint64_t alpha_bits = 0; + static_assert(sizeof(alpha_bits) == sizeof(adapter.alpha)); + std::memcpy(&alpha_bits, &adapter.alpha, sizeof(alpha_bits)); + hash.add_u64(alpha_bits); + hash.add_u64(adapter.target_layers.size()); + for (const auto layer : adapter.target_layers) hash.add_u64(layer); + hash.add_u64(adapter.target_modules.size()); + for (const auto& module : adapter.target_modules) hash.add_string(module); + hash.add_u64(adapter.params.size()); + for (const auto& [layer, pairs] : adapter.params) { + hash.add_u64(layer); + hash.add_u64(pairs.size()); + for (const auto& [a, b] : pairs) { + hash_tensor_layout(hash, a); + hash_tensor_layout(hash, b); + } + } +} + +static void validate_adapter_collective_registry( + TrainingContext* ctx, + const int64_t* requested_ids, + int64_t requested_count, + int64_t requested_rank, + bool use_registered_order +) { + if (!ctx || (!ctx->nccl_comm && !ctx->tp_comm)) return; + + AdapterRegistryHash hash; + hash.add_u64(requested_count); + hash.add_u64(requested_rank); + hash.add_u64(use_registered_order); + int64_t found_count = 0; + if (use_registered_order) { + hash.add_u64(ctx->adapters.size()); + for (const auto& adapter : ctx->adapters) { + hash_adapter_layout(hash, adapter); + ++found_count; + } + } else if (requested_ids && requested_count > 0) { + for (int64_t index = 0; index < requested_count; ++index) { + const int64_t requested_id = requested_ids[index]; + hash.add_u64(requested_id); + const auto it = std::find_if( + ctx->adapters.begin(), ctx->adapters.end(), + [&](const auto& adapter) { return adapter.id == requested_id; }); + if (it == ctx->adapters.end()) { + hash.add_u64(0x6d697373696e67ULL); + continue; + } + hash.add_u64(0x666f756e64ULL); + hash_adapter_layout(hash, *it); + ++found_count; + } + } else { + hash.add_u64(0x6e756c6cULL); + } + + constexpr uint64_t kPositiveInt64Mask = + static_cast(std::numeric_limits::max()); + const std::vector signature_values{ + requested_count, + requested_rank, + found_count, + static_cast(hash.first & kPositiveInt64Mask), + static_cast(hash.second & kPositiveInt64Mask), + }; + c10::cuda::set_device(ctx->cuda_device); + cudaSetDevice(ctx->cuda_device); + auto options = at::TensorOptions().dtype(at::kLong).device( + at::kCUDA, ctx->cuda_device); + auto minimum = at::tensor(signature_values, options); + auto maximum = minimum.clone(); + auto stream = c10::cuda::getCurrentCUDAStream(ctx->cuda_device).stream(); + + auto reduce_axis = [&](ncclComm_t communicator, const char* axis) { + if (!communicator) return; + auto err = ncclAllReduce( + minimum.data_ptr(), minimum.data_ptr(), + minimum.numel(), ncclInt64, ncclMin, communicator, stream); + TORCH_CHECK(err == ncclSuccess, axis, + " adapter registry minimum all-reduce failed: ", + ncclGetErrorString(err)); + err = ncclAllReduce( + maximum.data_ptr(), maximum.data_ptr(), + maximum.numel(), ncclInt64, ncclMax, communicator, stream); + TORCH_CHECK(err == ncclSuccess, axis, + " adapter registry maximum all-reduce failed: ", + ncclGetErrorString(err)); + }; + + // A reduction over one axis followed by the orthogonal axis propagates + // the extrema over the full TP x DP/EP grid. + reduce_axis(ctx->nccl_comm, "DP/EP"); + reduce_axis(ctx->tp_comm, "TP"); + const auto minimum_cpu = minimum.to(at::kCPU); + const auto maximum_cpu = maximum.to(at::kCPU); + const auto* minimum_data = minimum_cpu.data_ptr(); + const auto* maximum_data = maximum_cpu.data_ptr(); + for (int64_t index = 0; index < minimum.numel(); ++index) { + TORCH_CHECK(minimum_data[index] == maximum_data[index], + "dynamic LoRA adapter registry mismatch across distributed ranks; " + "all ranks must select the same ordered IDs, ranks, optimizer " + "clocks, targets, and tensor layouts"); + } +} + static bool harvest_leaf_grad( at::Tensor& param, at::Tensor& fp32_accumulator ) { @@ -2327,11 +2471,30 @@ static void tp_broadcast_lora_parameter( "NCCL LoRA parameter broadcast failed: ", ncclGetErrorString(err)); } +static void dp_broadcast_lora_parameter( + TrainingContext* ctx, at::Tensor& tensor +) { + if (!ctx || !ctx->data_parallel || ctx->ep_world_size <= 1 || + !tensor.defined()) return; + TORCH_CHECK(ctx->nccl_comm, + "LoRA DP communicator is not initialized for parameter broadcast"); + TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous(), + "LoRA DP parameter broadcast requires a contiguous CUDA tensor"); + const int dev = tensor.device().index(); + cudaSetDevice(dev); + auto stream = c10::cuda::getCurrentCUDAStream(dev).stream(); + auto err = ncclBroadcast( + tensor.data_ptr(), tensor.data_ptr(), tensor.numel(), + nccl_dtype_for(tensor), 0, ctx->nccl_comm, stream); + TORCH_CHECK(err == ncclSuccess, + "NCCL LoRA DP parameter broadcast failed: ", ncclGetErrorString(err)); +} + static void synchronize_adapter_replicated_lora_parameters( TrainingContext* ctx, TrainingContext::LoRAAdapter& adapter); static void synchronize_fixed_replicated_lora_parameters(TrainingContext* ctx) { - if (!ctx || !ctx->base_tp_attention || ctx->tp_world_size <= 1) return; + if (!ctx) return; for (int64_t layer = 0; layer < ctx->num_layers; ++layer) { const int64_t offset = ctx->lora_layer_offset[layer]; const int64_t pairs = lora_pair_count(ctx->layer_configs[layer]); @@ -2339,10 +2502,14 @@ static void synchronize_fixed_replicated_lora_parameters(TrainingContext* ctx) { const int64_t slot = offset + pair; if (!legacy_lora_slot_active(ctx, slot)) continue; const auto layout = lora_tp_layout(ctx, layer, pair); - if (layout == LoraTpLayout::ColumnParallel) + if (ctx->base_tp_attention && ctx->tp_world_size > 1 && + layout == LoraTpLayout::ColumnParallel) tp_broadcast_lora_parameter(ctx, ctx->lora_a[slot]); - else if (layout == LoraTpLayout::RowParallel) + else if (ctx->base_tp_attention && ctx->tp_world_size > 1 && + layout == LoraTpLayout::RowParallel) tp_broadcast_lora_parameter(ctx, ctx->lora_b[slot]); + dp_broadcast_lora_parameter(ctx, ctx->lora_a[slot]); + dp_broadcast_lora_parameter(ctx, ctx->lora_b[slot]); } } for (auto& adapter : ctx->adapters) @@ -2352,17 +2519,24 @@ static void synchronize_fixed_replicated_lora_parameters(TrainingContext* ctx) { static void synchronize_adapter_replicated_lora_parameters( TrainingContext* ctx, TrainingContext::LoRAAdapter& adapter ) { - if (!ctx || !ctx->base_tp_attention || ctx->tp_world_size <= 1) return; - if (!ctx->tp_comm) return; // qwen36_init_nccl synchronizes deferred adapters. + if (!ctx) return; + if (ctx->base_tp_attention && ctx->tp_world_size > 1 && !ctx->tp_comm) + return; // qwen36_init_nccl synchronizes deferred adapters. + if (ctx->data_parallel && ctx->ep_world_size > 1 && !ctx->nccl_comm) + return; for (auto& [layer, pairs] : adapter.params) { for (int64_t pair = 0; pair < static_cast(pairs.size()); ++pair) { auto& [a, b] = pairs[pair]; if (!a.requires_grad() && !b.requires_grad()) continue; const auto layout = lora_tp_layout(ctx, layer, pair); - if (layout == LoraTpLayout::ColumnParallel) + if (ctx->base_tp_attention && ctx->tp_world_size > 1 && + layout == LoraTpLayout::ColumnParallel) tp_broadcast_lora_parameter(ctx, a); - else if (layout == LoraTpLayout::RowParallel) + else if (ctx->base_tp_attention && ctx->tp_world_size > 1 && + layout == LoraTpLayout::RowParallel) tp_broadcast_lora_parameter(ctx, b); + dp_broadcast_lora_parameter(ctx, a); + dp_broadcast_lora_parameter(ctx, b); } } } @@ -4590,10 +4764,11 @@ static at::Tensor mtp_compute_loss( extern "C" { __attribute__((visibility("default"))) int64_t qwen36_kernel_abi_version() { - return 15; + return 16; } static constexpr int32_t QWEN36_CONTEXT_BASE_TP_ATTENTION = 1 << 0; +static constexpr int32_t QWEN36_CONTEXT_DATA_PARALLEL = 1 << 1; // Create training context — called once at startup // lora_rank: LoRA rank (from config) @@ -4630,29 +4805,50 @@ static void* qwen36_create_training_context_impl( const char* world_size_env = getenv("WORLD_SIZE"); const int configured_world_size = world_size_env ? atoi(world_size_env) : 1; TORCH_CHECK(configured_world_size > 0, "WORLD_SIZE must be positive"); - const bool data_parallel_requested = env_enabled("RUSTRAIN_DATA_PARALLEL"); + const bool data_parallel_requested = + (context_flags & QWEN36_CONTEXT_DATA_PARALLEL) != 0 || + env_enabled("RUSTRAIN_DATA_PARALLEL"); + const char* dp_size_env = getenv("DP_SIZE"); + if (!dp_size_env) dp_size_env = getenv("RUSTRAIN_DP_SIZE"); + int configured_dp_size = dp_size_env ? atoi(dp_size_env) : 1; + if (data_parallel_requested && !dp_size_env && + configured_world_size % ctx->tp_world_size == 0) { + configured_dp_size = configured_world_size / ctx->tp_world_size; + } + const char* ep_size_env = getenv("EP_SIZE"); + if (!ep_size_env) ep_size_env = getenv("RUSTRAIN_EP_SIZE"); + const int configured_ep_size = ep_size_env ? atoi(ep_size_env) : 1; const char* pp_size_env = getenv("PP_SIZE"); if (!pp_size_env) pp_size_env = getenv("RUSTRAIN_PP_SIZE"); const char* cp_size_env = getenv("CP_SIZE"); if (!cp_size_env) cp_size_env = getenv("RUSTRAIN_CP_SIZE"); const int configured_pp_size = pp_size_env ? atoi(pp_size_env) : 1; const int configured_cp_size = cp_size_env ? atoi(cp_size_env) : 1; - TORCH_CHECK(configured_pp_size > 0 && configured_cp_size > 0, - "PP_SIZE and CP_SIZE must be positive"); - TORCH_CHECK(configured_pp_size == 1 && configured_cp_size == 1, - "native Qwen LoRA does not implement PP/CP yet; ", + TORCH_CHECK(configured_pp_size > 0 && configured_cp_size > 0 && + configured_dp_size > 0 && configured_ep_size > 0, + "PP_SIZE, CP_SIZE, DP_SIZE, and EP_SIZE must be positive"); + TORCH_CHECK(configured_pp_size == 1 && configured_cp_size == 1 && + (!data_parallel_requested || configured_ep_size == 1), + "native Qwen LoRA does not implement PP/CP or mixed DP/EP yet; ", "PP_SIZE=", configured_pp_size, " CP_SIZE=", configured_cp_size); - TORCH_CHECK( - ctx->tp_world_size <= 1 || - (!data_parallel_requested && configured_world_size == ctx->tp_world_size), - "native Qwen LoRA supports TP-only topology when TP_SIZE>1; " - "TP_SIZE=", ctx->tp_world_size, " WORLD_SIZE=", configured_world_size, - " DATA_PARALLEL=", data_parallel_requested ? 1 : 0, - " is an incompatible mixed TP/DP/EP topology"); - ctx->ep_world_size = configured_world_size; + if (data_parallel_requested) { + TORCH_CHECK( + configured_world_size == ctx->tp_world_size * configured_dp_size, + "native Qwen LoRA requires WORLD_SIZE=TP_SIZE*DP_SIZE for dense DP; ", + "TP_SIZE=", ctx->tp_world_size, " DP_SIZE=", configured_dp_size, + " WORLD_SIZE=", configured_world_size); + } else if (ctx->tp_world_size > 1) { + TORCH_CHECK(configured_world_size == ctx->tp_world_size && + configured_ep_size == 1, + "native Qwen LoRA TP without DP requires WORLD_SIZE=TP_SIZE and EP_SIZE=1"); + } + ctx->ep_world_size = data_parallel_requested + ? configured_dp_size : configured_world_size; ctx->data_parallel = data_parallel_requested; const char* rank_env = getenv("RANK"); const int global_rank = rank_env ? atoi(rank_env) : 0; + ctx->ep_rank = data_parallel_requested + ? global_rank / ctx->tp_world_size : global_rank; ctx->tp_rank = global_rank % ctx->tp_world_size; TORCH_CHECK(lora_rank > 0 && lora_rank % ctx->tp_world_size == 0, "LoRA rank ", lora_rank, " must be divisible by TP_SIZE=", @@ -5094,7 +5290,7 @@ __attribute__((visibility("default"))) double qwen36_train_micro_step( TORCH_CHECK(gradient_scale > 0.0 && std::isfinite(gradient_scale), "gradient_scale must be finite and positive"); // Set CUDA device for EP - if (ctx->nccl_comm) { + if (ctx->nccl_comm || ctx->tp_comm) { c10::cuda::set_device(ctx->cuda_device); cudaSetDevice(ctx->cuda_device); } @@ -5487,6 +5683,8 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( auto& target_mask = *reinterpret_cast(target_mask_ptr); int64_t total_adapters = (int64_t)ctx->adapters.size(); + validate_adapter_collective_registry( + ctx, nullptr, n_total, lora_rank, true); if (total_adapters == 0) return -1.0; TORCH_CHECK(n_total > 0 && total_adapters == n_total, "n_total must equal the number of registered adapters (n_total=", @@ -5564,40 +5762,42 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora( ~AdapterRegistryChunkGuard() { restore(); } }; - // Compute N_max from available GPU memory. All workers must agree on - // the chunk schedule to keep later collectives in the same order, so - // rank 0 publishes its value through the existing communicator. + // Compute N_max from available GPU memory. All workers in the TP x DP + // grid must agree on the chunk schedule to keep collectives ordered. size_t free_mem, total_mem; cudaMemGetInfo(&free_mem, &total_mem); int64_t n_max = 0; - if (ctx->nccl_comm && ctx->ep_world_size > 1) { - if (ctx->ep_rank == 0) { - n_max = compute_n_max( - (int64_t)free_mem, lora_rank, - input_ids.size(-1), 2048, - ctx->group_size, ctx->num_layers - ); - n_max = std::min(n_max, total_adapters); - if (n_max < 1) n_max = 1; - } + n_max = compute_n_max( + (int64_t)free_mem, lora_rank, + input_ids.size(-1), 2048, + ctx->group_size, ctx->num_layers + ); + n_max = std::min(n_max, total_adapters); + if (n_max < 1) n_max = 1; + if ((ctx->nccl_comm && ctx->ep_world_size > 1) || + (ctx->tp_comm && ctx->tp_world_size > 1)) { auto published_n_max = at::full( {1}, n_max, input_ids.options().dtype(at::kLong)); auto stream = c10::cuda::getCurrentCUDAStream( input_ids.device().index()).stream(); - auto err = ncclBroadcast( - published_n_max.data_ptr(), - published_n_max.data_ptr(), 1, ncclInt64, 0, - reinterpret_cast(ctx->nccl_comm), stream); - TORCH_CHECK(err == ncclSuccess, "n_max broadcast failed: ", - ncclGetErrorString(err)); + if (ctx->nccl_comm && ctx->ep_world_size > 1) { + auto err = ncclAllReduce( + published_n_max.data_ptr(), + published_n_max.data_ptr(), 1, ncclInt64, ncclMin, + reinterpret_cast(ctx->nccl_comm), stream); + TORCH_CHECK(err == ncclSuccess, "DP/EP n_max all-reduce failed: ", + ncclGetErrorString(err)); + } + if (ctx->tp_comm && ctx->tp_world_size > 1) { + auto err = ncclAllReduce( + published_n_max.data_ptr(), + published_n_max.data_ptr(), 1, ncclInt64, ncclMin, + reinterpret_cast(ctx->tp_comm), stream); + TORCH_CHECK(err == ncclSuccess, "TP n_max all-reduce failed: ", + ncclGetErrorString(err)); + } n_max = published_n_max.to( at::TensorOptions().device(at::kCPU)).item(); - } else { - n_max = compute_n_max( - (int64_t)free_mem, lora_rank, - input_ids.size(-1), 2048, - ctx->group_size, ctx->num_layers - ); } n_max = std::min(n_max, total_adapters); if (n_max < 1) n_max = 1; @@ -5895,7 +6095,13 @@ __attribute__((visibility("default"))) double qwen36_train_multi_lora_selected( registry_detached = false; }; try { - TORCH_CHECK(ctx && adapter_ids && n_adapters > 0, + TORCH_CHECK(ctx, + "selected multi-LoRA requires a valid training context"); + TORCH_CHECK(!ctx->topology_invalid, + "native Qwen context rejected an incompatible TP/DP/EP topology"); + validate_adapter_collective_registry( + ctx, adapter_ids, n_adapters, lora_rank, false); + TORCH_CHECK(adapter_ids && n_adapters > 0, "selected multi-LoRA requires at least one adapter ID"); const auto original_count = ctx->adapters.size(); original.reserve(original_count); @@ -5944,7 +6150,46 @@ static ncclComm_t g_nccl_comm = nullptr; static cudaStream_t g_nccl_stream = nullptr; static ncclComm_t g_tp_comm = nullptr; static cudaStream_t g_tp_stream = nullptr; +static ncclComm_t g_dp_comm = nullptr; +static cudaStream_t g_dp_stream = nullptr; static bool g_nccl_initialized = false; +static bool g_nccl_cleanup_registered = false; +static int g_parallel_rank = 0; +static int g_parallel_world_size = 1; +static int g_parallel_tp_rank = 0; +static int g_parallel_tp_size = 1; +static int g_parallel_tp_color = 0; +static int g_parallel_dp_rank = 0; +static int g_parallel_dp_size = 1; +static int g_parallel_dp_color = 0; +static bool g_parallel_data_parallel = false; + +static void qwen36_destroy_process_communicators() { + if (g_dp_comm) ncclCommDestroy(g_dp_comm); + if (g_tp_comm) ncclCommDestroy(g_tp_comm); + if (g_nccl_comm) ncclCommDestroy(g_nccl_comm); + g_dp_comm = nullptr; + g_tp_comm = nullptr; + g_nccl_comm = nullptr; + g_nccl_initialized = false; +} + +static bool same_cached_parallel_topology( + int rank, int world_size, + int tp_rank, int tp_size, int tp_color, + int dp_rank, int dp_size, int dp_color, + bool data_parallel +) { + return g_parallel_rank == rank && + g_parallel_world_size == world_size && + g_parallel_tp_rank == tp_rank && + g_parallel_tp_size == tp_size && + g_parallel_tp_color == tp_color && + g_parallel_dp_rank == dp_rank && + g_parallel_dp_size == dp_size && + g_parallel_dp_color == dp_color && + g_parallel_data_parallel == data_parallel; +} static int g_cuda_device = 0; // Set CUDA device — called from Rust worker before any GPU operation. @@ -5962,80 +6207,68 @@ __attribute__((visibility("default"))) void qwen36_set_cuda_device(int32_t devic dummy.sizes(); // touch to ensure materialization } -__attribute__((visibility("default"))) int32_t qwen36_init_nccl( - void* ctx_ptr +static int32_t qwen36_init_parallel_nccl_impl( + void* ctx_ptr, + int rank, int world_size, + int tp_rank, int tp_size, int tp_color, + int dp_rank, int dp_size, int dp_color, + bool data_parallel ) { auto* ctx = reinterpret_cast(ctx_ptr); + if (rank < 0 || rank >= world_size || world_size <= 0 || + tp_rank < 0 || tp_rank >= tp_size || tp_size <= 0 || tp_color < 0 || + dp_rank < 0 || dp_rank >= dp_size || dp_size <= 0 || dp_color < 0 || + (data_parallel && world_size != tp_size * dp_size) || + (!data_parallel && tp_size > 1 && world_size != tp_size)) { + ctx->topology_invalid = true; + fprintf(stderr, + "[parallel_nccl] invalid topology: rank=%d world=%d tp=%d/%d color=%d " + "dp=%d/%d color=%d data_parallel=%d\n", + rank, world_size, tp_rank, tp_size, tp_color, + dp_rank, dp_size, dp_color, data_parallel ? 1 : 0); + return -1; + } + // If already initialized, just set the pointer on this context if (g_nccl_initialized) { - ctx->nccl_comm = g_nccl_comm; - ctx->nccl_stream = g_nccl_stream; - ctx->data_parallel = env_enabled("RUSTRAIN_DATA_PARALLEL"); - // CRITICAL: also set ep_rank/ep_world_size — needed for new TrainingContext - // created by subsequent CreateSession commands. The fast path previously - // skipped this, leaving ep_rank=0 → cudaSetDevice(0) on all ranks → crash. - const char* rank_str2 = getenv("RANK"); - const char* world_str2 = getenv("WORLD_SIZE"); - if (rank_str2) ctx->ep_rank = atoi(rank_str2); - if (world_str2) ctx->ep_world_size = atoi(world_str2); - const char* local_rank_str2 = getenv("LOCAL_RANK"); - ctx->cuda_device = local_rank_str2 ? atoi(local_rank_str2) : g_cuda_device; - const char* tp_size_str2 = getenv("TP_SIZE"); - if (!tp_size_str2) tp_size_str2 = getenv("RUSTRAIN_TP_SIZE"); - ctx->tp_world_size = tp_size_str2 ? atoi(tp_size_str2) : 1; - ctx->tp_rank = ctx->tp_world_size > 0 - ? ctx->ep_rank % ctx->tp_world_size : 0; - if (ctx->tp_world_size <= 0 || - (ctx->tp_world_size > 1 && - (ctx->data_parallel || ctx->ep_world_size != ctx->tp_world_size))) { + if (!same_cached_parallel_topology( + rank, world_size, tp_rank, tp_size, tp_color, + dp_rank, dp_size, dp_color, data_parallel)) { ctx->topology_invalid = true; fprintf(stderr, - "[tp_nccl] reject mixed topology: TP_SIZE=%d WORLD_SIZE=%d DATA_PARALLEL=%d\n", - ctx->tp_world_size, ctx->ep_world_size, ctx->data_parallel ? 1 : 0); + "[parallel_nccl] process communicator topology cannot change after initialization\n"); return -1; } + ctx->data_parallel = data_parallel; + ctx->ep_rank = data_parallel ? dp_rank : rank; + ctx->ep_world_size = data_parallel ? dp_size : world_size; + const char* local_rank_str2 = getenv("LOCAL_RANK"); + ctx->cuda_device = local_rank_str2 ? atoi(local_rank_str2) : g_cuda_device; + ctx->tp_world_size = tp_size; + ctx->tp_rank = tp_rank; ctx->topology_invalid = false; - ctx->tp_comm = ctx->tp_world_size > 1 ? g_tp_comm : nullptr; - ctx->tp_stream = ctx->tp_world_size > 1 ? g_tp_stream : nullptr; - // In TP-only mode the parent communicator is reserved for the TP - // split; EP layer collectives must remain disabled on replicated MoE. - if (ctx->tp_world_size <= 1) { - void* layer_comm = ctx->data_parallel - ? nullptr : (void*)g_nccl_comm; - void* layer_stream = ctx->data_parallel - ? nullptr : (void*)g_nccl_stream; - for (auto& lc : ctx->layer_configs) { - lc.nccl_comm = layer_comm; - lc.nccl_stream = layer_stream; - } - for (auto& lc : ctx->mtp_layer_configs) { - lc.nccl_comm = layer_comm; - lc.nccl_stream = layer_stream; - } + ctx->tp_comm = tp_size > 1 ? g_tp_comm : nullptr; + ctx->tp_stream = tp_size > 1 ? g_tp_stream : nullptr; + ctx->nccl_comm = data_parallel + ? (dp_size > 1 ? g_dp_comm : nullptr) + : (tp_size > 1 ? nullptr : g_nccl_comm); + ctx->nccl_stream = data_parallel ? g_dp_stream : g_nccl_stream; + void* layer_comm = !data_parallel && tp_size <= 1 + ? (void*)g_nccl_comm : nullptr; + void* layer_stream = !data_parallel && tp_size <= 1 + ? (void*)g_nccl_stream : nullptr; + for (auto& lc : ctx->layer_configs) { + lc.nccl_comm = layer_comm; + lc.nccl_stream = layer_stream; + } + for (auto& lc : ctx->mtp_layer_configs) { + lc.nccl_comm = layer_comm; + lc.nccl_stream = layer_stream; } synchronize_fixed_replicated_lora_parameters(ctx); return 0; } - - const char* rank_str = getenv("RANK"); - const char* world_str = getenv("WORLD_SIZE"); - if (!rank_str || !world_str) return -1; - int rank = atoi(rank_str); - int world_size = atoi(world_str); - const char* tp_size_str = getenv("TP_SIZE"); - if (!tp_size_str) tp_size_str = getenv("RUSTRAIN_TP_SIZE"); - const int configured_tp_size = tp_size_str ? atoi(tp_size_str) : 1; - const bool data_parallel_requested = env_enabled("RUSTRAIN_DATA_PARALLEL"); - if (configured_tp_size <= 0 || - (configured_tp_size > 1 && - (data_parallel_requested || world_size != configured_tp_size))) { - ctx->topology_invalid = true; - fprintf(stderr, - "[tp_nccl] reject mixed topology: TP_SIZE=%d WORLD_SIZE=%d DATA_PARALLEL=%d\n", - configured_tp_size, world_size, data_parallel_requested ? 1 : 0); - return -1; - } if (world_size <= 1) return 0; // no EP needed // Set CUDA device and initialize PyTorch CUDA context on this device. @@ -6134,56 +6367,120 @@ __attribute__((visibility("default"))) int32_t qwen36_init_nccl( // Store as process-level singleton g_nccl_comm = comm; g_nccl_stream = nccl_stream; - const int tp_size = tp_size_str ? atoi(tp_size_str) : 1; if (tp_size <= 0 || world_size % tp_size != 0) { fprintf(stderr, "[tp_nccl] invalid TP_SIZE=%d for WORLD_SIZE=%d\n", tp_size, world_size); return -1; } if (tp_size > 1) { - // The default rank order makes TP the least-significant axis. A - // future multi-axis implementation must pass an explicit color/key - // mapping instead of reusing this world-rank split. ncclResult_t tp_err = ncclCommSplit( - comm, rank / tp_size, rank % tp_size, &g_tp_comm, nullptr); + comm, tp_color, tp_rank, &g_tp_comm, nullptr); if (tp_err != ncclSuccess) { fprintf(stderr, "[tp_nccl] ncclCommSplit failed: %d (%s)\n", tp_err, ncclGetErrorString(tp_err)); + ncclCommDestroy(comm); + g_nccl_comm = nullptr; return -1; } g_tp_stream = nccl_stream; } + if (data_parallel && dp_size > 1) { + ncclResult_t dp_err = ncclCommSplit( + comm, dp_color, dp_rank, &g_dp_comm, nullptr); + if (dp_err != ncclSuccess) { + fprintf(stderr, "[dp_nccl] ncclCommSplit failed: %d (%s)\n", + dp_err, ncclGetErrorString(dp_err)); + if (g_tp_comm) ncclCommDestroy(g_tp_comm); + ncclCommDestroy(comm); + g_tp_comm = nullptr; + g_nccl_comm = nullptr; + return -1; + } + g_dp_stream = nccl_stream; + } + g_parallel_rank = rank; + g_parallel_world_size = world_size; + g_parallel_tp_rank = tp_rank; + g_parallel_tp_size = tp_size; + g_parallel_tp_color = tp_color; + g_parallel_dp_rank = dp_rank; + g_parallel_dp_size = dp_size; + g_parallel_dp_color = dp_color; + g_parallel_data_parallel = data_parallel; g_nccl_initialized = true; - - ctx->nccl_comm = comm; - ctx->nccl_stream = nccl_stream; - ctx->ep_rank = rank; - ctx->ep_world_size = world_size; - ctx->data_parallel = env_enabled("RUSTRAIN_DATA_PARALLEL"); + if (!g_nccl_cleanup_registered) { + std::atexit(qwen36_destroy_process_communicators); + g_nccl_cleanup_registered = true; + } + + ctx->nccl_comm = data_parallel + ? (dp_size > 1 ? g_dp_comm : nullptr) + : (tp_size > 1 ? nullptr : comm); + ctx->nccl_stream = data_parallel ? g_dp_stream : nccl_stream; + ctx->ep_rank = data_parallel ? dp_rank : rank; + ctx->ep_world_size = data_parallel ? dp_size : world_size; + ctx->data_parallel = data_parallel; ctx->tp_world_size = tp_size; - ctx->tp_rank = rank % tp_size; + ctx->tp_rank = tp_rank; ctx->tp_comm = tp_size > 1 ? g_tp_comm : nullptr; ctx->tp_stream = tp_size > 1 ? g_tp_stream : nullptr; // Propagate to layer configs - if (tp_size <= 1) { - void* layer_comm = ctx->data_parallel ? nullptr : (void*)comm; - void* layer_stream = ctx->data_parallel - ? nullptr : (void*)nccl_stream; - for (auto& lc : ctx->layer_configs) { - lc.nccl_comm = layer_comm; - lc.nccl_stream = layer_stream; - } - for (auto& lc : ctx->mtp_layer_configs) { - lc.nccl_comm = layer_comm; - lc.nccl_stream = layer_stream; - } + void* layer_comm = !data_parallel && tp_size <= 1 + ? (void*)comm : nullptr; + void* layer_stream = !data_parallel && tp_size <= 1 + ? (void*)nccl_stream : nullptr; + for (auto& lc : ctx->layer_configs) { + lc.nccl_comm = layer_comm; + lc.nccl_stream = layer_stream; + } + for (auto& lc : ctx->mtp_layer_configs) { + lc.nccl_comm = layer_comm; + lc.nccl_stream = layer_stream; } synchronize_fixed_replicated_lora_parameters(ctx); return 0; } +__attribute__((visibility("default"))) int32_t qwen36_init_parallel_nccl( + void* ctx_ptr, + int32_t rank, int32_t world_size, + int32_t tp_rank, int32_t tp_size, int32_t tp_color, + int32_t dp_rank, int32_t dp_size, int32_t dp_color +) { + return qwen36_init_parallel_nccl_impl( + ctx_ptr, rank, world_size, + tp_rank, tp_size, tp_color, + dp_rank, dp_size, dp_color, + dp_size > 1); +} + +__attribute__((visibility("default"))) int32_t qwen36_init_nccl( + void* ctx_ptr +) { + const char* rank_str = getenv("RANK"); + const char* world_str = getenv("WORLD_SIZE"); + if (!rank_str || !world_str) return -1; + const int rank = atoi(rank_str); + const int world_size = atoi(world_str); + const char* tp_size_str = getenv("TP_SIZE"); + if (!tp_size_str) tp_size_str = getenv("RUSTRAIN_TP_SIZE"); + const int tp_size = tp_size_str ? atoi(tp_size_str) : 1; + const bool data_parallel = env_enabled("RUSTRAIN_DATA_PARALLEL"); + const char* dp_size_str = getenv("DP_SIZE"); + if (!dp_size_str) dp_size_str = getenv("RUSTRAIN_DP_SIZE"); + const int dp_size = dp_size_str + ? atoi(dp_size_str) + : (data_parallel && tp_size > 0 ? world_size / tp_size : 1); + const int tp_rank = tp_size > 0 ? rank % tp_size : 0; + const int dp_rank = data_parallel && tp_size > 0 ? rank / tp_size : 0; + return qwen36_init_parallel_nccl_impl( + ctx_ptr, rank, world_size, + tp_rank, tp_size, rank / std::max(tp_size, 1), + dp_rank, dp_size, tp_rank, data_parallel); +} + // Set NCCL communicator for Expert Parallel all-reduce (legacy, from Rust) __attribute__((visibility("default"))) void qwen36_set_nccl_comm( void* ctx_ptr, void* comm_ptr, void* stream_ptr, diff --git a/crates/rustrain-qwen3-6/src/kernel.rs b/crates/rustrain-qwen3-6/src/kernel.rs index 7c000d45..72ac2a84 100644 --- a/crates/rustrain-qwen3-6/src/kernel.rs +++ b/crates/rustrain-qwen3-6/src/kernel.rs @@ -4,7 +4,7 @@ //! Rust only handles: weight loading, data loading, training loop orchestration. use crate::lora::Qwen36LoraTargetModule; -use anyhow::{bail, Result}; +use anyhow::{bail, Context, Result}; use std::ffi::c_void; use std::sync::OnceLock; use tch::{Kind, Tensor}; @@ -78,6 +78,8 @@ type FnSetMtpWeights = unsafe extern "C" fn( type FnSetCheckpoint = unsafe extern "C" fn(*mut c_void, i32, i64); type FnSetNcclComm = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut c_void, i32, i32); type FnInitNccl = unsafe extern "C" fn(*mut c_void) -> i32; +type FnInitParallelNccl = + unsafe extern "C" fn(*mut c_void, i32, i32, i32, i32, i32, i32, i32, i32) -> i32; type FnSetCudaDevice = unsafe extern "C" fn(i32); type FnSetBaseTpMlp = unsafe extern "C" fn(*mut c_void, i32) -> i32; type FnAddLora = unsafe extern "C" fn(*mut c_void, i64, f64, *const i64, i64, *const i8) -> i64; @@ -145,6 +147,7 @@ struct KernelHandles { set_checkpoint: FnSetCheckpoint, set_nccl_comm: FnSetNcclComm, init_nccl: FnInitNccl, + init_parallel_nccl: FnInitParallelNccl, set_cuda_device: FnSetCudaDevice, set_base_tp_mlp: FnSetBaseTpMlp, add_lora: FnAddLora, @@ -197,7 +200,7 @@ unsafe fn load_kernels() -> Option { }}; } let abi_version: FnKernelAbiVersion = sym!("qwen36_kernel_abi_version"); - if abi_version() != 15 { + if abi_version() != 16 { return None; } Some(KernelHandles { @@ -224,6 +227,7 @@ unsafe fn load_kernels() -> Option { set_checkpoint: sym!("qwen36_set_checkpoint"), set_nccl_comm: sym!("qwen36_set_nccl_comm"), init_nccl: sym!("qwen36_init_nccl"), + init_parallel_nccl: sym!("qwen36_init_parallel_nccl"), set_cuda_device: sym!("qwen36_set_cuda_device"), set_base_tp_mlp: sym!("qwen36_set_base_tp_mlp"), add_lora: sym!("qwen36_add_lora"), @@ -635,6 +639,7 @@ impl CppTrainingContext { lora_rank: i64, base_tp_attention: bool, base_tp_mlp: bool, + data_parallel: bool, target_layers: &[usize], target_modules: &[Qwen36LoraTargetModule], expert_start: usize, @@ -710,7 +715,7 @@ impl CppTrainingContext { tl_ptr, tl_len, modules_ptr, - i32::from(base_tp_attention), + i32::from(base_tp_attention) | (i32::from(data_parallel) << 1), ) }; if ptr.is_null() { @@ -976,6 +981,44 @@ impl CppTrainingContext { unsafe { (kh.init_nccl)(self.ptr) } } + /// Initialize orthogonal TP and DP process groups from validated topology metadata. + pub fn init_parallel_nccl( + &self, + rank: usize, + world_size: usize, + tp_rank: usize, + tp_size: usize, + tp_color: usize, + dp_rank: usize, + dp_size: usize, + dp_color: usize, + ) -> Result<()> { + let values = [ + rank, world_size, tp_rank, tp_size, tp_color, dp_rank, dp_size, dp_color, + ] + .map(|value| i32::try_from(value).context("parallel topology exceeds i32")); + let [rank, world_size, tp_rank, tp_size, tp_color, dp_rank, dp_size, dp_color] = + values; + let kh = get_kernels().expect("kernels not loaded"); + let status = unsafe { + (kh.init_parallel_nccl)( + self.ptr, + rank?, + world_size?, + tp_rank?, + tp_size?, + tp_color?, + dp_rank?, + dp_size?, + dp_color?, + ) + }; + if status != 0 { + bail!("C++ parallel NCCL init failed (code {status})"); + } + Ok(()) + } + /// Set CUDA device and force PyTorch CUDA context initialization. /// Must be called before any GPU operation in worker processes. pub fn set_cuda_device(device: i32) { diff --git a/crates/rustrain-qwen3-6/src/session.rs b/crates/rustrain-qwen3-6/src/session.rs index 50d725c9..b9c3fad5 100644 --- a/crates/rustrain-qwen3-6/src/session.rs +++ b/crates/rustrain-qwen3-6/src/session.rs @@ -16,6 +16,7 @@ use crate::lora::{ use crate::sft::SftDataset; use rustrain_checkpoint::safetensors::read_safetensors_dir_filtered; use rustrain_core::runtime::{Config, RunPaths}; +use rustrain_parallel::topology::{DEFAULT_RANK_ORDER, ParallelTopology}; // ────────────────────────────────────────────────────────────────────── // EP Shard @@ -278,6 +279,7 @@ fn train_impl( let tp_size = config.parallel.tensor_model_parallel_size; let env_tp_size = std::env::var("TP_SIZE") .or_else(|_| std::env::var("RUSTRAIN_TP_SIZE")) + .or_else(|_| std::env::var("TENSOR_MODEL_PARALLEL_SIZE")) .ok() .and_then(|s| s.parse::().ok()) .unwrap_or(1); @@ -286,39 +288,87 @@ fn train_impl( "TP_SIZE environment ({env_tp_size}) does not match config tensor_model_parallel_size ({tp_size})" ); } - if tp_size > 1 { - if world_size != tp_size - || config.parallel.pipeline_model_parallel_size != 1 - || config.parallel.data_parallel_size != 1 + let configured_dp_size = std::env::var("DP_SIZE") + .or_else(|_| std::env::var("RUSTRAIN_DP_SIZE")) + .or_else(|_| std::env::var("DATA_PARALLEL_SIZE")) + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(config.parallel.data_parallel_size); + let dp_size = if !is_ep && configured_dp_size == 1 && world_size > tp_size { + world_size + .checked_div(tp_size) + .filter(|value| value * tp_size == world_size) + .ok_or_else(|| { + anyhow!("WORLD_SIZE={world_size} is not divisible by TP_SIZE={tp_size}") + })? + } else { + configured_dp_size + }; + let rank_order = std::env::var("RUSTRAIN_PARALLEL_ORDER") + .or_else(|_| std::env::var("PARALLEL_ORDER")) + .unwrap_or_else(|_| DEFAULT_RANK_ORDER.to_string()); + let dense_topology = if is_ep { + None + } else { + if config.parallel.pipeline_model_parallel_size != 1 || config.parallel.expert_model_parallel_size != 1 || config.parallel.context_parallel_size != 1 { bail!( - "native Qwen LoRA currently supports TP-only topology: TP={} WORLD_SIZE={} PP={} DP={} EP={} CP={}", + "native dense Qwen LoRA supports TPxDP only: TP={} WORLD_SIZE={} PP={} DP={} EP={} CP={}", tp_size, world_size, config.parallel.pipeline_model_parallel_size, - config.parallel.data_parallel_size, + dp_size, config.parallel.expert_model_parallel_size, config.parallel.context_parallel_size ); } + Some(ParallelTopology::with_order( + tp_size, + 1, + dp_size, + 1, + 1, + &rank_order, + )?) + }; + if let Some(topology) = dense_topology.as_ref() { + topology.validate_world_size(world_size)?; + topology.coordinates(rank)?; + } + if tp_size > 1 { if lora_config.rank % tp_size as i64 != 0 { bail!( "LoRA rank {} must be divisible by TP_SIZE={tp_size}", lora_config.rank ); } - unsafe { - std::env::set_var("TP_SIZE", tp_size.to_string()); - } } - let is_data_parallel = !is_ep && world_size > 1 && tp_size == 1; + let is_data_parallel = !is_ep && dp_size > 1; if is_data_parallel && runtime_config.is_moe { bail!( "replicated Qwen data parallelism is only supported for dense/linear-attention models; use *_ep for MoE" ); } + unsafe { + std::env::set_var("TP_SIZE", tp_size.to_string()); + std::env::set_var("DP_SIZE", dp_size.to_string()); + std::env::set_var( + "RUSTRAIN_DATA_PARALLEL", + if is_data_parallel { "1" } else { "0" }, + ); + } + let tp_rank = dense_topology + .as_ref() + .map(|topology| topology.tensor_rank(rank)) + .transpose()? + .unwrap_or(0); + let dp_rank = dense_topology + .as_ref() + .map(|topology| topology.data_rank(rank)) + .transpose()? + .unwrap_or(rank); if is_data_parallel || tp_size > 1 { crate::kernel::CppTrainingContext::set_cuda_device( std::env::var("LOCAL_RANK") @@ -466,10 +516,7 @@ fn train_impl( for (name, tensor) in &weights { let local_shard = if base_tp_attention { let full_attention_shard = crate::kernel::shard_full_attention_weight_for_tp( - name, - tensor, - tp_size, - rank % tp_size, + name, tensor, tp_size, tp_rank, )?; let attention_shard = if full_attention_shard.is_some() { full_attention_shard @@ -478,7 +525,7 @@ fn train_impl( name, tensor, tp_size, - rank % tp_size, + tp_rank, runtime_config.linear_num_key_heads, runtime_config.linear_key_head_dim, runtime_config.linear_num_value_heads, @@ -488,12 +535,7 @@ fn train_impl( if attention_shard.is_some() || !base_tp_mlp { attention_shard } else { - crate::kernel::shard_dense_mlp_weight_for_tp( - name, - tensor, - tp_size, - rank % tp_size, - )? + crate::kernel::shard_dense_mlp_weight_for_tp(name, tensor, tp_size, tp_rank)? } } else { None @@ -508,7 +550,7 @@ fn train_impl( if base_tp_attention { info!( tp_size, - tp_rank = rank % tp_size, + tp_rank, base_tp_mlp, "frozen base TP enabled: full-attention/GDN head shards and optional dense MLP shards" ); @@ -571,6 +613,7 @@ fn train_impl( lora_config.rank as i64, base_tp_attention, base_tp_mlp, + is_data_parallel, &lora_config.target_layers, &lora_config.target_modules, expert_start, @@ -578,18 +621,25 @@ fn train_impl( )?; if world_size > 1 { - // Tell the native reducer whether this communicator is replicated DP - // or expert-parallel. EP already all-reduces routed activations and - // must not average replicated LoRA gradients a second time. - unsafe { - std::env::set_var( - "RUSTRAIN_DATA_PARALLEL", - if is_data_parallel { "1" } else { "0" }, - ); - } - let ret = ctx.init_nccl(); - if ret != 0 { - bail!("C++ NCCL init failed (code {})", ret); + if let Some(topology) = dense_topology.as_ref() { + let tp_color = *topology + .tensor_group(rank)? + .iter() + .min() + .ok_or_else(|| anyhow!("empty TP process group"))?; + let dp_color = *topology + .data_group(rank)? + .iter() + .min() + .ok_or_else(|| anyhow!("empty DP process group"))?; + ctx.init_parallel_nccl( + rank, world_size, tp_rank, tp_size, tp_color, dp_rank, dp_size, dp_color, + )?; + } else { + let ret = ctx.init_nccl(); + if ret != 0 { + bail!("C++ NCCL init failed (code {})", ret); + } } info!( rank, @@ -623,7 +673,12 @@ fn train_impl( for accumulation_index in 0..gradient_accumulation_steps { let micro_step = step * gradient_accumulation_steps + accumulation_index; let data_start = if is_data_parallel || ep_a2a_sharded { - (micro_step * batch_size * world_size + rank * batch_size) % data.len() + let (replica_rank, replica_count) = if is_data_parallel { + (dp_rank, dp_size) + } else { + (rank, world_size) + }; + (micro_step * batch_size * replica_count + replica_rank * batch_size) % data.len() } else { (micro_step * batch_size) % data.len() }; diff --git a/crates/rustrain-qwen3-6/tests/native_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_smoke.cpp index 36114967..bee42145 100644 --- a/crates/rustrain-qwen3-6/tests/native_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_smoke.cpp @@ -366,7 +366,7 @@ static int run_dynamic_dp_smoke( } int main() { - assert(qwen36_kernel_abi_version() == 15); + assert(qwen36_kernel_abi_version() == 16); const int world = std::atoi(std::getenv("WORLD_SIZE") ? std::getenv("WORLD_SIZE") : "1"); const int process_rank = std::atoi(std::getenv("RANK") ? std::getenv("RANK") : "0"); const int local_rank = std::atoi(std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); diff --git a/crates/rustrain-qwen3-6/tests/native_tp_dp_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_dp_smoke.cpp new file mode 100644 index 00000000..10bd6813 --- /dev/null +++ b/crates/rustrain-qwen3-6/tests/native_tp_dp_smoke.cpp @@ -0,0 +1,534 @@ +#include + +#include +#include +#include +#include +#include +#include +#include + +struct LayerConfig { + int64_t layer_type, num_heads, num_kv_heads, head_dim; + int64_t num_k_heads, key_dim, num_v_heads, val_dim, conv_kernel; + double partial_rotary_factor, rope_theta, rms_eps; + int64_t num_experts, top_k, moe_intermediate, expert_start, expert_count; + int64_t intermediate_size; + int32_t norm_topk_prob; + void* nccl_comm; + void* nccl_stream; +}; + +extern "C" int64_t qwen36_kernel_abi_version(); +extern "C" void qwen36_set_cuda_device(int32_t); +extern "C" void* qwen36_create_training_context( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*); +extern "C" void* qwen36_create_training_context_ex( + void**, int64_t, void*, void*, void*, void*, int64_t, int32_t, + double, double, double, double, double, int64_t, double, int64_t, + const int64_t*, int64_t, const char*, int32_t); +extern "C" int32_t qwen36_init_nccl(void*); +extern "C" int32_t qwen36_init_parallel_nccl( + void*, int32_t, int32_t, int32_t, int32_t, int32_t, + int32_t, int32_t, int32_t); +extern "C" int32_t qwen36_set_lora_tensor(void*, int64_t, int32_t, void*); +extern "C" void* qwen36_get_lora_a(void*, int64_t); +extern "C" void* qwen36_get_lora_b(void*, int64_t); +extern "C" int64_t qwen36_export_optimizer_state( + void*, void**, void**, int64_t); +extern "C" double qwen36_train_step(void*, void*, void*, void*); +extern "C" int64_t qwen36_add_lora( + void*, int64_t, double, const int64_t*, int64_t, const char*); +extern "C" double qwen36_train_multi_lora_selected( + void*, void*, void*, void*, const int64_t*, int32_t, int32_t); +extern "C" int64_t qwen36_get_adapter_step_count(void*, int64_t); +extern "C" void qwen36_free_training_context(void*); + +static constexpr double kLearningRate = 1e-3; +static constexpr double kBeta1 = 0.9; +static constexpr double kBeta2 = 0.999; +static constexpr double kAdamEps = 1e-8; +static constexpr int64_t kOptimizerSlots = 14; + +static at::Tensor deterministic(std::initializer_list shape, double scale) { + int64_t count = 1; + for (int64_t dim : shape) count *= dim; + return ((at::arange(count, + at::TensorOptions().device(at::kCUDA).dtype(at::kFloat)) + .remainder(19) - 9.0) * scale) + .reshape(shape).to(at::kBFloat16); +} + +static std::vector pointers(std::vector& tensors) { + std::vector result; + result.reserve(tensors.size()); + for (auto& tensor : tensors) result.push_back(&tensor); + return result; +} + +static double max_diff(const at::Tensor& lhs, const at::Tensor& rhs) { + return (lhs - rhs).abs().max().item(); +} + +struct Batch { + at::Tensor input_ids; + at::Tensor target_mask; + at::Tensor attention_mask; +}; + +static Batch make_batch(int dp_rank) { + auto options_long = at::TensorOptions().device(at::kCUDA).dtype(at::kLong); + auto options_float = at::TensorOptions().device(at::kCUDA).dtype(at::kFloat); + auto options_bool = at::TensorOptions().device(at::kCUDA).dtype(at::kBool); + if (dp_rank == 0) { + return { + at::tensor({1, 2, 3, 4}, options_long).reshape({1, 4}), + at::tensor({1.0, 1.0, 1.0, 1.0}, options_float).reshape({1, 4}), + at::ones({1, 4}, options_bool), + }; + } + return { + at::tensor({4, 2, 5, 1}, options_long).reshape({1, 4}), + at::tensor({1.0, 1.0, 0.0, 0.0}, options_float).reshape({1, 4}), + at::ones({1, 4}, options_bool), + }; +} + +struct LoraWeights { + at::Tensor q_a, q_b; + at::Tensor k_a, k_b; + at::Tensor v_a, v_b; + at::Tensor o_a, o_b; +}; + +static LoraWeights make_lora_weights( + int64_t hidden, int64_t heads, int64_t kv_heads, + int64_t head_dim, int64_t lora_rank +) { + return { + deterministic({lora_rank, hidden}, 0.0020), + deterministic({2 * heads * head_dim, lora_rank}, 0.0010), + deterministic({lora_rank, hidden}, 0.0018), + deterministic({kv_heads * head_dim, lora_rank}, 0.0011), + deterministic({lora_rank, hidden}, 0.0016), + deterministic({kv_heads * head_dim, lora_rank}, 0.0009), + deterministic({lora_rank, heads * head_dim}, 0.0015), + deterministic({hidden, lora_rank}, 0.0012), + }; +} + +static LoraWeights clone_lora_weights(const LoraWeights& weights) { + return { + weights.q_a.clone(), weights.q_b.clone(), + weights.k_a.clone(), weights.k_b.clone(), + weights.v_a.clone(), weights.v_b.clone(), + weights.o_a.clone(), weights.o_b.clone(), + }; +} + +static void offset_lora_weights(LoraWeights& weights, double offset) { + std::array tensors = { + &weights.q_a, &weights.q_b, &weights.k_a, &weights.k_b, + &weights.v_a, &weights.v_b, &weights.o_a, &weights.o_b, + }; + for (auto* tensor : tensors) tensor->add_(offset); +} + +static void install_full_lora(void* context, LoraWeights& weights) { + std::array tensors = { + &weights.q_a, &weights.q_b, &weights.k_a, &weights.k_b, + &weights.v_a, &weights.v_b, &weights.o_a, &weights.o_b, + }; + for (int64_t index = 0; index < 4; ++index) { + assert(qwen36_set_lora_tensor(context, index, 0, tensors[2 * index]) == 0); + assert(qwen36_set_lora_tensor(context, index, 1, tensors[2 * index + 1]) == 0); + } +} + +static LoraWeights shard_lora_for_tp( + const LoraWeights& full, int tp_rank, int64_t local_heads, + int64_t local_kv_heads, int64_t head_dim +) { + return { + full.q_a, + full.q_b.narrow( + 0, tp_rank * 2 * local_heads * head_dim, + 2 * local_heads * head_dim).contiguous(), + full.k_a, + full.k_b.narrow( + 0, tp_rank * local_kv_heads * head_dim, + local_kv_heads * head_dim).contiguous(), + full.v_a, + full.v_b.narrow( + 0, tp_rank * local_kv_heads * head_dim, + local_kv_heads * head_dim).contiguous(), + full.o_a.narrow( + 1, tp_rank * local_heads * head_dim, + local_heads * head_dim).contiguous(), + full.o_b, + }; +} + +static std::array lora_parameters(void* context) { + std::array result{}; + for (int64_t index = 0; index < 4; ++index) { + result[2 * index] = reinterpret_cast( + qwen36_get_lora_a(context, index)); + result[2 * index + 1] = reinterpret_cast( + qwen36_get_lora_b(context, index)); + assert(result[2 * index] && result[2 * index + 1]); + } + return result; +} + +struct OptimizerState { + std::vector m; + std::vector v; +}; + +static OptimizerState optimizer_state(void* context) { + OptimizerState result{ + std::vector(kOptimizerSlots), + std::vector(kOptimizerSlots), + }; + assert(qwen36_export_optimizer_state( + context, result.m.data(), result.v.data(), kOptimizerSlots) == + kOptimizerSlots); + return result; +} + +static at::Tensor& state_tensor(std::vector& tensors, int64_t index) { + auto* tensor = reinterpret_cast(tensors[index]); + assert(tensor); + return *tensor; +} + +struct FullReference { + void* context; + double loss; + OptimizerState optimizer; +}; + +static FullReference run_full_reference( + std::vector& full_ptrs, at::Tensor& embed, at::Tensor& final_norm, + at::Tensor& lm_head, LayerConfig& config, LoraWeights& lora, + Batch& batch, int64_t vocab, int64_t lora_rank +) { + const int64_t target_layer = 0; + void* context = qwen36_create_training_context( + full_ptrs.data(), full_ptrs.size(), &embed, &final_norm, &lm_head, + &config, 1, static_cast(at::kBFloat16), + 1.0, kLearningRate, kBeta1, kBeta2, kAdamEps, + vocab, 1e-5, lora_rank, &target_layer, 1, + "q_proj,k_proj,v_proj,o_proj"); + assert(context); + install_full_lora(context, lora); + const double loss = qwen36_train_step( + context, &batch.input_ids, &batch.target_mask, &batch.attention_mask); + assert(loss > 0.0 && std::isfinite(loss)); + return {context, loss, optimizer_state(context)}; +} + +int main() { + assert(qwen36_kernel_abi_version() == 16); + const int rank = std::atoi(std::getenv("RANK")); + const int world = std::atoi(std::getenv("WORLD_SIZE")); + const int local_rank = std::atoi( + std::getenv("LOCAL_RANK") ? std::getenv("LOCAL_RANK") : "0"); + assert(world == 4 && rank >= 0 && rank < world); + const int tp_rank = rank % 2; + const int dp_rank = rank / 2; + qwen36_set_cuda_device(local_rank); + + constexpr int64_t hidden = 8; + constexpr int64_t heads = 4; + constexpr int64_t kv_heads = 2; + constexpr int64_t head_dim = 2; + constexpr int64_t intermediate = 12; + constexpr int64_t vocab = 16; + constexpr int64_t lora_rank = 4; + constexpr int64_t local_heads = heads / 2; + constexpr int64_t local_kv_heads = kv_heads / 2; + + std::vector full_weights; + full_weights.push_back(at::ones( + {hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + full_weights.push_back(at::ones( + {hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + full_weights.push_back(deterministic({2 * heads * head_dim, hidden}, 0.010)); + full_weights.push_back(at::ones( + {head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + full_weights.push_back(deterministic({kv_heads * head_dim, hidden}, 0.012)); + full_weights.push_back(at::ones( + {head_dim}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16))); + full_weights.push_back(deterministic({kv_heads * head_dim, hidden}, 0.008)); + full_weights.push_back(deterministic({hidden, heads * head_dim}, 0.011)); + full_weights.push_back(deterministic({intermediate, hidden}, 0.009)); + full_weights.push_back(deterministic({intermediate, hidden}, 0.007)); + full_weights.push_back(deterministic({hidden, intermediate}, 0.010)); + for (auto& weight : full_weights) weight.set_requires_grad(false); + + std::vector local_weights; + local_weights.push_back(full_weights[0]); + local_weights.push_back(full_weights[1]); + local_weights.push_back(full_weights[2].narrow( + 0, tp_rank * 2 * local_heads * head_dim, + 2 * local_heads * head_dim).contiguous()); + local_weights.push_back(full_weights[3]); + local_weights.push_back(full_weights[4].narrow( + 0, tp_rank * local_kv_heads * head_dim, + local_kv_heads * head_dim).contiguous()); + local_weights.push_back(full_weights[5]); + local_weights.push_back(full_weights[6].narrow( + 0, tp_rank * local_kv_heads * head_dim, + local_kv_heads * head_dim).contiguous()); + local_weights.push_back(full_weights[7].narrow( + 1, tp_rank * local_heads * head_dim, + local_heads * head_dim).contiguous()); + local_weights.insert( + local_weights.end(), full_weights.begin() + 8, full_weights.end()); + + auto embed = deterministic({vocab, hidden}, 0.020); + auto final_norm = at::ones( + {hidden}, at::TensorOptions().device(at::kCUDA).dtype(at::kBFloat16)); + auto lm_head = deterministic({vocab, hidden}, 0.015); + LayerConfig config{}; + config.layer_type = 0; + config.num_heads = heads; + config.num_kv_heads = kv_heads; + config.head_dim = head_dim; + config.partial_rotary_factor = 1.0; + config.rope_theta = 10000.0; + config.rms_eps = 1e-5; + config.intermediate_size = intermediate; + + auto local_ptrs = pointers(local_weights); + auto full_ptrs = pointers(full_weights); + const int64_t target_layer = 0; + + // Native validation must reject a topology whose declared axes do not + // cover the process world before any communicator is initialized. + setenv("TP_SIZE", "2", 1); + setenv("DP_SIZE", "3", 1); + setenv("RUSTRAIN_DATA_PARALLEL", "1", 1); + void* invalid = qwen36_create_training_context_ex( + local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head, + &config, 1, static_cast(at::kBFloat16), + 1.0, kLearningRate, kBeta1, kBeta2, kAdamEps, + vocab, 1e-5, lora_rank, &target_layer, 1, + "q_proj,k_proj,v_proj,o_proj", 3); + assert(invalid == nullptr); + + setenv("DP_SIZE", "2", 1); + auto full_lora = make_lora_weights( + hidden, heads, kv_heads, head_dim, lora_rank); + auto local_lora = shard_lora_for_tp( + full_lora, tp_rank, local_heads, local_kv_heads, head_dim); + auto synchronized_lora = clone_lora_weights(local_lora); + if (dp_rank == 1) { + // DP must restore every local shard from the matching dp_rank=0 peer. + offset_lora_weights(local_lora, 0.25); + } else if (tp_rank == 1) { + // TP must restore projection-replicated factors from tp_rank=0. + local_lora.q_a.add_(0.25); + local_lora.k_a.add_(0.25); + local_lora.v_a.add_(0.25); + local_lora.o_b.add_(0.25); + } + void* distributed = qwen36_create_training_context_ex( + local_ptrs.data(), local_ptrs.size(), &embed, &final_norm, &lm_head, + &config, 1, static_cast(at::kBFloat16), + 1.0, kLearningRate, kBeta1, kBeta2, kAdamEps, + vocab, 1e-5, lora_rank, &target_layer, 1, + "q_proj,k_proj,v_proj,o_proj", 3); + assert(distributed); + install_full_lora(distributed, local_lora); + assert(qwen36_init_parallel_nccl( + distributed, rank, world, + tp_rank, 2, dp_rank * 2, + dp_rank, 2, tp_rank) == 0); + + const auto synchronized_params = lora_parameters(distributed); + const std::array expected_synchronized = { + &synchronized_lora.q_a, &synchronized_lora.q_b, + &synchronized_lora.k_a, &synchronized_lora.k_b, + &synchronized_lora.v_a, &synchronized_lora.v_b, + &synchronized_lora.o_a, &synchronized_lora.o_b, + }; + double broadcast_diff = 0.0; + for (int64_t index = 0; index < 8; ++index) { + broadcast_diff = std::max( + broadcast_diff, + max_diff(*synchronized_params[index], *expected_synchronized[index])); + } + assert(broadcast_diff == 0.0); + + auto batch0 = make_batch(0); + auto batch1 = make_batch(1); + Batch full_batch{ + at::cat({batch0.input_ids, batch1.input_ids}, 0), + at::cat({batch0.target_mask, batch1.target_mask}, 0), + at::cat({batch0.attention_mask, batch1.attention_mask}, 0), + }; + Batch& local_batch = dp_rank == 0 ? batch0 : batch1; + const double distributed_loss = qwen36_train_step( + distributed, &local_batch.input_ids, &local_batch.target_mask, + &local_batch.attention_mask); + assert(distributed_loss > 0.0 && std::isfinite(distributed_loss)); + + // Reference contexts stay process-local. The concatenated batch is the + // hard oracle for the DP-reduced update; the two local references prove + // that unequal token counts are not accidentally averaged by replica. + setenv("TP_SIZE", "1", 1); + setenv("DP_SIZE", "1", 1); + unsetenv("RUSTRAIN_DATA_PARALLEL"); + auto reference_lora = make_lora_weights( + hidden, heads, kv_heads, head_dim, lora_rank); + auto reference = run_full_reference( + full_ptrs, embed, final_norm, lm_head, config, reference_lora, + full_batch, vocab, lora_rank); + auto local0_lora = make_lora_weights( + hidden, heads, kv_heads, head_dim, lora_rank); + auto local0 = run_full_reference( + full_ptrs, embed, final_norm, lm_head, config, local0_lora, + batch0, vocab, lora_rank); + auto local1_lora = make_lora_weights( + hidden, heads, kv_heads, head_dim, lora_rank); + auto local1 = run_full_reference( + full_ptrs, embed, final_norm, lm_head, config, local1_lora, + batch1, vocab, lora_rank); + + auto distributed_params = lora_parameters(distributed); + auto reference_params = lora_parameters(reference.context); + auto distributed_optimizer = optimizer_state(distributed); + std::array before = { + synchronized_lora.q_a, synchronized_lora.q_b, + synchronized_lora.k_a, synchronized_lora.k_b, + synchronized_lora.v_a, synchronized_lora.v_b, + synchronized_lora.o_a, synchronized_lora.o_b, + }; + auto reference_slice = [&](const at::Tensor& tensor, int64_t index) { + switch (index) { + case 1: + return tensor.narrow( + 0, tp_rank * 2 * local_heads * head_dim, + 2 * local_heads * head_dim); + case 3: + case 5: + return tensor.narrow( + 0, tp_rank * local_kv_heads * head_dim, + local_kv_heads * head_dim); + case 6: + return tensor.narrow( + 1, tp_rank * local_heads * head_dim, + local_heads * head_dim); + default: + return tensor; + } + }; + + double parameter_diff = 0.0; + double optimizer_m_diff = 0.0; + double optimizer_v_diff = 0.0; + double adam_error = 0.0; + double weighted_reference_error = 0.0; + double equal_replica_gap = 0.0; + constexpr double token_count0 = 3.0; + constexpr double token_count1 = 1.0; + for (int64_t index = 0; index < 8; ++index) { + auto expected_parameter = reference_slice(*reference_params[index], index); + auto expected_m = reference_slice( + state_tensor(reference.optimizer.m, index), index); + auto expected_v = reference_slice( + state_tensor(reference.optimizer.v, index), index); + parameter_diff = std::max( + parameter_diff, max_diff(*distributed_params[index], expected_parameter)); + optimizer_m_diff = std::max( + optimizer_m_diff, + max_diff(state_tensor(distributed_optimizer.m, index), expected_m)); + optimizer_v_diff = std::max( + optimizer_v_diff, + max_diff(state_tensor(distributed_optimizer.v, index), expected_v)); + + auto expected_adam_parameter = ( + before[index].to(at::kFloat) - kLearningRate * + (state_tensor(distributed_optimizer.m, index) / (1.0 - kBeta1)) / + ((state_tensor(distributed_optimizer.v, index) / + (1.0 - kBeta2)).sqrt() + kAdamEps)) + .to(at::kBFloat16); + adam_error = std::max( + adam_error, + max_diff(*distributed_params[index], expected_adam_parameter)); + + auto local_m0 = state_tensor(local0.optimizer.m, index); + auto local_m1 = state_tensor(local1.optimizer.m, index); + auto token_weighted_m = + (local_m0 * token_count0 + local_m1 * token_count1) / + (token_count0 + token_count1); + auto equal_replica_m = (local_m0 + local_m1) * 0.5; + weighted_reference_error = std::max( + weighted_reference_error, max_diff( + state_tensor(reference.optimizer.m, index), token_weighted_m)); + equal_replica_gap = std::max( + equal_replica_gap, max_diff( + state_tensor(reference.optimizer.m, index), equal_replica_m)); + } + + const double expected_local_loss = + dp_rank == 0 ? local0.loss : local1.loss; + const double local_loss_diff = std::abs(distributed_loss - expected_local_loss); + std::printf( + "native_tp_dp_smoke rank=%d tp_rank=%d dp_rank=%d " + "broadcast_diff=%0.8e local_loss_diff=%0.8e " + "parameter_diff=%0.8e m_diff=%0.8e " + "v_diff=%0.8e adam_error=%0.8e weighted_reference_error=%0.8e " + "equal_replica_gap=%0.8e\n", + rank, tp_rank, dp_rank, broadcast_diff, local_loss_diff, parameter_diff, + optimizer_m_diff, optimizer_v_diff, adam_error, + weighted_reference_error, equal_replica_gap); + std::fflush(stdout); + + // TP row-parallel BF16 rounds each local matmul before the collective. + assert(local_loss_diff < 5e-3); + assert(parameter_diff <= 2e-3); + assert(optimizer_m_diff < 5e-5 && optimizer_v_diff < 5e-8); + assert(adam_error < 1e-8); + assert(weighted_reference_error < 5e-5); + assert(equal_replica_gap > 1e-8); + + // A fixed-size registry signature must reject a different tenant order + // before any adapter-shaped collective can mix gradients or deadlock. + const int64_t adapter_one = qwen36_add_lora( + distributed, lora_rank, lora_rank, &target_layer, 1, "q_proj"); + const int64_t adapter_two = qwen36_add_lora( + distributed, lora_rank, lora_rank, &target_layer, 1, "q_proj"); + assert(adapter_one > 0 && adapter_two > adapter_one); + std::array mismatched_ids = dp_rank == 0 + ? std::array{adapter_one, adapter_two} + : std::array{adapter_two, adapter_one}; + auto dynamic_input = local_batch.input_ids.repeat({2, 1}); + auto dynamic_target = local_batch.target_mask.repeat({2, 1}); + auto dynamic_attention = local_batch.attention_mask.repeat({2, 1}); + const double mismatched_loss = qwen36_train_multi_lora_selected( + distributed, &dynamic_input, &dynamic_target, &dynamic_attention, + mismatched_ids.data(), mismatched_ids.size(), lora_rank); + assert(mismatched_loss < 0.0); + assert(qwen36_get_adapter_step_count(distributed, adapter_one) == 0); + assert(qwen36_get_adapter_step_count(distributed, adapter_two) == 0); + + const std::array ordered_ids{adapter_one, adapter_two}; + const double dynamic_loss = qwen36_train_multi_lora_selected( + distributed, &dynamic_input, &dynamic_target, &dynamic_attention, + ordered_ids.data(), ordered_ids.size(), lora_rank); + assert(dynamic_loss > 0.0 && std::isfinite(dynamic_loss)); + assert(qwen36_get_adapter_step_count(distributed, adapter_one) == 1); + assert(qwen36_get_adapter_step_count(distributed, adapter_two) == 1); + + qwen36_free_training_context(local1.context); + qwen36_free_training_context(local0.context); + qwen36_free_training_context(reference.context); + qwen36_free_training_context(distributed); + return 0; +} diff --git a/crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp b/crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp index 5d03651e..356c55de 100644 --- a/crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp +++ b/crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp @@ -42,7 +42,7 @@ extern "C" void qwen36_free_training_context(void*); namespace { -constexpr int64_t kAbiVersion = 15; +constexpr int64_t kAbiVersion = 16; constexpr int32_t kBaseTpAttention = 1 << 0; static int env_int(const char* name, int fallback) { diff --git a/crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp index aabd341b..5159967e 100644 --- a/crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp @@ -58,7 +58,7 @@ extern "C" void qwen36_free_training_context(void*); namespace { -constexpr int64_t kAbiVersion = 15; +constexpr int64_t kAbiVersion = 16; constexpr int32_t kBaseTpAttention = 1 << 0; constexpr int64_t kLayers = 2; constexpr int64_t kHidden = 32; diff --git a/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp b/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp index 07157a9a..e6d6970f 100644 --- a/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp +++ b/crates/rustrain-qwen3-6/tests/native_tp_latent_smoke.cpp @@ -130,7 +130,7 @@ int main() { configs, 2, static_cast(at::kBFloat16), 1.0, 1e-3, 0.9, 0.999, 1e-8, vocab, 1e-5, lora_rank, // This regression exercises latent-rank-only TP with replicated GDN - // base weights. ABI15 bit 0 now explicitly enables base GDN head TP. + // base weights. ABI16 bit 0 explicitly enables base GDN head TP. target_layers, 2, "in_proj_qkv", 0); assert(distributed && qwen36_init_nccl(distributed) == 0); diff --git a/crates/rustrain-server/Cargo.toml b/crates/rustrain-server/Cargo.toml index 3b5f52d1..fdb2da4f 100644 --- a/crates/rustrain-server/Cargo.toml +++ b/crates/rustrain-server/Cargo.toml @@ -9,6 +9,7 @@ rustrain-qwen3-6 = { path = "../rustrain-qwen3-6" } rustrain-checkpoint = { path = "../rustrain-checkpoint" } rustrain-train = { path = "../rustrain-train" } rustrain-ipc = { path = "../rustrain-ipc" } +rustrain-parallel = { path = "../rustrain-parallel" } anyhow.workspace = true axum.workspace = true base64 = "0.22" diff --git a/crates/rustrain-server/src/api.rs b/crates/rustrain-server/src/api.rs index 21d820ae..c0ae59dd 100644 --- a/crates/rustrain-server/src/api.rs +++ b/crates/rustrain-server/src/api.rs @@ -286,21 +286,20 @@ fn decode_int64_vec(t: &TensorHttp) -> Result, String> { Ok(values) } -fn validate_multi_lora_http_shapes( +fn validate_train_http_shapes( input_ids: &TensorHttp, target_mask: &TensorHttp, attention_mask: &TensorHttp, - n_total: i32, -) -> Result { - if n_total <= 0 { - return Err(format!("n_total must be positive, got {n_total}")); - } +) -> Result<(usize, usize), String> { if input_ids.shape.len() != 1 && input_ids.shape.len() != 2 { - return Err(format!("input_ids must have shape [seq] or [batch, seq], got {:?}", input_ids.shape)); + return Err(format!( + "input_ids must have shape [seq] or [batch, seq], got {:?}", + input_ids.shape + )); } if target_mask.shape != input_ids.shape || attention_mask.shape != input_ids.shape { return Err(format!( - "multi-LoRA masks must have the same shape as input_ids: input={:?} target={:?} attention={:?}", + "masks must have the same shape as input_ids: input={:?} target={:?} attention={:?}", input_ids.shape, target_mask.shape, attention_mask.shape )); } @@ -311,16 +310,37 @@ fn validate_multi_lora_http_shapes( if seq_len <= 0 { return Err(format!("sequence length must be positive, got {seq_len}")); } - if input_ids.shape.len() == 2 { - let batch = input_ids.shape[0]; - if batch != 1 && batch != i64::from(n_total) { - return Err(format!( - "multi-LoRA batch must be 1 or n_total={}, got {}", - n_total, batch - )); - } + let batch_size = if input_ids.shape.len() == 2 { + input_ids.shape[0] + } else { + 1 + }; + if batch_size <= 0 { + return Err(format!("batch size must be positive, got {batch_size}")); + } + Ok(( + usize::try_from(batch_size).map_err(|_| format!("invalid batch size {batch_size}"))?, + usize::try_from(seq_len).map_err(|_| format!("invalid sequence length {seq_len}"))?, + )) +} + +fn validate_multi_lora_http_shapes( + input_ids: &TensorHttp, + target_mask: &TensorHttp, + attention_mask: &TensorHttp, + n_total: i32, +) -> Result<(usize, usize), String> { + if n_total <= 0 { + return Err(format!("n_total must be positive, got {n_total}")); } - usize::try_from(seq_len).map_err(|_| format!("invalid sequence length {seq_len}")) + let (batch_size, seq_len) = validate_train_http_shapes(input_ids, target_mask, attention_mask)?; + let n_total = n_total as usize; + if batch_size != 1 && batch_size % n_total != 0 { + return Err(format!( + "multi-LoRA batch must be 1 or a positive multiple of n_total={n_total}, got {batch_size}" + )); + } + Ok((batch_size, seq_len)) } fn decode_tensor(t: &TensorHttp) -> Result { @@ -761,16 +781,19 @@ async fn ep_train_step( Path(id): Path, Json(req): Json, ) -> Result, (StatusCode, Json)> { + let (batch_size, seq_len) = + validate_train_http_shapes(&req.input_ids, &req.target_mask, &req.attention_mask) + .map_err(|e| err_resp(&e))?; let input_ids = decode_int64_vec(&req.input_ids).map_err(|e| err_resp(&e))?; let target_mask = decode_int64_vec(&req.target_mask).map_err(|e| err_resp(&e))?; let attention_mask = decode_int64_vec(&req.attention_mask).map_err(|e| err_resp(&e))?; - let seq_len = input_ids.len(); let cmd = rustrain_ipc::EpCommand::TrainStep { session_id: id, input_ids, target_mask, attention_mask, + batch_size, seq_len, }; match state.coordinator.dispatch(&cmd) { @@ -812,7 +835,7 @@ async fn ep_train_multi_lora( let input_ids = decode_int64_vec(&req.input_ids).map_err(|e| err_resp(&e))?; let target_mask = decode_int64_vec(&req.target_mask).map_err(|e| err_resp(&e))?; let attention_mask = decode_int64_vec(&req.attention_mask).map_err(|e| err_resp(&e))?; - let seq_len = validate_multi_lora_http_shapes( + let (batch_size, seq_len) = validate_multi_lora_http_shapes( &req.input_ids, &req.target_mask, &req.attention_mask, @@ -837,6 +860,7 @@ async fn ep_train_multi_lora( input_ids, target_mask, attention_mask, + batch_size, seq_len, n_total: req.n_total, lora_rank: req.lora_rank, @@ -849,6 +873,67 @@ async fn ep_train_multi_lora( } } +#[cfg(test)] +mod tensor_http_shape_tests { + use super::{TensorHttp, validate_multi_lora_http_shapes, validate_train_http_shapes}; + + fn tensor(shape: &[i64]) -> TensorHttp { + TensorHttp { + data: String::new(), + shape: shape.to_vec(), + dtype: "int64".into(), + } + } + + #[test] + fn train_shape_keeps_batch_and_sequence_dimensions() { + let input = tensor(&[8, 128]); + let target = tensor(&[8, 128]); + let attention = tensor(&[8, 128]); + + assert_eq!( + validate_train_http_shapes(&input, &target, &attention).unwrap(), + (8, 128) + ); + } + + #[test] + fn train_shape_rejects_mismatched_masks_and_non_positive_dimensions() { + let input = tensor(&[4, 16]); + assert!(validate_train_http_shapes(&input, &tensor(&[2, 16]), &tensor(&[4, 16])).is_err()); + assert!( + validate_train_http_shapes(&tensor(&[0, 16]), &tensor(&[0, 16]), &tensor(&[0, 16])) + .is_err() + ); + assert!( + validate_train_http_shapes(&tensor(&[4, 0]), &tensor(&[4, 0]), &tensor(&[4, 0])) + .is_err() + ); + } + + #[test] + fn multi_lora_shape_accepts_replica_and_global_row_layouts() { + for rows in [1, 3, 6] { + let input = tensor(&[rows, 32]); + assert_eq!( + validate_multi_lora_http_shapes( + &input, + &tensor(&[rows, 32]), + &tensor(&[rows, 32]), + 3 + ) + .unwrap(), + (rows as usize, 32) + ); + } + let invalid = tensor(&[5, 32]); + assert!( + validate_multi_lora_http_shapes(&invalid, &tensor(&[5, 32]), &tensor(&[5, 32]), 3) + .is_err() + ); + } +} + async fn ep_add_lora( State(state): State>, Path(id): Path, diff --git a/crates/rustrain-server/src/ep.rs b/crates/rustrain-server/src/ep.rs index 70e6b865..0c82af33 100644 --- a/crates/rustrain-server/src/ep.rs +++ b/crates/rustrain-server/src/ep.rs @@ -7,12 +7,13 @@ //! //! Workers use `worker_main()` to enter the wait loop. -use std::collections::HashMap; use std::io; +use std::ops::Range; use std::path::PathBuf; use std::sync::Arc; use rustrain_ipc::{EpChannel, EpCommand, EpResult, EpWorker}; +use rustrain_parallel::topology::ParallelTopology; use tch::{Device, Kind}; use crate::session::{ @@ -247,11 +248,46 @@ fn execute_command(session: &mut Qwen36Session, cmd: &EpCommand) -> EpResult { EpCommand::ListLora { .. } => { EpResult::AdapterIds(session.list_lora()) } - EpCommand::TrainStep { input_ids, target_mask, attention_mask, seq_len, .. } => { + EpCommand::TrainStep { + input_ids, + target_mask, + attention_mask, + batch_size, + seq_len, + .. + } => { + let dp_shard = match dense_dp_shard_from_env() { + Ok(shard) => shard, + Err(error) => return EpResult::Error(error), + }; + let rows = match train_step_rows(*batch_size, dp_shard) { + Ok(rows) => rows, + Err(error) => return EpResult::Error(error), + }; + let elements = match tensor_element_range(*batch_size, *seq_len, rows.clone()) { + Ok(elements) => elements, + Err(error) => return EpResult::Error(error), + }; + if let Err(error) = validate_flat_tensor_lengths( + *batch_size, + *seq_len, + input_ids.len(), + target_mask.len(), + attention_mask.len(), + ) { + return EpResult::Error(error); + } let sl = *seq_len as i64; - let input_ids_tensor = tch::Tensor::from_slice(input_ids).reshape(&[1, sl]).to_device(session.device()); - let target_mask_tensor = tch::Tensor::from_slice(target_mask).reshape(&[1, sl]).to_device(session.device()); - let attention_mask_tensor = tch::Tensor::from_slice(attention_mask).reshape(&[1, sl]).to_device(session.device()); + let local_batch = rows.len() as i64; + let input_ids_tensor = tch::Tensor::from_slice(&input_ids[elements.clone()]) + .reshape(&[local_batch, sl]) + .to_device(session.device()); + let target_mask_tensor = tch::Tensor::from_slice(&target_mask[elements.clone()]) + .reshape(&[local_batch, sl]) + .to_device(session.device()); + let attention_mask_tensor = tch::Tensor::from_slice(&attention_mask[elements]) + .reshape(&[local_batch, sl]) + .to_device(session.device()); match session.train_step(TrainInput { input_ids: input_ids_tensor, @@ -266,43 +302,56 @@ fn execute_command(session: &mut Qwen36Session, cmd: &EpCommand) -> EpResult { input_ids, target_mask, attention_mask, + batch_size, seq_len, n_total, lora_rank, adapter_ids, .. } => { - let sl = *seq_len as i64; - let batch = if *n_total > 0 - && input_ids.len() == (*n_total as usize).saturating_mul(*seq_len) - { - *n_total as i64 - } else if input_ids.len() == *seq_len { - 1 - } else { - return EpResult::Error(format!( - "multi-LoRA input length {} is incompatible with n_total={} seq_len={}", - input_ids.len(), n_total, seq_len - )); - }; - let expected = (batch as usize).saturating_mul(*seq_len); - if target_mask.len() != expected || attention_mask.len() != expected { - return EpResult::Error(format!( - "multi-LoRA mask lengths must equal {}, got target={} attention={}", - expected, - target_mask.len(), - attention_mask.len() - )); + if let Err(error) = validate_flat_tensor_lengths( + *batch_size, + *seq_len, + input_ids.len(), + target_mask.len(), + attention_mask.len(), + ) { + return EpResult::Error(error); } - let input_ids_tensor = tch::Tensor::from_slice(input_ids).reshape(&[batch, sl]).to_device(session.device()); - let target_mask_tensor = tch::Tensor::from_slice(target_mask).reshape(&[batch, sl]).to_device(session.device()); - let attention_mask_tensor = tch::Tensor::from_slice(attention_mask).reshape(&[batch, sl]).to_device(session.device()); + let dp_shard = match dense_dp_shard_from_env() { + Ok(shard) => shard, + Err(error) => return EpResult::Error(error), + }; + let rows = match multi_lora_rows(*batch_size, *n_total, dp_shard) { + Ok(rows) => rows, + Err(error) => return EpResult::Error(error), + }; + let elements = match tensor_element_range(*batch_size, *seq_len, rows.clone()) { + Ok(elements) => elements, + Err(error) => return EpResult::Error(error), + }; + let sl = *seq_len as i64; + let local_batch = rows.len() as i64; + let input_ids_tensor = tch::Tensor::from_slice(&input_ids[elements.clone()]) + .reshape(&[local_batch, sl]) + .to_device(session.device()); + let target_mask_tensor = tch::Tensor::from_slice(&target_mask[elements.clone()]) + .reshape(&[local_batch, sl]) + .to_device(session.device()); + let attention_mask_tensor = tch::Tensor::from_slice(&attention_mask[elements]) + .reshape(&[local_batch, sl]) + .to_device(session.device()); - match session.train_multi_lora(TrainInput { - input_ids: input_ids_tensor, - target_mask: target_mask_tensor, - attention_mask: attention_mask_tensor, - }, *n_total, *lora_rank, adapter_ids) { + match session.train_multi_lora( + TrainInput { + input_ids: input_ids_tensor, + target_mask: target_mask_tensor, + attention_mask: attention_mask_tensor, + }, + *n_total, + *lora_rank, + adapter_ids, + ) { Ok(TrainOutput { loss, .. }) => EpResult::Loss(loss), Err(e) => EpResult::Error(e.to_string()), } @@ -344,3 +393,216 @@ fn execute_command(session: &mut Qwen36Session, cmd: &EpCommand) -> EpResult { } } } + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct DpShard { + rank: usize, + size: usize, +} + +fn dense_dp_shard_from_env() -> Result, String> { + let data_parallel = std::env::var("RUSTRAIN_DATA_PARALLEL") + .map(|value| !value.is_empty() && value != "0") + .unwrap_or(false); + if !data_parallel { + return Ok(None); + } + + let topology = ParallelTopology::from_env() + .map_err(|error| format!("invalid dense DP topology: {error}"))?; + let global_rank = std::env::var("RANK") + .map_err(|_| "RANK is required for dense data parallel training".to_string())? + .parse::() + .map_err(|_| "RANK must be a non-negative integer".to_string())?; + dp_shard_for_topology(&topology, global_rank) +} + +fn dp_shard_for_topology( + topology: &ParallelTopology, + global_rank: usize, +) -> Result, String> { + let size = topology.data_parallel_size(); + if size <= 1 { + return Ok(None); + } + let rank = topology + .data_rank(global_rank) + .map_err(|error| format!("invalid dense DP rank: {error}"))?; + Ok(Some(DpShard { rank, size })) +} + +fn validate_flat_tensor_lengths( + batch_size: usize, + seq_len: usize, + input_len: usize, + target_len: usize, + attention_len: usize, +) -> Result<(), String> { + if batch_size == 0 || seq_len == 0 { + return Err(format!( + "batch_size and seq_len must be positive, got batch_size={batch_size} seq_len={seq_len}" + )); + } + let expected = batch_size + .checked_mul(seq_len) + .ok_or_else(|| "batch tensor element count overflowed usize".to_string())?; + if input_len != expected || target_len != expected || attention_len != expected { + return Err(format!( + "tensor lengths must match batch_size={batch_size} seq_len={seq_len} (expected {expected}), got input={input_len} target={target_len} attention={attention_len}" + )); + } + Ok(()) +} + +fn train_step_rows(batch_size: usize, shard: Option) -> Result, String> { + let Some(shard) = shard else { + return Ok(0..batch_size); + }; + if shard.size == 0 || shard.rank >= shard.size { + return Err(format!( + "invalid DP shard rank={}/size={}", + shard.rank, shard.size + )); + } + if batch_size % shard.size != 0 { + return Err(format!( + "global batch_size={batch_size} must be divisible by DP size {}", + shard.size + )); + } + let local_batch = batch_size / shard.size; + let start = shard.rank * local_batch; + Ok(start..start + local_batch) +} + +fn multi_lora_rows( + batch_size: usize, + n_total: i32, + shard: Option, +) -> Result, String> { + if n_total <= 0 { + return Err(format!("n_total must be positive, got {n_total}")); + } + let n_total = n_total as usize; + if batch_size == 1 || batch_size == n_total { + return Ok(0..batch_size); + } + let Some(shard) = shard else { + return Err(format!( + "multi-LoRA batch_size={batch_size} must be 1 or n_total={n_total} when DP is disabled" + )); + }; + if shard.size == 0 || shard.rank >= shard.size { + return Err(format!( + "invalid DP shard rank={}/size={}", + shard.rank, shard.size + )); + } + let global_rows = n_total + .checked_mul(shard.size) + .ok_or_else(|| "multi-LoRA global row count overflowed usize".to_string())?; + if batch_size != global_rows { + return Err(format!( + "multi-LoRA batch_size={batch_size} must be n_total={n_total} (replicated) or n_total*dp_size={global_rows} (global)" + )); + } + let start = shard.rank * n_total; + Ok(start..start + n_total) +} + +fn tensor_element_range( + batch_size: usize, + seq_len: usize, + rows: Range, +) -> Result, String> { + if rows.start > rows.end || rows.end > batch_size { + return Err(format!( + "row range {:?} is outside batch_size={batch_size}", + rows + )); + } + let start = rows + .start + .checked_mul(seq_len) + .ok_or_else(|| "tensor slice start overflowed usize".to_string())?; + let end = rows + .end + .checked_mul(seq_len) + .ok_or_else(|| "tensor slice end overflowed usize".to_string())?; + Ok(start..end) +} + +#[cfg(test)] +mod tests { + use super::{ + DpShard, dp_shard_for_topology, multi_lora_rows, tensor_element_range, + train_step_rows, validate_flat_tensor_lengths, + }; + use rustrain_parallel::topology::ParallelTopology; + + #[test] + fn train_step_slices_global_batch_contiguously_by_dp_rank() { + assert_eq!( + train_step_rows(8, Some(DpShard { rank: 0, size: 2 })).unwrap(), + 0..4 + ); + assert_eq!( + train_step_rows(8, Some(DpShard { rank: 1, size: 2 })).unwrap(), + 4..8 + ); + assert!(train_step_rows(7, Some(DpShard { rank: 0, size: 2 })).is_err()); + assert_eq!(train_step_rows(7, None).unwrap(), 0..7); + } + + #[test] + fn tp_peers_with_the_same_dp_rank_select_the_same_rows() { + let topology = ParallelTopology::new(2, 1, 2, 1, 1).unwrap(); + let tp_rank_zero_rows = train_step_rows( + 12, + dp_shard_for_topology(&topology, 2).unwrap(), + ) + .unwrap(); + let tp_rank_one_rows = train_step_rows( + 12, + dp_shard_for_topology(&topology, 3).unwrap(), + ) + .unwrap(); + assert_eq!(tp_rank_zero_rows, tp_rank_one_rows); + assert_eq!(tp_rank_zero_rows, 6..12); + } + + #[test] + fn custom_rank_order_keeps_tp_peers_on_the_same_dp_shard() { + let topology = ParallelTopology::with_order(2, 1, 2, 1, 1, "dp-tp").unwrap(); + let first_tp_peer = train_step_rows( + 12, + dp_shard_for_topology(&topology, 1).unwrap(), + ) + .unwrap(); + let second_tp_peer = train_step_rows( + 12, + dp_shard_for_topology(&topology, 3).unwrap(), + ) + .unwrap(); + assert_eq!(first_tp_peer, second_tp_peer); + assert_eq!(first_tp_peer, 6..12); + } + + #[test] + fn multi_lora_supports_replicated_and_global_dp_rows() { + let shard = DpShard { rank: 1, size: 2 }; + assert_eq!(multi_lora_rows(3, 3, Some(shard)).unwrap(), 0..3); + assert_eq!(multi_lora_rows(6, 3, Some(shard)).unwrap(), 3..6); + assert_eq!(multi_lora_rows(1, 3, Some(shard)).unwrap(), 0..1); + assert!(multi_lora_rows(9, 3, Some(shard)).is_err()); + assert!(multi_lora_rows(6, 3, None).is_err()); + } + + #[test] + fn validates_flat_shapes_before_tensor_slicing() { + validate_flat_tensor_lengths(4, 8, 32, 32, 32).unwrap(); + assert!(validate_flat_tensor_lengths(4, 8, 31, 32, 32).is_err()); + assert!(validate_flat_tensor_lengths(0, 8, 0, 0, 0).is_err()); + assert_eq!(tensor_element_range(4, 8, 1..3).unwrap(), 8..24); + } +} diff --git a/crates/rustrain-server/src/session.rs b/crates/rustrain-server/src/session.rs index 8ea9568c..c6d7483d 100644 --- a/crates/rustrain-server/src/session.rs +++ b/crates/rustrain-server/src/session.rs @@ -8,6 +8,7 @@ use tokio::sync::Mutex; use crate::checkpoint; use crate::metrics::{FileMetricsSink, MetricsSink, StepMetric}; +use rustrain_parallel::topology::ParallelTopology; use rustrain_qwen3_6::lora::{Qwen36AdapterArtifact, Qwen36LoraConfig, Qwen36LoraTargetModule}; fn lora_tp_shard_layout( @@ -339,24 +340,60 @@ impl TrainingSession for Qwen36Session { .unwrap_or(0); let tp_size = std::env::var("TP_SIZE") .or_else(|_| std::env::var("RUSTRAIN_TP_SIZE")) + .or_else(|_| std::env::var("TENSOR_MODEL_PARALLEL_SIZE")) .ok() .and_then(|s| s.parse::().ok()) .unwrap_or(1); - if tp_size > 1 { - if ep_world_size != tp_size || ep_rank >= ep_world_size { + let is_ep = ep_world_size > 1 && runtime_config.is_moe && tp_size == 1; + let dense_topology = if runtime_config.is_moe { + None + } else { + let topology = ParallelTopology::from_env_with_world_size(ep_world_size)?; + if topology.pipeline_model_parallel_size() != 1 + || topology.expert_model_parallel_size() != 1 + || topology.context_parallel_size() != 1 + { return Err(anyhow!( - "native Qwen server TP-only mode requires WORLD_SIZE=TP_SIZE and a valid global rank (world={}, tp={}, rank={})", - ep_world_size, - tp_size, - ep_rank + "native dense Qwen server supports TPxDP only (tp={} pp={} dp={} ep={} cp={})", + topology.tensor_model_parallel_size(), + topology.pipeline_model_parallel_size(), + topology.data_parallel_size(), + topology.expert_model_parallel_size(), + topology.context_parallel_size(), )); } - unsafe { - std::env::set_var("TP_SIZE", tp_size.to_string()); + if topology.tensor_model_parallel_size() != tp_size { + return Err(anyhow!( + "Qwen server topology TP={} does not match TP_SIZE={tp_size}", + topology.tensor_model_parallel_size() + )); } + topology.coordinates(ep_rank)?; + Some(topology) + }; + let dp_size = dense_topology + .as_ref() + .map(ParallelTopology::data_parallel_size) + .unwrap_or(1); + let tp_rank = dense_topology + .as_ref() + .map(|topology| topology.tensor_rank(ep_rank)) + .transpose()? + .unwrap_or(0); + let dp_rank = dense_topology + .as_ref() + .map(|topology| topology.data_rank(ep_rank)) + .transpose()? + .unwrap_or(ep_rank); + let is_data_parallel = !runtime_config.is_moe && dp_size > 1; + unsafe { + std::env::set_var("TP_SIZE", tp_size.to_string()); + std::env::set_var("DP_SIZE", dp_size.to_string()); + std::env::set_var( + "RUSTRAIN_DATA_PARALLEL", + if is_data_parallel { "1" } else { "0" }, + ); } - let is_ep = ep_world_size > 1 && runtime_config.is_moe && tp_size == 1; - let is_data_parallel = ep_world_size > 1 && !runtime_config.is_moe && tp_size == 1; let base_tp_attention = tp_size > 1; let base_tp_mlp = tp_size > 1 && !runtime_config.is_moe; if base_tp_attention { @@ -483,10 +520,7 @@ impl TrainingSession for Qwen36Session { let local_shard = if base_tp_attention { let full_attention_shard = rustrain_qwen3_6::kernel::shard_full_attention_weight_for_tp( - &name, - &tensor, - tp_size, - ep_rank % tp_size, + &name, &tensor, tp_size, tp_rank, )?; let attention_shard = if full_attention_shard.is_some() { full_attention_shard @@ -495,7 +529,7 @@ impl TrainingSession for Qwen36Session { &name, &tensor, tp_size, - ep_rank % tp_size, + tp_rank, runtime_config.linear_num_key_heads, runtime_config.linear_key_head_dim, runtime_config.linear_num_value_heads, @@ -506,10 +540,7 @@ impl TrainingSession for Qwen36Session { attention_shard } else { rustrain_qwen3_6::kernel::shard_dense_mlp_weight_for_tp( - &name, - &tensor, - tp_size, - ep_rank % tp_size, + &name, &tensor, tp_size, tp_rank, )? } } else { @@ -558,24 +589,42 @@ impl TrainingSession for Qwen36Session { req.rank, base_tp_attention, base_tp_mlp, + is_data_parallel, &all_layers, &target_modules, expert_start, expert_count, )?; - // Initialize NCCL directly in C++. The parent communicator handles - // EP/DP, while TP-only LoRA uses a split communicator for deltas. + // Initialize NCCL directly in C++. Dense TPxDP receives explicit + // orthogonal group metadata; legacy EP keeps its world communicator. let nccl_ep = if ep_world_size > 1 { - unsafe { - std::env::set_var( - "RUSTRAIN_DATA_PARALLEL", - if is_data_parallel { "1" } else { "0" }, - ); - } - let ret = ctx.init_nccl(); - if ret != 0 { - return Err(anyhow!("C++ NCCL init failed (code {})", ret)); + if let Some(topology) = dense_topology.as_ref() { + let tp_color = *topology + .tensor_group(ep_rank)? + .iter() + .min() + .ok_or_else(|| anyhow!("empty TP process group"))?; + let dp_color = *topology + .data_group(ep_rank)? + .iter() + .min() + .ok_or_else(|| anyhow!("empty DP process group"))?; + ctx.init_parallel_nccl( + ep_rank, + ep_world_size, + tp_rank, + tp_size, + tp_color, + dp_rank, + dp_size, + dp_color, + )?; + } else { + let ret = ctx.init_nccl(); + if ret != 0 { + return Err(anyhow!("C++ NCCL init failed (code {})", ret)); + } } tracing::info!( ep_rank, diff --git a/scripts/run_qwen36_native_gdn_tp.sh b/scripts/run_qwen36_native_gdn_tp.sh index d3ae0ebd..e80beed8 100755 --- a/scripts/run_qwen36_native_gdn_tp.sh +++ b/scripts/run_qwen36_native_gdn_tp.sh @@ -2,13 +2,13 @@ set -euo pipefail usage() { - echo "usage: $0 {smoke|bench-single|bench-tp2}" >&2 + echo "usage: $0 {smoke|tpdp-smoke|bench-single|bench-tp2}" >&2 exit 2 } mode="${1:-}" case "$mode" in - smoke|bench-single|bench-tp2) ;; + smoke|tpdp-smoke|bench-single|bench-tp2) ;; *) usage ;; esac @@ -16,8 +16,15 @@ repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" cd "$repo_root" python_bin="${PYTHON:-python3}" -readarray -t torch_config < <( - "$python_bin" - <<'PY' +if [[ -n "${TORCH_INCLUDE_PATH:-}" && -n "${TORCH_LIB_PATH:-}" && -n "${GLIBCXX_USE_CXX11_ABI:-}" ]]; then + torch_include="$TORCH_INCLUDE_PATH" + torch_lib="$TORCH_LIB_PATH" + cxx11_abi="$GLIBCXX_USE_CXX11_ABI" + site_packages="${PYTHON_SITE_PACKAGES:-${torch_include%/torch/include}}" + torch_version="${TORCH_VERSION:-prebuilt}" +else + readarray -t torch_config < <( + "$python_bin" - <<'PY' import pathlib import torch @@ -28,12 +35,13 @@ print(int(torch._C._GLIBCXX_USE_CXX11_ABI)) print(root.parent) print(torch.__version__) PY -) -torch_include="${TORCH_INCLUDE_PATH:-${torch_config[0]}}" -torch_lib="${TORCH_LIB_PATH:-${torch_config[1]}}" -cxx11_abi="${GLIBCXX_USE_CXX11_ABI:-${torch_config[2]}}" -site_packages="${torch_config[3]}" -torch_version="${torch_config[4]}" + ) + torch_include="${torch_config[0]}" + torch_lib="${torch_config[1]}" + cxx11_abi="${torch_config[2]}" + site_packages="${torch_config[3]}" + torch_version="${torch_config[4]}" +fi cuda_home="${CUDA_HOME:-/usr/local/cuda}" cuda_include="${CUDA_INCLUDE_PATH:-$cuda_home/include}" @@ -163,12 +171,17 @@ common_flags=( ) smoke_bin="$native_dir/native_tp_gdn_smoke" +tpdp_smoke_bin="$native_dir/native_tp_dp_smoke" bench_bin="$native_dir/native_tp_gdn_bench" smoke_source=crates/rustrain-qwen3-6/tests/native_tp_gdn_smoke.cpp +tpdp_smoke_source=crates/rustrain-qwen3-6/tests/native_tp_dp_smoke.cpp bench_source=crates/rustrain-qwen3-6/tests/native_tp_gdn_bench.cpp if [[ ! -e "$smoke_bin" || "$smoke_source" -nt "$smoke_bin" || "$kernel_lib" -nt "$smoke_bin" ]]; then g++ "$smoke_source" -o "$smoke_bin" "${common_flags[@]}" fi +if [[ ! -e "$tpdp_smoke_bin" || "$tpdp_smoke_source" -nt "$tpdp_smoke_bin" || "$kernel_lib" -nt "$tpdp_smoke_bin" ]]; then + g++ "$tpdp_smoke_source" -o "$tpdp_smoke_bin" "${common_flags[@]}" +fi if [[ ! -e "$bench_bin" || "$bench_source" -nt "$bench_bin" || "$kernel_lib" -nt "$bench_bin" ]]; then g++ "$bench_source" -o "$bench_bin" "${common_flags[@]}" fi @@ -182,6 +195,11 @@ case "$mode" in TP_SIZE=2 "$python_bin" -m torch.distributed.run --standalone \ --nnodes=1 --nproc-per-node=2 --no-python "$smoke_bin" ;; + tpdp-smoke) + TP_SIZE=2 DP_SIZE=2 RUSTRAIN_DATA_PARALLEL=1 \ + "$python_bin" -m torch.distributed.run --standalone \ + --nnodes=1 --nproc-per-node=4 --no-python "$tpdp_smoke_bin" + ;; bench-single) BENCH_MODE=single WORLD_SIZE=1 RANK=0 LOCAL_RANK=0 "$bench_bin" ;; diff --git a/src/main.rs b/src/main.rs index cb5a3c85..d302506d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -902,9 +902,11 @@ paths=["/tmp/qwen3_6_test.jsonl"] input_ids: ids.clone(), target_mask: mask.clone(), attention_mask: attn.clone(), + batch_size: 1, seq_len, n_total: n_adapters, lora_rank, + adapter_ids: vec![], }) { EpResult::Loss(l) => l, EpResult::Error(e) => { @@ -941,9 +943,11 @@ paths=["/tmp/qwen3_6_test.jsonl"] input_ids: ids.clone(), target_mask: mask.clone(), attention_mask: attn.clone(), + batch_size: 1, seq_len, n_total: n_adapters, lora_rank, + adapter_ids: vec![], }) { EpResult::Loss(l) => { losses.push(l); From b13bbbdf058f56f82b3a87409071c54bbe2221dc Mon Sep 17 00:00:00 2001 From: NolanHo Date: Fri, 17 Jul 2026 16:21:53 +0800 Subject: [PATCH 028/156] feat: add vocabulary tensor parallelism --- .../kernels/qwen3_6_kernels.cpp | 324 ++++++++++++++++-- crates/rustrain-qwen3-6/src/kernel.rs | 81 ++++- crates/rustrain-qwen3-6/src/session.rs | 33 +- .../rustrain-qwen3-6/tests/native_smoke.cpp | 2 +- .../tests/native_tp_dp_smoke.cpp | 29 +- .../tests/native_tp_gdn_bench.cpp | 12 +- .../tests/native_tp_gdn_smoke.cpp | 51 ++- crates/rustrain-server/src/session.rs | 25 +- docs/plans/qwen-lora-megatron-progress.md | 12 +- 9 files changed, 502 insertions(+), 67 deletions(-) diff --git a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp index 191015e7..d6bfab66 100644 --- a/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp +++ b/crates/rustrain-qwen3-6/kernels/qwen3_6_kernels.cpp @@ -136,7 +136,8 @@ struct NcclAllReduceFunction : public torch::autograd::Function