diff --git a/CMakeLists.txt b/CMakeLists.txt index e4b3d3f..9cc576e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -260,6 +260,7 @@ if (BUILD_TESTS) src/testing/test-rope.cu src/testing/test-classifier.cu src/testing/test-gemm.cpp + src/testing/test-transformer-config.cpp ) target_link_libraries(unit-tests PRIVATE llmq-common Catch2::Catch2 CLI11::CLI11) target_compile_options(unit-tests PUBLIC $<$:--expt-relaxed-constexpr> $<$:-lineinfo>) diff --git a/scripts/create_tiny_test_model.py b/scripts/create_tiny_test_model.py new file mode 100644 index 0000000..ea0bbb7 --- /dev/null +++ b/scripts/create_tiny_test_model.py @@ -0,0 +1,155 @@ +# /// script +# requires-python = ">=3.12" +# dependencies = ["torch", "transformers"] +# /// +"""Create a tiny random-weight model in the local HF cache for integration tests. + +The generated model lands under models--test--tiny- in the HF hub cache, so +both `transformers` (with HF_HUB_OFFLINE=1) and llmq can load it as `test/tiny-`. + +The vocabulary is chosen to match one of the tokenizers that `tokenize_data.py` +supports, so an existing tokenized dataset can be reused: + * qwen3 -> data/tiny-shakespeare-qwen + * llama* / mistral -> data/tiny-shakespeare-llama +""" +import argparse +from pathlib import Path + +import torch +import transformers + +# tiny-shakespeare-llama is tokenized with the llama-2 tokenizer +LLAMA_VOCAB = 32000 +QWEN_VOCAB = 151936 + + +def qwen3_config(): + # head_dim != hidden_size / num_attention_heads, to exercise the decoupled path + return transformers.Qwen3Config( + hidden_size=256, + intermediate_size=512, + num_hidden_layers=4, + num_attention_heads=8, + num_key_value_heads=4, + head_dim=64, + max_position_embeddings=2048, + rope_theta=1_000_000.0, + rms_norm_eps=1e-6, + tie_word_embeddings=False, + vocab_size=QWEN_VOCAB, + bos_token_id=151643, + eos_token_id=151645, + torch_dtype=torch.bfloat16, + ) + + +def _llama_config(*, attention_bias: bool, tie_word_embeddings: bool): + # 4 heads over 256 channels gives head_dim 64; the cuDNN attention backend + # rejects the head_dim 32 that 8 heads would produce. + return transformers.LlamaConfig( + hidden_size=256, + intermediate_size=512, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=2048, + rope_theta=10_000.0, + rms_norm_eps=1e-5, + tie_word_embeddings=tie_word_embeddings, + vocab_size=LLAMA_VOCAB, + bos_token_id=1, + eos_token_id=2, + attention_bias=attention_bias, + mlp_bias=False, + torch_dtype=torch.bfloat16, + ) + + +def llama_config(): + return _llama_config(attention_bias=False, tie_word_embeddings=False) + + +def llama_bias_config(): + # negative fixture: attention_bias also biases o_proj, which we cannot represent + return _llama_config(attention_bias=True, tie_word_embeddings=False) + + +def llama_tied_config(): + return _llama_config(attention_bias=False, tie_word_embeddings=True) + + +def llama_rope_scaling_config(): + # negative fixture: Llama-3.1 style scaling, where we only implement plain rope_theta + config = _llama_config(attention_bias=False, tie_word_embeddings=False) + config.rope_scaling = { + "rope_type": "llama3", + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 1024, + } + return config + + +def mistral_config(): + # sliding_window must stay disabled; llmq rejects an active one + return transformers.MistralConfig( + hidden_size=256, + intermediate_size=512, + num_hidden_layers=4, + num_attention_heads=8, + num_key_value_heads=4, + head_dim=64, + max_position_embeddings=2048, + rope_theta=10_000.0, + rms_norm_eps=1e-5, + tie_word_embeddings=False, + vocab_size=LLAMA_VOCAB, + bos_token_id=1, + eos_token_id=2, + sliding_window=None, + torch_dtype=torch.bfloat16, + ) + + +CONFIGS = { + "qwen3": qwen3_config, + "llama": llama_config, + "llama-bias": llama_bias_config, + "llama-tied": llama_tied_config, + "llama-rope-scaling": llama_rope_scaling_config, + "mistral": mistral_config, +} + + +def create(arch: str, seed: int = 42) -> Path: + torch.manual_seed(seed) + config = CONFIGS[arch]() + model = transformers.AutoModelForCausalLM.from_config(config, torch_dtype=torch.bfloat16) + + from huggingface_hub.constants import HF_HUB_CACHE + hub = Path(HF_HUB_CACHE) + base = hub / f"models--test--tiny-{arch}" + snapshot = base / "snapshots" / "main" + snapshot.mkdir(parents=True, exist_ok=True) + (base / "refs").mkdir(exist_ok=True) + (base / "refs" / "main").write_text("main") + + model.save_pretrained(snapshot) + return snapshot + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--arch", choices=[*sorted(CONFIGS), "all"], default="qwen3") + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + arches = sorted(CONFIGS) if args.arch == "all" else [args.arch] + for arch in arches: + snapshot = create(arch, args.seed) + print(f"saved test/tiny-{arch} to {snapshot}") + + +if __name__ == "__main__": + main() diff --git a/src/binding/py_train.cpp b/src/binding/py_train.cpp index 94f6378..b9fc01c 100644 --- a/src/binding/py_train.cpp +++ b/src/binding/py_train.cpp @@ -279,32 +279,21 @@ std::vector> MultiGPUPyTrainer::get_gradients(int using namespace LLamaWeightID; std::vector> result; - // TODO make this work with generalized gradients run_work([&result](sThreadContext& ctx) { - const auto& config = ctx.Model->config(); auto& grads = ctx.Model->grads(); CUDA_CHECK(cudaDeviceSynchronize()); - result.emplace_back("model.embed_tokens.weight", grads.get_non_block_shard(LLamaWeightID::EMBEDDING, nullptr)); - if (!config.TiedWordEmbeddings) { - result.emplace_back("lm_head.weight", grads.get_non_block_shard(LM_HEAD, nullptr)); + for (unsigned id = 0; id < ctx.Model->num_non_block_tensors(); ++id) { + if (const Tensor& tensor = grads.get_non_block_shard(id, nullptr)) { + result.emplace_back(non_block_weight_name(id), tensor); + } } - result.emplace_back("model.norm.weight", grads.get_non_block_shard(LNF_W, nullptr)); - for (int l = 0; l < config.NumLayers; l++) { - - std::string prefix = "model.layers." + std::to_string(l); + for (int l = 0; l < ctx.Model->config().NumLayers; l++) { auto& block = grads.get_block_shard(l, nullptr); - result.emplace_back(prefix + ".self_attn.qkv.weight", block.get_tensor(QKV_W)); - if (block.get_tensor(QKV_B)) - result.emplace_back(prefix + ".self_attn.qkv.bias", block.get_tensor(QKV_B)); - result.emplace_back(prefix + ".self_attn.o_proj.weight", block.get_tensor(ATTO_W)); - if (block.get_tensor(QNORM_W)) - result.emplace_back(prefix + ".self_attn.q_norm.weight", block.get_tensor(QNORM_W)); - if (block.get_tensor(KNORM_W)) - result.emplace_back(prefix + ".self_attn.k_norm.weight", block.get_tensor(KNORM_W)); - result.emplace_back(prefix + ".mlp.up.weight", block.get_tensor(UP_W)); - result.emplace_back(prefix + ".mlp.down_proj.weight", block.get_tensor(DOWN_W)); - result.emplace_back(prefix + ".input_layernorm.weight", block.get_tensor(LN1_W)); - result.emplace_back(prefix + ".post_attention_layernorm.weight", block.get_tensor(LN2_W)); + for (unsigned id = 0; id < block.num_tensors(); ++id) { + if (const Tensor& tensor = block.get_tensor(id)) { + result.emplace_back(block_weight_name(l, id), tensor); + } + } } CUDA_CHECK(cudaDeviceSynchronize()); }, gpu_id); diff --git a/src/binding/python/tests/run.py b/src/binding/python/tests/run.py index 8057a6c..577cd60 100644 --- a/src/binding/python/tests/run.py +++ b/src/binding/python/tests/run.py @@ -100,13 +100,17 @@ def run_training(config: TrainingConfig) -> RunResult: options = _create_options(config) - # Create trainer + if config.batch_size % config.gpus != 0: + raise ValueError(f"batch size {config.batch_size} must be divisible by the number of GPUs ({config.gpus})") + + # Create trainer. `batch_size` is the per-GPU micro-batch; `step()` takes the global batch, + # so results are comparable across GPU counts. trainer = pyllmq.LLMQTrainer.from_pretrained( name=config.model, - ngpu=1, + ngpu=config.gpus, dtype=config.model_dtype, options=options, - batch_size=config.batch_size, + batch_size=config.batch_size // config.gpus, seq_len=config.seq_len, grad_accum=config.grad_accumulation, memcpy_all_gather=config.memcpy_all_gather, diff --git a/src/binding/python/tests/torch_reference.py b/src/binding/python/tests/torch_reference.py index 5efa46f..40efe83 100644 --- a/src/binding/python/tests/torch_reference.py +++ b/src/binding/python/tests/torch_reference.py @@ -48,13 +48,17 @@ def torch_grad_one_step(config: TrainingConfig): def llmq_grad_one_step(config: TrainingConfig): options = _create_options(config) - # Create trainer + if config.batch_size % config.gpus != 0: + raise ValueError(f"batch size {config.batch_size} must be divisible by the number of GPUs ({config.gpus})") + + # Create trainer. `batch_size` is the per-GPU micro-batch; `step()` takes the global batch, + # so the total amount of data (and thus the gradient) is independent of the GPU count. trainer = pyllmq.LLMQTrainer.from_pretrained( name=config.model, ngpu=config.gpus, dtype=config.model_dtype, options=options, - batch_size=config.batch_size, + batch_size=config.batch_size // config.gpus, seq_len=config.seq_len, grad_accum=config.grad_accumulation, memcpy_all_gather=config.memcpy_all_gather, @@ -75,7 +79,15 @@ def llmq_grad_one_step(config: TrainingConfig): for j in range(config.grad_accumulation): train_loader.load_batch(in_tokens, out_tokens) trainer.step(in_tokens, out_tokens) - return {k: torch.from_dlpack(v).cpu().to(torch.float32).numpy() for k, v in trainer.get_gradients(0).items()} + + # each GPU returns its (dim-0) shard of the gradients; concatenate to reconstruct the + # full tensors. GPUs whose shard of a tensor is empty do not report it at all. + grads = {} + for g in range(config.gpus): + for k, v in trainer.get_gradients(g).items(): + arr = torch.from_dlpack(v).cpu().to(torch.float32).numpy().flatten() + grads.setdefault(k, []).append(arr) + return {k: np.concatenate(parts) for k, parts in grads.items()} def compare_single_step(config, file=None): diff --git a/src/kernels/adamw.cu b/src/kernels/adamw.cu index d4c45bb..c4ae79e 100644 --- a/src/kernels/adamw.cu +++ b/src/kernels/adamw.cu @@ -122,6 +122,7 @@ __global__ void adamw_kernel(floatX* params_memory, const floatX* grads_memory, if(threadIdx.x == 0) { block_abs_max = 0.f; } + __syncthreads(); float thread_abs_max = 0.0f; vec_x_t p_new = adamw_update(params_memory, diff --git a/src/kernels/gemm_mma.cu b/src/kernels/gemm_mma.cu index 6349542..de3f90a 100644 --- a/src/kernels/gemm_mma.cu +++ b/src/kernels/gemm_mma.cu @@ -75,11 +75,13 @@ __device__ void gemm_mma_tn_impl(nv_bfloat16* __restrict__ out, const uint4* g_ptr; uint4* s_ptr; + // 64 bit on purpose: `m * k` (and `m * n` in the epilogue) exceeds INT_MAX at ordinary + // large-model shapes, and the wrapped product becomes an illegal address. if(wid < 2) { - g_ptr = reinterpret_cast(a) + (bi + wid) * TI * stride; + g_ptr = reinterpret_cast(a) + (long)(bi + wid) * TI * stride; s_ptr = input_tiles + wid * ROW_OFFSET; } else { - g_ptr = reinterpret_cast(b) + (bj + wid - 2) * TJ * stride; + g_ptr = reinterpret_cast(b) + (long)(bj + wid - 2) * TJ * stride; s_ptr = input_tiles + (wid - 2) * ROW_OFFSET + DEPTH * PIPE_OFFSET; } @@ -211,24 +213,26 @@ __device__ void gemm_mma_tn_impl(nv_bfloat16* __restrict__ out, __syncwarp(); int c = threadIdx.x % 2; int r = threadIdx.x / 2; + // 64 bit: (row index) * n overflows int once m * n exceeds 2^31 + long out_offset = ((long)(i + ii) * TI + r) * n + (j + jj) * TJ + 8 * c; if(accumulate) { - auto old = GenericVector::load(out + ((i + ii) * TI + r) * n + (j + jj) * TJ + 8 * c); + auto old = GenericVector::load(out + out_offset); auto upd = GenericVector::load(out_shared + (c + 2 * r) * 8); for(int l = 0; l < 8; ++l) { old[l] += upd[l]; } - old.store(out + ((i + ii) * TI + r) * n + (j + jj) * TJ + 8 * c); + old.store(out + out_offset); } else if (bias != nullptr) { auto old = GenericVector::load(bias + (j + jj) * TJ + 8 * c); auto upd = GenericVector::load(out_shared + (c + 2 * r) * 8); for(int l = 0; l < 8; ++l) { old[l] += (nv_bfloat16)upd[l]; } - old.store(out + ((i + ii) * TI + r) * n + (j + jj) * TJ + 8 * c); + old.store(out + out_offset); } else { uint4 load = reinterpret_cast(out_shared)[c + 2 * r]; - *reinterpret_cast(out + ((i + ii) * TI + r) * n + (j + jj) * TJ + 8 * c) = load; + *reinterpret_cast(out + out_offset) = load; } } } diff --git a/src/kernels/rmsnorm.cu b/src/kernels/rmsnorm.cu index 381be75..a64e8b9 100644 --- a/src/kernels/rmsnorm.cu +++ b/src/kernels/rmsnorm.cu @@ -25,7 +25,7 @@ __device__ void rmsnorm_forward_kernel(floatX* __restrict__ out, float* __restri // load weights into shared memory // do this before we allow any threads to exit! - extern __shared__ char* params[]; + extern __shared__ __align__(16) unsigned char params[]; __shared__ float block_abs_max; // load128/store128 sometimes generated multiple instructions when the types here were floatX*, so // let's keep everything as x128 @@ -98,7 +98,7 @@ __device__ void fused_residual_rmsnorm_forward_kernel(floatX* residual, floatX* // load weights and biases into shared memory // do this before we allow any threads to exit! - extern __shared__ char* params[]; + extern __shared__ __align__(16) unsigned char params[]; __shared__ float block_abs_max; // load128/store128 sometimes generated multiple instructions when the types here were floatX*, so // let's keep everything as x128 diff --git a/src/models/llama_optimizer.cpp b/src/models/llama_optimizer.cpp index e66b245..7a3def0 100644 --- a/src/models/llama_optimizer.cpp +++ b/src/models/llama_optimizer.cpp @@ -23,28 +23,20 @@ struct OptStateWrapper : ITensorContainer { }; void OptStateWrapper::iterate_tensors(const std::function& callback) { - auto cb = [&callback](std::string name, const Tensor& t) { - if (t) { - callback(std::move(name), t); + using namespace LLamaWeightID; + for(unsigned id = 0; id < NonBlock->num_tensors(); ++id) { + if(const Tensor& tensor = NonBlock->get_tensor(id)) { + callback(non_block_weight_name(id), tensor); } - }; - - cb("model.embed_tokens.weight", NonBlock->get_tensor(LLamaWeightID::EMBEDDING)); - cb("lm_head.weight", NonBlock->get_tensor(LLamaWeightID::LM_HEAD)); - cb("model.norm.weight", NonBlock->get_tensor(LLamaWeightID::LNF_W)); + } for(int i = 0; i < Blocks->size(); i++) { auto& layer = Blocks->at(i); - std::string prefix = "model.layers." + std::to_string(i); - cb(prefix + ".self_attn.qkv.weight", layer.get_tensor(LLamaWeightID::QKV_W)); - cb(prefix + ".self_attn.qkv.bias", layer.get_tensor(LLamaWeightID::QKV_B)); - cb(prefix + ".self_attn.o_proj.weight", layer.get_tensor(LLamaWeightID::ATTO_W)); - cb(prefix + ".self_attn.q_norm.weight", layer.get_tensor(LLamaWeightID::QNORM_W)); - cb(prefix + ".self_attn.k_norm.weight", layer.get_tensor(LLamaWeightID::KNORM_W)); - cb(prefix + ".mlp.up.weight", layer.get_tensor(LLamaWeightID::UP_W)); - cb(prefix + ".mlp.down_proj.weight", layer.get_tensor(LLamaWeightID::DOWN_W)); - cb(prefix + ".input_layernorm.weight", layer.get_tensor(LLamaWeightID::LN1_W)); - cb(prefix + ".post_attention_layernorm.weight", layer.get_tensor(LLamaWeightID::LN2_W)); + for(unsigned id = 0; id < layer.num_tensors(); ++id) { + if(const Tensor& tensor = layer.get_tensor(id)) { + callback(block_weight_name(i, id), tensor); + } + } } } diff --git a/src/models/llama_weights.cpp b/src/models/llama_weights.cpp index 756f0ec..c53c64d 100644 --- a/src/models/llama_weights.cpp +++ b/src/models/llama_weights.cpp @@ -4,6 +4,8 @@ #include "llama_weights.h" +#include + #include "kernels/kernels.h" #include "llama_model.h" #include "llama_run_state.h" @@ -229,32 +231,32 @@ LLamaWeightsManager::~LLamaWeightsManager() { } } +// Each tensor gets two floats of stats (abs-max and scale), assigned in tensor-ID order. +static constexpr int STATS_PER_TENSOR = 2; +static constexpr int NON_BLOCK_STATS = STATS_PER_TENSOR * LLamaWeightID::NUM_NON_BLOCK_TENSORS; +static constexpr int BLOCK_STATS = STATS_PER_TENSOR * LLamaWeightID::NUM_BLOCK_TENSORS; + void LLamaWeightsManager::setup_scales(TensorAllocator& alloc) { int layers = mMaster.Blocks.size(); - mAbsMaxes = alloc.allocate(ETensorDType::FP32, "abs_maxes", EAllocationType::ON_DEVICE, {6 + layers * 18}); + mAbsMaxes = alloc.allocate(ETensorDType::FP32, "abs_maxes", EAllocationType::ON_DEVICE, + {NON_BLOCK_STATS + layers * BLOCK_STATS}); float* abs_maxes = mAbsMaxes.get(); - mMaster.NonBlocks.Embeddings.Stats = abs_maxes + 0; - mMaster.NonBlocks.LNF_w.Stats = abs_maxes + 2; - mMaster.NonBlocks.LMHead.Stats = abs_maxes + 4; + for(unsigned id = 0; id < LLamaWeightID::NUM_NON_BLOCK_TENSORS; ++id) { + mMaster.NonBlocks.get_tensor(id).Stats = abs_maxes + STATS_PER_TENSOR * id; + } for(int i = 0; i < layers; ++i) { - float* a = abs_maxes + 6 + i * 14; - mMaster.Blocks[i].Attn_QKV_w.Stats = a + 0; - mMaster.Blocks[i].Attn_Out_w.Stats = a + 2; - mMaster.Blocks[i].MLP_Up_w.Stats = a + 4; - mMaster.Blocks[i].MLP_Down_w.Stats = a + 6; - mMaster.Blocks[i].Attn_QKV_b.Stats = a + 8; - mMaster.Blocks[i].LN1_w.Stats = a + 10; - mMaster.Blocks[i].LN2_w.Stats = a + 12; - mMaster.Blocks[i].QNorm_w.Stats = a + 14; - mMaster.Blocks[i].KNorm_w.Stats = a + 16; + float* a = abs_maxes + NON_BLOCK_STATS + i * BLOCK_STATS; + for(unsigned id = 0; id < LLamaWeightID::NUM_BLOCK_TENSORS; ++id) { + mMaster.Blocks[i].get_tensor(id).Stats = a + STATS_PER_TENSOR * id; + } } } std::pair LLamaWeightsManager::get_scales_for_block(int layer_idx) { float* abs_maxes = mAbsMaxes.get(); - float* begin = abs_maxes + 6 + layer_idx * 18; - return {begin + 0, begin + 18}; + float* begin = abs_maxes + NON_BLOCK_STATS + layer_idx * BLOCK_STATS; + return {begin + 0, begin + BLOCK_STATS}; } @@ -552,35 +554,47 @@ void LLamaWeightsManager::release_head(cudaStream_t stream) { release_status(mEmbStatus, mHeadID, stream); } +namespace LLamaWeightID { +std::string block_weight_name(int layer, unsigned id) { + constexpr std::array suffixes = { + ".input_layernorm.weight", // LN1_W + ".post_attention_layernorm.weight", // LN2_W + ".self_attn.qkv.weight", // QKV_W + ".self_attn.qkv.bias", // QKV_B + ".self_attn.o_proj.weight", // ATTO_W + ".mlp.up.weight", // UP_W + ".mlp.down_proj.weight", // DOWN_W + ".self_attn.q_norm.weight", // QNORM_W + ".self_attn.k_norm.weight", // KNORM_W + }; + return "model.layers." + std::to_string(layer) + suffixes.at(id); +} + +std::string non_block_weight_name(unsigned id) { + constexpr std::array names = { + "model.embed_tokens.weight", // EMBEDDING + "lm_head.weight", // LM_HEAD + "model.norm.weight", // LNF_W + }; + return names.at(id); +} +} + void sLLamaWeights::iterate_tensors(const std::function& callback) { - callback("model.embed_tokens.weight", NonBlocks.Embeddings); + using namespace LLamaWeightID; + callback(non_block_weight_name(EMBEDDING), NonBlocks.Embeddings); if(NonBlocks.LMHead.Data != NonBlocks.Embeddings.Data) { - callback("lm_head.weight", NonBlocks.LMHead); + callback(non_block_weight_name(LM_HEAD), NonBlocks.LMHead); } - callback("model.norm.weight", NonBlocks.LNF_w); + callback(non_block_weight_name(LNF_W), NonBlocks.LNF_w); - const auto& Layers = Blocks; - for(int i = 0; i < Layers.size(); i++) { - auto& layer = Layers[i]; - const Tensor& qkv_w = layer.Attn_QKV_w; - const Tensor& up_proj = layer.MLP_Up_w; - std::string prefix = "model.layers." + std::to_string(i); - callback(prefix + ".self_attn.qkv.weight", qkv_w); - if (layer.Attn_QKV_b) { - callback(prefix + ".self_attn.qkv.bias", layer.Attn_QKV_b); - } - if (layer.QNorm_w) { - callback(prefix + ".self_attn.q_norm.weight", layer.QNorm_w); - } - if (layer.KNorm_w) { - callback(prefix + ".self_attn.k_norm.weight", layer.KNorm_w); + for(int i = 0; i < Blocks.size(); i++) { + auto& layer = Blocks[i]; + for(unsigned id = 0; id < layer.num_tensors(); ++id) { + if(const TensorShard& tensor = layer.get(id)) { + callback(block_weight_name(i, id), tensor); + } } - - callback(prefix + ".self_attn.o_proj.weight", layer.Attn_Out_w); - callback(prefix + ".mlp.up.weight", up_proj); - callback(prefix + ".mlp.down_proj.weight", layer.MLP_Down_w); - callback(prefix + ".input_layernorm.weight", layer.LN1_w); - callback(prefix + ".post_attention_layernorm.weight", layer.LN2_w); } } diff --git a/src/models/llama_weights.h b/src/models/llama_weights.h index 9ecf717..d46c37b 100644 --- a/src/models/llama_weights.h +++ b/src/models/llama_weights.h @@ -32,9 +32,16 @@ namespace LLamaWeightID { inline constexpr unsigned DOWN_W = 6; inline constexpr unsigned QNORM_W = 7; inline constexpr unsigned KNORM_W = 8; + inline constexpr unsigned NUM_BLOCK_TENSORS = 9; inline constexpr unsigned EMBEDDING = 0; inline constexpr unsigned LM_HEAD = 1; inline constexpr unsigned LNF_W = 2; + inline constexpr unsigned NUM_NON_BLOCK_TENSORS = 3; + + //! HF-checkpoint name of the block tensor `id` in layer `layer` + std::string block_weight_name(int layer, unsigned id); + //! HF-checkpoint name of the non-block tensor `id` + std::string non_block_weight_name(unsigned id); }; template @@ -49,7 +56,7 @@ struct sLLamaBlockWeights : public SimpleTensorContainer { TTensor QNorm_w; // Hd; optional TTensor KNorm_w; // Hd; optional - std::size_t num_tensors() const noexcept override { return 9; } + std::size_t num_tensors() const noexcept override { return LLamaWeightID::NUM_BLOCK_TENSORS; } const Tensor& get_tensor(std::size_t idx) const override { using namespace LLamaWeightID; @@ -68,6 +75,11 @@ struct sLLamaBlockWeights : public SimpleTensorContainer { } } + //! Like `get_tensor`, but preserves the field type; valid because all fields are `TTensor`s. + const TTensor& get(std::size_t idx) const { + return static_cast(get_tensor(idx)); + } + using SimpleTensorContainer::get_tensor; }; @@ -77,7 +89,7 @@ struct sLLamaNonBlockWeights : public SimpleTensorContainer { TTensor LMHead; // V, C TTensor LNF_w; // C - std::size_t num_tensors() const noexcept override { return 3; } + std::size_t num_tensors() const noexcept override { return LLamaWeightID::NUM_NON_BLOCK_TENSORS; } const Tensor& get_tensor(std::size_t idx) const override { using namespace LLamaWeightID; diff --git a/src/testing/test-gemm.cpp b/src/testing/test-gemm.cpp index 8a055e0..4377397 100644 --- a/src/testing/test-gemm.cpp +++ b/src/testing/test-gemm.cpp @@ -260,3 +260,50 @@ TEST_CASE("matmul fp8 x fp8 -> bfloat16", "[gemm][fp8]") { } run_test<__nv_fp8_e4m3, __nv_fp8_e4m3, nv_bfloat16>(m, n, k, 4.0f / k, accumulate, bias); } + +// The epilogue indexes `out` as (row * n), which overflowed int once m * n passed 2^31 and +// faulted rather than returning wrong results. Needs ~5 GB, so it skips where that does not fit. +TEST_CASE("matmul indexes outputs larger than INT_MAX", "[gemm][fp8][large]") { + const long m = 16384; + const long n = 147456; // m * n = 2.42e9, comfortably past INT_MAX + const long k = 128; + const std::size_t out_bytes = (std::size_t)m * n * sizeof(nv_bfloat16); + + std::size_t free_mem = 0, total_mem = 0; + CUDA_CHECK(cudaMemGetInfo(&free_mem, &total_mem)); + if(free_mem < out_bytes + (std::size_t)256 * 1024 * 1024) { + SUCCEED("not enough device memory for the >INT_MAX output test"); + return; + } + REQUIRE((double)m * n > 2147483647.0); + + __nv_fp8_e4m3 *a, *b; + nv_bfloat16* c; + float *scale_a, *scale_b; + CUDA_CHECK(cudaMalloc(&a, (std::size_t)m * k)); + CUDA_CHECK(cudaMalloc(&b, (std::size_t)n * k)); + CUDA_CHECK(cudaMalloc(&c, out_bytes)); + CUDA_CHECK(cudaMalloc(&scale_a, sizeof(float))); + CUDA_CHECK(cudaMalloc(&scale_b, sizeof(float))); + CUDA_CHECK(cudaMemset(a, 0x38, (std::size_t)m * k)); + CUDA_CHECK(cudaMemset(b, 0x38, (std::size_t)n * k)); + CUDA_CHECK(cudaMemset(c, 0, out_bytes)); + float one = 1.f; + CUDA_CHECK(cudaMemcpy(scale_a, &one, sizeof(float), cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(scale_b, &one, sizeof(float), cudaMemcpyHostToDevice)); + + cublasLtHandle_t handle = create_cublaslt_handle(); + std::byte* workspace; + std::size_t workspace_size = 32 * 1024 * 1024; + CUDA_CHECK(cudaMalloc(&workspace, workspace_size)); + + matmul(c, a, b, (const nv_bfloat16*)nullptr, scale_a, scale_b, handle, workspace, workspace_size, + (int)m, (int)n, (int)k, EMMTranspose::TN, false, nullptr, EMatmulBackend::Custom); + CHECK(cudaDeviceSynchronize() == cudaSuccess); + + CUDA_CHECK(cudaFree(a)); CUDA_CHECK(cudaFree(b)); + CUDA_CHECK(cudaFree(c)); + CUDA_CHECK(cudaFree(scale_a)); CUDA_CHECK(cudaFree(scale_b)); + CUDA_CHECK(cudaFree(workspace)); + cublasLtDestroy(handle); +} diff --git a/src/testing/test-transformer-config.cpp b/src/testing/test-transformer-config.cpp new file mode 100644 index 0000000..0bb48e0 --- /dev/null +++ b/src/testing/test-transformer-config.cpp @@ -0,0 +1,152 @@ +// Copyright (c) 2026, IST Austria, developed by Erik Schultheis +// SPDX-License-Identifier: Apache-2.0 +// + +#include "training/transformer_config.h" + +#include +#include + +#include +#include + +namespace { + +//! a config.json in the temp directory that removes itself again +class ScratchFile { +public: + ScratchFile() { + static int counter = 0; + mPath = std::filesystem::temp_directory_path() / + ("llmq-test-config-" + std::to_string(counter++) + ".json"); + std::filesystem::remove(mPath); + } + ~ScratchFile() { std::filesystem::remove(mPath); } + + [[nodiscard]] const char* c_str() const { return mPath.c_str(); } + +private: + std::filesystem::path mPath; +}; + +TransformerConfig make_config(TransformerConfig::EArchitecture arch) { + TransformerConfig config{}; + config.Architecture = arch; + config.BosTokenId = 1; + config.EosTokenId = 2; + config.HiddenSize = 256; + config.IntermediateSize = 512; + config.VocabSize = 32000; + config.NumQueryHeads = 4; + config.NumKeyValHeads = 2; + config.NumLayers = 4; + config.MaxPositionEmbeddings = 2048; + config.RopeTheta = 10000.f; + config.RmsNormEps = 1e-5f; + config.TiedWordEmbeddings = false; + config.UseQKVBias = arch == TransformerConfig::QWEN2; + config.UseQKNorm = arch == TransformerConfig::QWEN3; + return config; +} + +//! save `config`, then patch `overrides` on top, to get a config.json that +//! save_transformer_config would never emit by itself. +void save_with(const TransformerConfig& config, const ScratchFile& file, const nlohmann::json& overrides) { + save_transformer_config(config, file.c_str()); + + std::ifstream in(file.c_str()); + auto json = nlohmann::json::parse(in); + in.close(); + for(const auto& [key, value] : overrides.items()) { + json[key] = value; + } + std::ofstream out(file.c_str()); + out << json.dump(4); +} + +} + +TEST_CASE("transformer config survives a save/load round trip") { + for(auto arch : {TransformerConfig::LLAMA, TransformerConfig::MISTRAL, + TransformerConfig::QWEN2, TransformerConfig::QWEN3}) { + auto original = make_config(arch); + ScratchFile file; + save_transformer_config(original, file.c_str()); + auto loaded = load_transformer_config(file.c_str(), original.DType); + + INFO("architecture " << original.model_name()); + CHECK(loaded.Architecture == original.Architecture); + CHECK(loaded.HiddenSize == original.HiddenSize); + CHECK(loaded.IntermediateSize == original.IntermediateSize); + CHECK(loaded.VocabSize == original.VocabSize); + CHECK(loaded.NumQueryHeads == original.NumQueryHeads); + CHECK(loaded.NumKeyValHeads == original.NumKeyValHeads); + CHECK(loaded.NumLayers == original.NumLayers); + CHECK(loaded.head_size() == original.head_size()); + CHECK(loaded.RopeTheta == original.RopeTheta); + CHECK(loaded.TiedWordEmbeddings == original.TiedWordEmbeddings); + CHECK(loaded.UseQKVBias == original.UseQKVBias); + CHECK(loaded.UseQKNorm == original.UseQKNorm); + } +} + +TEST_CASE("a decoupled head_dim round trips") { + auto original = make_config(TransformerConfig::QWEN3); + original.HeadDim = 128; // HiddenSize / NumQueryHeads would be 256 / 4 == 64 + REQUIRE(original.head_size() == 128); + + ScratchFile file; + save_transformer_config(original, file.c_str()); + auto loaded = load_transformer_config(file.c_str(), original.DType); + CHECK(loaded.HeadDim == 128); + CHECK(loaded.head_size() == 128); +} + +TEST_CASE("saving a llama-family config with QKV biases is refused") { + // neither spelling of attention_bias is correct for a q/k/v-only bias, so we must + // not write a config.json that we would refuse to load again + for(auto arch : {TransformerConfig::LLAMA, TransformerConfig::MISTRAL}) { + auto config = make_config(arch); + config.UseQKVBias = true; + ScratchFile file; + CHECK_THROWS(save_transformer_config(config, file.c_str())); + } +} + +TEST_CASE("configs we cannot represent are rejected on load") { + auto base = make_config(TransformerConfig::LLAMA); + + auto rejects = [&](const nlohmann::json& overrides) { + ScratchFile file; + save_with(base, file, overrides); + CHECK_THROWS(load_transformer_config(file.c_str(), base.DType)); + }; + + rejects({{"attention_bias", true}}); + rejects({{"mlp_bias", true}}); + rejects({{"hidden_act", "gelu"}}); + rejects({{"attention_dropout", 0.1f}}); + rejects({{"rope_scaling", {{"rope_type", "llama3"}, {"factor", 8.0}}}}); + rejects({{"sliding_window", 1024}}); + rejects({{"architectures", nlohmann::json::array({"GemmaForCausalLM"})}}); +} + +TEST_CASE("a sliding window is only rejected when it is active") { + auto base = make_config(TransformerConfig::QWEN2); + + ScratchFile inactive; + save_with(base, inactive, {{"sliding_window", 32768}, {"use_sliding_window", false}}); + CHECK_NOTHROW(load_transformer_config(inactive.c_str(), base.DType)); + + ScratchFile active; + save_with(base, active, {{"sliding_window", 32768}, {"use_sliding_window", true}}); + CHECK_THROWS(load_transformer_config(active.c_str(), base.DType)); +} + +TEST_CASE("a null-valued key is treated as absent") { + // Mistral writes `attention_bias: null`, which value() would throw on + auto base = make_config(TransformerConfig::MISTRAL); + ScratchFile file; + save_with(base, file, {{"attention_bias", nullptr}, {"rope_scaling", nullptr}}); + CHECK_NOTHROW(load_transformer_config(file.c_str(), base.DType)); +} diff --git a/src/training/transformer_config.cpp b/src/training/transformer_config.cpp index f20f4ec..4fb9dc9 100644 --- a/src/training/transformer_config.cpp +++ b/src/training/transformer_config.cpp @@ -25,6 +25,8 @@ TransformerConfig load_transformer_config(const char* file_name, ETensorDType dt TransformerConfig::EArchitecture arch_id; if(archs.front() == "LlamaForCausalLM") { arch_id = TransformerConfig::LLAMA; + } else if(archs.front() == "MistralForCausalLM") { + arch_id = TransformerConfig::MISTRAL; } else if(archs.front() == "Qwen2ForCausalLM") { arch_id = TransformerConfig::QWEN2; } else if(archs.front() == "Qwen3ForCausalLM") { @@ -60,9 +62,48 @@ TransformerConfig load_transformer_config(const char* file_name, ETensorDType dt result.RmsNormEps = result.Architecture == TransformerConfig::LLAMA ? 1e-5 : 1e-6; } + // value() throws on a present-but-null key, and Mistral writes `attention_bias: null`. + auto get_or = [&](const char* key, auto fallback) { + auto it = config_json.find(key); + if(it == config_json.end() || it->is_null()) { + return fallback; + } + return it->template get(); + }; + result.UseQKNorm = arch_id == TransformerConfig::QWEN3; + // Qwen2 biases q/k/v only and carries no flag for it; that is not HF's + // attention_bias, which also biases o_proj, so the flag does not feed in here. result.UseQKVBias = arch_id == TransformerConfig::QWEN2; + // Anything we cannot represent exactly has to fail here: silently training a model + // that differs from the checkpoint only shows up as degraded quality later. + auto reject = [&](std::string_view key, std::string_view value) { + throw std::runtime_error(fmt::format("config {}: cannot represent '{}' = {}", file_name, key, value)); + }; + + if(auto it = config_json.find("rope_scaling"); it != config_json.end() && !it->is_null()) { + reject("rope_scaling", it->dump()); + } + if(get_or("mlp_bias", false)) { + reject("mlp_bias", "true"); + } + if(get_or("attention_bias", false)) { + reject("attention_bias", "true, which implies an o_proj bias we have no tensor for"); + } + if(auto act = get_or("hidden_act", std::string{"silu"}); act != "silu") { + reject("hidden_act", act); + } + if(float dropout = get_or("attention_dropout", 0.f); dropout != 0.f) { + reject("attention_dropout", fmt::format("{}", dropout)); + } + // Qwen's sliding_window stays inactive unless use_sliding_window is set; Mistral has + // no such flag, so a non-null window there is always active. + if(auto it = config_json.find("sliding_window"); it != config_json.end() && !it->is_null() + && get_or("use_sliding_window", true)) { + reject("sliding_window", it->dump()); + } + return result; } @@ -74,6 +115,8 @@ TransformerConfig load_transformer_config(const char* file_name, ETensorDType dt return "Qwen2"; case TransformerConfig::LLAMA: return "LLaMA"; + case TransformerConfig::MISTRAL: + return "Mistral"; default: throw std::logic_error("Unknown architecture"); } @@ -92,6 +135,8 @@ void save_transformer_config(const TransformerConfig& config, const char* file_n archs = {"Qwen3ForCausalLM"}; } else if (config.Architecture == TransformerConfig::LLAMA) { archs = {"LlamaForCausalLM"}; + } else if (config.Architecture == TransformerConfig::MISTRAL) { + archs = {"MistralForCausalLM"}; } nlohmann::json config_json; @@ -123,10 +168,22 @@ void save_transformer_config(const TransformerConfig& config, const char* file_n config_json["sliding_window"] = config.MaxPositionEmbeddings; config_json["use_sliding_window"] = false; config_json["use_mrope"] = false; - } else if (config.Architecture == TransformerConfig::LLAMA) { - config_json["model_type"] = "llama"; + } else if (config.Architecture == TransformerConfig::LLAMA || config.Architecture == TransformerConfig::MISTRAL) { + bool is_llama = config.Architecture == TransformerConfig::LLAMA; + config_json["model_type"] = is_llama ? "llama" : "mistral"; + // A q/k/v-only bias has no faithful spelling here: true would promise an o_proj + // bias we lack, false would disclaim biases we have. + if(config.UseQKVBias) { + throw std::runtime_error(fmt::format( + "cannot save a {} config with QKV biases: HF's attention_bias also implies an o_proj bias", + config.model_name())); + } config_json["attention_bias"] = false; config_json["mlp_bias"] = false; + if(!is_llama) { + // we only support Mistral variants that keep the window disabled + config_json["sliding_window"] = nullptr; + } } file << config_json.dump(4); diff --git a/src/training/transformer_config.h b/src/training/transformer_config.h index 9c93c13..ff54603 100644 --- a/src/training/transformer_config.h +++ b/src/training/transformer_config.h @@ -15,6 +15,7 @@ struct TransformerConfig { enum EArchitecture { LLAMA, + MISTRAL, QWEN2, QWEN3, } Architecture; diff --git a/src/utilities/comm.cpp b/src/utilities/comm.cpp index cebaa6e..e9e6666 100644 --- a/src/utilities/comm.cpp +++ b/src/utilities/comm.cpp @@ -182,7 +182,7 @@ void NCCLCommunicator::schedule_reduce_scatter(Tensor& tensor) { mCmdBuf->Commands.emplace_back(CommandBuffer::ScatterReduce{.DType = tensor.DType, .Tensor = tensor.Data, .Elements = tensor.nelem()}); } -void NCCLCommunicator::schedule_all_gather(const TensorShard& src, Tensor& tgt) { +void NCCLCommunicator::schedule_all_gather(const Tensor& src, Tensor& tgt) { if (src.Data == nullptr) { throw std::runtime_error("gather: Source tensor is null"); } diff --git a/src/utilities/comm.h b/src/utilities/comm.h index e293dbc..9095853 100644 --- a/src/utilities/comm.h +++ b/src/utilities/comm.h @@ -21,7 +21,6 @@ namespace std } struct Tensor; -struct TensorShard; typedef struct ncclComm* ncclComm_t; typedef struct CUevent_st* cudaEvent_t; @@ -47,7 +46,7 @@ class NCCLCommunicator { void begin_transaction(cudaEvent_t ready); void begin_transaction(cudaStream_t wait_for_stream); void schedule_reduce_scatter(Tensor& tensor); - void schedule_all_gather(const TensorShard& src, Tensor& tgt); + void schedule_all_gather(const Tensor& src, Tensor& tgt); // like all-to-all, except the local shard will *not* be preserved, and results will be shifted cyclically void schedule_destructive_all_to_all(Tensor& tensor); void execute_transaction(cudaEvent_t signal); diff --git a/test/test_architectures.py b/test/test_architectures.py new file mode 100644 index 0000000..bb66395 --- /dev/null +++ b/test/test_architectures.py @@ -0,0 +1,70 @@ +"""Gradient-parity checks for the model architectures llmq claims to support. + +Each case builds a tiny random-weight model with `scripts/create_tiny_test_model.py` +and compares every parameter gradient after one forward+backward against +transformers; forward-only agreement would not catch a mis-wired backward. + +Needs a GPU and tokenized data: + uv run --extra scripts python scripts/tokenize_data.py --dataset tiny-shakespeare --model llama + uv run --extra scripts python scripts/tokenize_data.py --dataset tiny-shakespeare --model qwen +""" +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +REPO = Path(__file__).resolve().parents[1] +REFERENCE = REPO / "src" / "binding" / "python" / "tests" / "torch_reference.py" +ENV = {**os.environ, "HF_HUB_OFFLINE": "1"} + +# fixture architecture -> tokenizer whose vocabulary it reuses +SUPPORTED = { + "llama": "llama", + "llama-tied": "llama", + "mistral": "llama", + "qwen3": "qwen", +} + +# fixtures llmq must refuse, and the text the error has to mention +REFUSED = { + "llama-bias": "attention_bias", + "llama-rope-scaling": "rope_scaling", +} + + +@pytest.fixture(scope="session") +def tiny_models(): + script = REPO / "scripts" / "create_tiny_test_model.py" + done = subprocess.run([sys.executable, str(script), "--arch", "all"], + capture_output=True, text=True, env=ENV) + if done.returncode != 0: + pytest.skip(f"could not create tiny models:\n{done.stderr}") + + +def _compare(arch: str, tokenizer: str): + train_file = REPO / "data" / f"tiny-shakespeare-{tokenizer}" / "train.bin" + if not train_file.exists(): + pytest.skip(f"missing {train_file}; see this module's docstring") + return subprocess.run( + [sys.executable, str(REFERENCE), + "--model", f"test/tiny-{arch}", "--train-file", str(train_file), + "--seq-len", "256", "--grad-accum", "2", + # fp32 keeps this about the architecture rather than quantization + "--model-dtype", "fp32", "--matmul-dtype", "fp32", "--gpus", "1"], + capture_output=True, text=True, env=ENV) + + +@pytest.mark.parametrize("arch,tokenizer", sorted(SUPPORTED.items())) +def test_matches_transformers(tiny_models, arch, tokenizer): + done = _compare(arch, tokenizer) + assert done.returncode == 0, f"gradients diverge:\n{done.stdout[-3000:]}{done.stderr[-2000:]}" + + +@pytest.mark.parametrize("arch,expected", sorted(REFUSED.items())) +def test_rejects_unrepresentable(tiny_models, arch, expected): + done = _compare(arch, "llama") + assert done.returncode != 0, "should have been refused, but loaded fine" + assert expected in done.stdout + done.stderr, \ + f"refused, but the error never mentions {expected!r}:\n{done.stderr[-2000:]}" diff --git a/test/test_multi_gpu.py b/test/test_multi_gpu.py new file mode 100644 index 0000000..ef965dc --- /dev/null +++ b/test/test_multi_gpu.py @@ -0,0 +1,52 @@ +"""End-to-end multi-GPU gradient parity tests. + +Each case runs torch_reference.py in a subprocess (fresh NCCL state per case) on the +tiny-qwen3 test model, comparing llmq gradients on 2 GPUs against a single-GPU torch +reference. Requires 2 GPUs and the tiny test model +(create with `scripts/create_tiny_test_model.py --arch qwen3`). +""" +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +REPO_ROOT = Path(__file__).resolve().parent.parent +SCRIPT = REPO_ROOT / "src" / "binding" / "python" / "tests" / "torch_reference.py" + + +def _tiny_model_available() -> bool: + try: + from huggingface_hub.constants import HF_HUB_CACHE + except ImportError: + return False + return (Path(HF_HUB_CACHE) / "models--test--tiny-qwen3").exists() + + +pytestmark = [ + pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs at least 2 GPUs"), + pytest.mark.skipif(not _tiny_model_available(), + reason="tiny-qwen3 missing; run scripts/create_tiny_test_model.py"), +] + + +@pytest.mark.parametrize("extra", [ + pytest.param([], id="data-parallel"), + pytest.param(["--shard-weights"], id="shard-weights"), + pytest.param(["--shard-gradients"], id="shard-gradients"), + pytest.param(["--shard-weights", "--shard-gradients"], id="shard-both"), +]) +def test_two_gpu_gradient_parity(extra): + result = subprocess.run( + [sys.executable, str(SCRIPT), + "--model", "test/tiny-qwen3", "--gpus", "2", + "--seq-len", "512", "--grad-accum", "2", + "--model-dtype", "bf16", "--matmul-dtype", "bf16", + *extra], + cwd=REPO_ROOT, + env={**os.environ, "HF_HUB_OFFLINE": "1"}, + capture_output=True, text=True, timeout=600, + ) + assert result.returncode == 0, result.stdout[-2000:] + result.stderr[-2000:]