mixed_quantize_int6_gptq() prints its summary unconditionally. In distributed runs this produces duplicated lines (once per rank), as visible in the provided logs. Please guard this print behind rank == 0 / master_process (pass rank into the function or log at the call site) to keep logs deterministic and easier to parse.

The artifact filename/logging uses final_model.int6.ptz and messages like int6+{_COMPRESSOR}, but the PR description/record metadata claim an int5 quantization scheme ([-15,15], 31 levels). This mismatch is confusing for reviewers and for any downstream tooling that expects naming to reflect the actual quantization format. Please align naming (filenames, log strings, helper function names) with the true quantization level being produced.

This record script writes final_model.int6.ptz and logs int6+{_COMPRESSOR}, but the submission/README claim an int5 scheme ([-15,15], 31 levels). Please align filenames/log strings and helper names with the actual quantization format to avoid confusion during verification.

When FlashAttention is available, the code calls flash_attn_3_func(q, k, v, ...) even if num_kv_heads != num_heads. Unlike the SDPA fallback (which enables GQA), FlashAttention may not accept mismatched head counts and can error if someone overrides NUM_KV_HEADS. Consider gating the FlashAttention path to num_kv_heads == num_heads (or using the FlashAttention API that explicitly supports GQA/MQA) and otherwise fall back to scaled_dot_product_attention(..., enable_gqa=True).

The QAT “int5” STE path uses scale = row_clip / 15.0 but clamps to [-16, 15], which is 32 levels and doesn’t match the stated [-15, 15] (31 levels) quantization used elsewhere (clip_range=15). This inconsistency makes it hard to reason about what quantization the model is being trained for. Please make the clamp range consistent with the intended scheme (and update the README/submission metadata accordingly).

The per-step EMA update iterates for name, t in base_model.state_dict().items() on every training step. Building a full state_dict each step can be a significant CPU overhead and allocates many temporary objects. Consider tracking EMA over named_parameters()/buffers directly (or caching the parameter list once) to avoid repeated state_dict() construction in the hot path.

eval_val_sliding() includes a final partial window whenever it has at least 1 token. When total_tokens is not an exact multiple of stride, that last window will typically re-score tokens that were already scored as the “last stride” of the previous window (double counting the tail). Consider using the same inclusion condition as the TTT evaluator (wlen >= stride or ws == 0) or otherwise adjusting the scoring slice for the final window to ensure each token is counted exactly once.

The root train_gpt.py no longer follows the repo’s stated goal of being a newcomer-friendly baseline (and it drops the existing guidance comment). The repository README explicitly says the best/SOTA configs should live under /records (README.md:204), and train_gpt_mlx.py documents a 1500-line hard stop. Please restore the baseline-oriented header and move record-specific/experimental code paths into the /records/.../train_gpt.py snapshot instead of upstreaming them into the root script.

Rotary.forward() computes scale ** (rd / (rd - 2)) when seq_len > train_seq_len. If ROPE_DIMS is set to 2 via env (or any value <= 2), this will divide by zero (or become numerically unstable). Please validate rope_dims at initialization (e.g., require even and > 2 when extrapolation is enabled) or guard the extrapolation formula for small rd.

In eval_val_sliding_ttt(), the model is left in eval() mode during the adaptation/training phase (there’s no base_model.train() before computing ttt_loss/backward()). That can silently disable any training-mode behaviors (and also disables the QAT STE path which checks self.training). Please switch to train() for Phase 2 (and then back to eval() for scoring) so the adaptation step is actually run under training semantics.