Skip to content

feat(hslm): Phase 2 — Temperature STE + EMA shadow + bfloat16 + RMSNorm - #573

Merged
gHashTag merged 1 commit into
mainfrom
feat/294-phase2-ste-ema-bfloat16
Apr 30, 2026
Merged

gHashTag merged 1 commit into
mainfrom
feat/294-phase2-ste-ema-bfloat16

Conversation

@gHashTag

Copy link
Copy Markdown
Owner

Summary

Phase 2 training optimizations: temperature-annealed STE, EMA shadow weights, bfloat16 accumulation, RMSNorm.

New file

  • src/b2t/ste_ema_bfloat16.zig — 276 LOC

Components

  1. StraightThroughEstimator: Cosine-annealed temperature (5.0→0.1) for soft ternary during training, hard ternary at inference
  2. EMAShadow: Exponential moving average (decay=0.999) with warmup correction for eval weights
  3. BFloat16: 16-bit format (1/8/7) with f32 roundtrip
  4. BFloat16Accumulator: BF16 gradient accumulation with periodic quantization
  5. RMSNorm: Zero-parameter normalization (BitNet b1.58 style)

Tests (7)

  • Temperature annealing, soft ternarize, STE gradient, EMA tracking, bf16 roundtrip, accumulator, RMSNorm

Closes #294

- Add src/b2t/ste_ema_bfloat16.zig
- StraightThroughEstimator: cosine-annealed temperature (5.0→0.1)
  soft ternary forward, hard ternary inference, STE gradient
- EMAShadow: decay=0.999, warmup-corrected, evaluation weights
- BFloat16: 16-bit format (1/8/7), roundtrip f32↔bf16
- BFloat16Accumulator: bf16 gradient accumulation + quantize
- RMSNorm: zero-param normalization (BitNet b1.58 style)
- 7 tests: temperature annealing, soft ternary, STE gradient,
  EMA tracking, bf16 roundtrip, accumulator, RMSNorm unit norm

Closes #294
@gHashTag
gHashTag merged commit a92443c into main Apr 30, 2026
13 of 25 checks passed
@gHashTag
gHashTag deleted the feat/294-phase2-ste-ema-bfloat16 branch April 30, 2026 01:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Phase 2: Temperature STE + EMA shadow weights + bfloat16 (+15-25% speedup)

1 participant