dbrx¶
DBRX (DbrxForCausalLM).
transformer.{wte, blocks[i].{norm_attn_norm.{norm_1, attn, norm_2}, ffn},
norm_f} and lm_head. The attention sits inside a norm_attn_norm
module that norms before it, adds the residual after it, and norms again for
the FFN; the block adds the FFN's output to the state norm_attn_norm
returns. So the attention module's own output is its contribution (the add
happens outside it, in norm_attn_norm) and the base classes hold; the
aliases reach through norm_attn_norm so the block reads like any other.
The attention runs on the shared interface. The FFN is
a mixture of experts (Moe) returning the hidden states. Its router's projection is
a child nn.Linear, router.layer, whose output is the logits; the FFN's own
route_tokens_to_experts takes a softmax top-k and p-normalizes the weights
(moe_normalize_expert_weights). The experts are DBRX's own loop over experts,
taking and returning [batch, seq, hidden_size], so the per-slot outputs are
unavailable.
module-attribute
¶
RENAME = {'transformer.wte': 'embed_tokens', 'transformer.blocks': 'layers', 'transformer.norm_f': 'norm', 'norm_attn_norm.attn': 'self_attn', 'norm_attn_norm.norm_1': 'input_layernorm', 'norm_attn_norm.norm_2': 'post_attention_layernorm', 'ffn': 'mlp', 'out_proj': 'o_proj'}
Bases: Moe
DBRX's mixture of experts; the residual is added in the block, so the base holds.
class-attribute
instance-attribute
¶
expert_outputs = unavailable("DBRX's experts loop over the experts in their own forward; no tensor holds each slot's output")
intermediate_size(model: StandardizedTransformer) -> int
The experts' width, ffn_config.ffn_hidden_size; every MLP is a mixture of experts.