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.