Skip to content

spec(ml): train_hslm speed=4 -- int8 trit forward with one scale per row, error bounded against speed=2 #10875

Closed

Description

@gHashTag

The HSLM trainer (specs/ml/train_hslm.t27) computes every forward ternary matmul in f64: each activation is multiplied by a trit (-1, 0, +1) stored as one byte and added into an f64 sum. Speed 3 (#10849, #10866) reads each matrix once per window, but the inner loop still works in f64, four lanes per AVX2 register.

This issue adds speed=4: speed 3's window with every forward ternary matmul done in integers. Each row of activations gets one scale, s = max |x| / 127, and is rounded to int8. Each trit is read as the int8 it already is (transposed once per call so the dot runs along contiguous memory, sixteen independent accumulators). The products add in an i32 accumulator, which is exact in any order (|sum| <= 127 * 729 < 2^17), and the output is s * sum. The backward stays speed 3's f64 code, on the activations this forward made. This changes the bits, so the tests bound the error instead of comparing bits. The bound is stated and proved in the spec: per output, |error| <= (max_j |x_j| / 254) * (nonzero trits in that column).

The default stays speed=2; train.cfg selects 4.

Boundary

  • specs/ml/train_hslm.t27
  • .trinity/seals/ml_train_hslm.json

User Scenarios

  • Given a real window's activations when block 0's K, block 1's down and the head are computed at speed 3 and at speed 4 then every output's difference is within (max_j |x_j| / 254) * nnz_c, and the largest ratio to that bound is above 0 (the int8 path really rounds).
  • Given one window when it runs at speed 2 and at speed 4 then the logits, the k-position loss and the gradients differ by less than the bounds stated in the spec.
  • Given speed=4 and speed=2 with the same seed when the trainer runs a few hundred steps on the Railway lab then both loss curves and both tok/s are constants in the spec, and a test re-derives the largest gap between the curves and the speed-up.

Requirements

  • FR-001: speed=4 MUST quantize each activation row to int8 with one scale per row and accumulate trit products in an integer accumulator.
  • FR-002: The kernel error bound MUST be stated in the spec and checked by a test on every output of three real matmuls.
  • FR-003: A test MUST bound the difference of logits, loss and gradients between speed 4 and speed 2, and state the bound.
  • FR-004: A lab loss-curve comparison of speed 4 against speed 2 over a few hundred steps with the same seed MUST be recorded as constants, with a test that re-derives the largest gap.
  • FR-005: The tok/s of both speeds MUST be constants with a test that re-derives the speed-up.
  • FR-006: The default speed MUST stay 2, and speeds 0 to 3 MUST be unchanged (their tests stay green).
  • FR-007: Edits MUST stay in section 1 (kernels), the window functions and the tests of the trainer.

Success Criteria

  • t27c test-report specs/ml/train_hslm.t27 on the Railway lab: exit 0, 0 failed, 0 vacuous.
  • The kernel bound holds on every checked output (ratio <= 1).
  • t27c seal specs/ml/train_hslm.t27 --verify: exit 0.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions