Refactor and fixes

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 $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr> $<$<COMPILE_LANGUAGE:CUDA>:-lineinfo>)

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-<arch> in the HF hub cache, so
+both `transformers` (with HF_HUB_OFFLINE=1) and llmq can load it as `test/tiny-<arch>`.
+
+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()

src/binding/py_train.cpp

@@ -279,32 +279,21 @@ std::vector<std::pair<std::string, Tensor>> MultiGPUPyTrainer::get_gradients(int
     using namespace LLamaWeightID;
 
     std::vector<std::pair<std::string, Tensor>> 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);

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,

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):

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<const uint4*>(a) + (bi + wid) * TI * stride;
+        g_ptr = reinterpret_cast<const uint4*>(a) + (long)(bi + wid) * TI * stride;
         s_ptr = input_tiles + wid * ROW_OFFSET;
     } else {
-        g_ptr = reinterpret_cast<const uint4*>(b) + (bj + wid - 2) * TJ * stride;
+        g_ptr = reinterpret_cast<const uint4*>(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<nv_bfloat16, 8>::load(out + ((i + ii) * TI + r) * n + (j + jj) * TJ + 8 * c);
+                auto old = GenericVector<nv_bfloat16, 8>::load(out + out_offset);
                 auto upd = GenericVector<nv_bfloat16, 8>::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<BiasType, 8>::load(bias + (j + jj) * TJ + 8 * c);
                 auto upd = GenericVector<nv_bfloat16, 8>::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<uint4*>(out_shared)[c + 2 * r];
-                *reinterpret_cast<uint4*>(out + ((i + ii) * TI + r) * n + (j + jj) * TJ + 8 * c) = load;
+                *reinterpret_cast<uint4*>(out + out_offset) = load;
             }
         }
     }

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

src/models/llama_optimizer.cpp

@@ -23,28 +23,20 @@ struct OptStateWrapper : ITensorContainer {
 };
 
 void OptStateWrapper::iterate_tensors(const std::function<void(std::string, const TensorShard&)>& 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);
+            }
+        }
     }
 }