Refactor and fixes

[ngc92]

some more refactoring, adding support for mistral (llama-like) architectures, fix matmul for larger shapes.

[Copilot]

Pull request overview

This PR refactors model/weight naming utilities, expands architecture support to include Mistral (LLaMA-like), hardens config loading against unsupported HF settings, and fixes GEMM address calculations for very large matrix shapes. It also adds integration tests to validate gradient parity across architectures and across 2-GPU training modes.

Changes:

• Add gradient-parity integration tests for supported/refused architectures and for 2-GPU runs with different sharding modes.
• Add Mistral architecture handling in transformer config load/save, plus explicit rejection of unsupported config features (rope_scaling, attention_bias, etc.).
• Refactor weight/gradient naming to use centralized name helpers; update GEMM kernel indexing to use wider integer arithmetic and add a large-output regression test.

Reviewed changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
test/test_multi_gpu.py	Adds end-to-end 2-GPU gradient parity test runner wrapper.
test/test_architectures.py	Adds per-architecture gradient parity + “must refuse” negative fixtures.
src/utilities/comm.h	Adjusts NCCL all-gather API to accept Tensor instead of TensorShard.
src/utilities/comm.cpp	Updates NCCL all-gather implementation signature accordingly.
src/training/transformer_config.h	Adds MISTRAL architecture enum.
src/training/transformer_config.cpp	Adds Mistral load/save support and explicit rejection of unsupported HF config fields.
src/testing/test-gemm.cpp	Adds large-output regression test for GEMM INT overflow bug.
src/models/llama_weights.h	Adds centralized HF weight-name helpers and a typed accessor for block tensors.
src/models/llama_weights.cpp	Refactors tensor iteration/naming to use centralized name helpers.
src/models/llama_optimizer.cpp	Refactors optimizer-state tensor naming/iteration to use centralized name helpers.
src/kernels/rmsnorm.cu	Fixes shared-memory declaration/alignment for RMSNorm kernels.
src/kernels/gemm_mma.cu	Uses wider index arithmetic to avoid overflow for large m*k / m*n.
src/binding/python/tests/torch_reference.py	Fixes multi-GPU gradient retrieval semantics (per-GPU shard concat) and per-GPU microbatch sizing.
src/binding/python/tests/run.py	Makes training test runner respect config.gpus and divides global batch into per-GPU microbatches.
src/binding/py_train.cpp	Generalizes multi-GPU gradient enumeration via centralized weight-name helpers.
scripts/create_tiny_test_model.py	Adds script to generate tiny HF-cache fixtures for integration tests (incl. mistral).

Comments suppressed due to low confidence (2)

src/testing/test-gemm.cpp:289

• cudaMemset sizes for a and b are also missing sizeof(__nv_fp8_e4m3), so only a fraction of each buffer is initialized (and if the allocations are fixed, these would become short writes).

    CUDA_CHECK(cudaMemset(a, 0x38, (std::size_t)m * k));
    CUDA_CHECK(cudaMemset(b, 0x38, (std::size_t)n * k));

src/kernels/gemm_mma.cu:217

• out_offset is intended to be 64-bit (per the comment), but it's declared as long, which may be 32-bit depending on platform/toolchain. Use an explicitly 64-bit type for out_offset so the INT overflow fix is robust.

            // 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;

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 18 out of 18 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated no new comments.

Suppressed comments (7)

src/utilities/comm.cpp:189

• schedule_all_gather takes a generic Tensor as the source, but gather_weight() assumes the source pointer addresses the local shard (it sends Bytes/world_size() bytes from Src). Passing a full tensor here would produce an incorrect all-gather (each rank sends the same prefix). Restoring the TensorShard parameter type (or at least validating src.bytes() * world_size() == tgt.bytes()) would prevent accidental misuse.

void NCCLCommunicator::schedule_all_gather(const Tensor& src, Tensor& tgt) {
    if (src.Data == nullptr) {
        throw std::runtime_error("gather: Source tensor is null");
    }

src/testing/test-gemm.cpp:287

• In the new large-output GEMM test, the cudaMalloc / cudaMemset sizes for a and b are missing sizeof(__nv_fp8_e4m3). This is inconsistent with the rest of the file (which uses * sizeof(T) everywhere) and would under-allocate / under-initialize if the fp8 storage type is not exactly 1 byte.

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

src/kernels/gemm_mma.cu:86

• This kernel relies on 64-bit indexing, but it uses long casts. long is not guaranteed to be 64-bit on all platforms (e.g., Windows), which can reintroduce the overflow the comment warns about. Prefer long long (or a fixed-width integer) for the intermediate products.

    // 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) + (long)(bi + wid) * TI * stride;
        s_ptr = input_tiles + wid * ROW_OFFSET;
    } else {
        g_ptr = reinterpret_cast<const uint4*>(b) + (long)(bj + wid - 2) * TJ * stride;
        s_ptr = input_tiles + (wid - 2) * ROW_OFFSET + DEPTH * PIPE_OFFSET;
    }

src/kernels/gemm_mma.cu:218

• Same portability concern here: out_offset is meant to be 64-bit, but long is not guaranteed to be 64-bit. Using long long (or int64_t) keeps the overflow fix effective across platforms.

            // 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;

src/training/transformer_config.cpp:105

• The comment says Qwen’s sliding_window is inactive unless use_sliding_window is set, but the code currently treats a missing use_sliding_window as true (active) via get_or("use_sliding_window", true). This makes the behavior contradict the comment and could incorrectly reject some Qwen configs. Consider making the default depend on architecture: default false for Qwen2/Qwen3, default true for others (no flag).

    // 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());
    }

src/utilities/comm.h:24

• schedule_all_gather is intended to gather a per-rank shard into a full tensor. If you switch the parameter type back to TensorShard for type-safety, comm.h also needs a forward declaration for TensorShard (it was removed in this PR).

struct Tensor;

src/utilities/comm.h:49

• schedule_all_gather assumes src points at the local shard that will be all-gathered into the full tgt (see gather_weight() which sends Bytes/world_size() bytes from Src). Taking const TensorShard& here keeps that contract enforceable at compile time and prevents accidental calls with a full tensor.

    void schedule_all_gather(const Tensor& src, Tensor& tgt);