xtrain

gahow/xtrain

Fork 0

Commit Graph

Author	SHA1	Message	Date
Gahow Wang	15f1e526c7	train: parameterize model size (scaling ladder) Add Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn) so the model size is a tunable rung instead of a hardcoded tiny config, and Config::core_params() (num_params minus the two vocab×dim tables) — the figure the ladder is sized against (the 50257-vocab embed+lm_head adds a fixed ~25M that is not capacity). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>	2026-06-15 18:34:39 +08:00
Gahow Wang	7a4f69e430	model: add per-head QK-norm (Qwen3-compat) for xserv export xserv's Qwen3 forward unconditionally applies per-head RMSNorm to Q and K (q_norm/k_norm, shape [head_dim]) before RoPE — even gamma=1 is a real RMS divide, not identity. xtrain never had this, so an exact xserv<->xtrain loop was structurally impossible. Add it (reusing the 2D rms_norm op on the [seq*nh, hd] head rows, inserted between reshape and rope to mirror qwen3.rs's order) so the trained model is genuinely Qwen3-compatible. params() inserts q_norm,k_norm after wv; num_params() counts them; the PyTorch parity refs (parity.py / adamw_parity.py) + their name lists add the same step so the dumps stay self-consistent. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>	2026-06-15 17:33:19 +08:00
Gahow Wang	e3912c2380	model: tiny RoPE+RMSNorm+SwiGLU transformer + overfit test New crate xtrain-model: a from-scratch decoder built entirely from the autodiff op set. - Config (tiny: dim=32, 2 layers, 2 heads, head_dim=16, ffn=64). - TinyTransformer: embedding -> N x {pre-RMSNorm -> multi-head causal attention (RoPE, additive causal mask, per-head SDPA) -> residual; pre-RMSNorm -> SwiGLU MLP -> residual} -> final RMSNorm -> LM head. x@W weight convention (engine GEMM is plain A@B); dim=n_headshead_dim. - params()/zero_grad-able leaves for the optimizer; param_to_host export. - overfit test: char-level bring-up (embedded text -> vocab -> shifted targets), minimal hand-written GD (p -= lrgrad) memorises one fixed batch -> loss ~0 + greedy argmax matches targets. End-to-end fwd+bwd correctness signal. Gated #![cfg(not(no_cuda))]. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>	2026-06-15 16:05:20 +08:00

Author

SHA1

Message

Date

Gahow Wang

15f1e526c7

train: parameterize model size (scaling ladder)

Add Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn) so the model
size is a tunable rung instead of a hardcoded tiny config, and Config::core_params()
(num_params minus the two vocab×dim tables) — the figure the ladder is sized
against (the 50257-vocab embed+lm_head adds a fixed ~25M that is not capacity).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

2026-06-15 18:34:39 +08:00

Gahow Wang

7a4f69e430

model: add per-head QK-norm (Qwen3-compat) for xserv export

xserv's Qwen3 forward unconditionally applies per-head RMSNorm to Q and K
(q_norm/k_norm, shape [head_dim]) before RoPE — even gamma=1 is a real RMS
divide, not identity. xtrain never had this, so an exact xserv<->xtrain loop
was structurally impossible. Add it (reusing the 2D rms_norm op on the
[seq*nh, hd] head rows, inserted between reshape and rope to mirror
qwen3.rs's order) so the trained model is genuinely Qwen3-compatible.

params() inserts q_norm,k_norm after wv; num_params() counts them; the
PyTorch parity refs (parity.py / adamw_parity.py) + their name lists add the
same step so the dumps stay self-consistent.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

2026-06-15 17:33:19 +08:00

Gahow Wang

e3912c2380

model: tiny RoPE+RMSNorm+SwiGLU transformer + overfit test

New crate xtrain-model: a from-scratch decoder built entirely from the
autodiff op set.
- Config (tiny: dim=32, 2 layers, 2 heads, head_dim=16, ffn=64).
- TinyTransformer: embedding -> N x {pre-RMSNorm -> multi-head causal
  attention (RoPE, additive causal mask, per-head SDPA) -> residual;
  pre-RMSNorm -> SwiGLU MLP -> residual} -> final RMSNorm -> LM head.
  x@W weight convention (engine GEMM is plain A@B); dim=n_heads*head_dim.
- params()/zero_grad-able leaves for the optimizer; param_to_host export.
- overfit test: char-level bring-up (embedded text -> vocab -> shifted
  targets), minimal hand-written GD (p -= lr*grad) memorises one fixed
  batch -> loss ~0 + greedy argmax matches targets. End-to-end fwd+bwd
  correctness signal. Gated #![cfg(not(no_cuda))].

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

2026-06-15 16:05:20 +08:00

3 Commits