Views
No views yet
transformers and vLLM without any code changes.RMSNorm → Linear that (i) folds the per-channel normalization weights into the following linear layer (W* = W · diag(g)) and (ii) defers the scalar 1/RMS(x) normalization to after the matmul. On hardware with distinct vector and matrix units, the matrix multiplication and the RMS reduction can execute in parallel.| Tensor | Source | This checkpoint |
|---|---|---|
model.layers.*.input_layernorm.weight | learned per-channel g | all ones |
model.layers.*.self_attn.{q,k,v}_proj.weight | W | W · diag(g_input_layernorm) |
model.layers.*.post_attention_layernorm.weight | learned per-channel g | all ones |
model.layers.*.mlp.{gate,up}_proj.weight | W | W · diag(g_post_attention_layernorm) |
bfloat16); the merged products are computed in float32 internally before casting back. model.norm.weight is unchanged (tied embeddings). Mathematical identity holds by construction.1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3tok = AutoTokenizer.from_pretrained('open-machine/Llama-3.2-1B-FlashNorm')
4model = AutoModelForCausalLM.from_pretrained('open-machine/Llama-3.2-1B-FlashNorm')
5
6ids = tok('Once upon a time there was', return_tensors='pt').input_ids
7out = model.generate(ids, max_new_tokens=50, do_sample=False)
8print(tok.decode(out[0], skip_special_tokens=True))vllm serve open-machine/Llama-3.2-1B-FlashNormx·g·W would.RMSNorm + QKV kernel (deferring g to runtime) eliminates the framework dependency and is in progress for vLLM / FlashInfer.