codegen¶
CodeGen (CodeGenForCausalLM).
transformer.{wte, drop, h[i].{ln_1, attn, mlp}, ln_f} and lm_head. GPT-J's
parallel block: ln_1 feeds both sublayers and x + attn + mlp is summed in
the block, which returns (hidden_states, attn_weights). The attention projects
queries, values and keys (in that order) out of one fused qkv_proj laid out for
four TPU cores, applies rotary embeddings to the first rotary_dim channels, and
does its own arithmetic in _attn: an fp32 q @ k^T, the mask added, then
the division by sqrt(head_dim), softmax, a cast to the values' dtype,
attn_dropout, @ v. The model has no attention implementation but that one,
so attn_implementation never routes around it and the interior needs no eager
load. drop has no standard name.
CodeGen2 checkpoints (Salesforce/codegen2-*) declare model_type codegen,
but their own code lays qkv_proj out in eight groups (mp_num = 8); loaded
through transformers' class, which splits it in four, the queries, keys and values
are cut wrong and nothing reports it.
module-attribute
¶
RENAME = {'transformer.wte': 'embed_tokens', 'transformer.h': 'layers', 'transformer.ln_f': 'norm', 'ln_1': 'input_layernorm', 'attn': 'self_attn', 'out_proj': 'o_proj', 'act': 'act_fn', 'fc_out': 'down_proj'}
module-attribute
¶
intermediate_size(model: StandardizedTransformer) -> int
The MLP width is n_inner, None meaning four times the hidden size.