implement qat and weight sharing

[kiratoyoshihara]

Note: This is a draft PR.

Approach Overview

To achieve extreme efficiency under the 16MB artifact constraint and the 10-minute training limit on 8xH100, I am introducing a custom architecture modifying the baseline train_gpt.py.

1. Depth Recurrence (Parameter Tying)

Standard transformer architectures can only hold a few million parameters within 16MB. To overcome this, I have replaced the standard nn.ModuleList of blocks with a single shared Transformer Block. By applying this block recursively across all layers, we simulate a much deeper network and maximize representational capacity without increasing the actual file size on disk.

2. Quantization-Aware Training (QAT)

To cram the maximum effective parameter count into the footprint, I implemented a Fake Quantization step (symmetric INT8 via Straight-Through Estimator) inside CastedLinear. The forward/backward passes are kept in high precision (bfloat16) for H100 throughput, but the weights are regularized to be robust to extreme post-training quantization.

Current Status

• Basic architecture design and local implementation
• QAT and Parameter Tying integrated into train_gpt.py
• Hyperparameter search (MuP) on a single GPU
• Full-scale scaling and evaluation on 8xH100 SXM (Pending compute)

I will update this PR with training logs and the final artifact once the large-scale runs are complete.