研究

研究#

自定义 Transformer#

为了实验，您可能需要创建更小的自定义 Transformer 版本，而不是使用默认的官方 “Gemma” 预训练版本。

这是一个创建仅文本的 12 层 Transformer 的示例

class MyTinyTransformer(gm.nn.Transformer):
  config: gm.nn.config.TransformerConfig = gm.nn.config.TransformerConfig(
      num_layers=12,
      final_logit_softcap=None,
      num_embed=262144,  # Vocab size, matching the tokenizer
      embed_dim=896,
      hidden_dim=4 * 896,
      num_heads=4,
      head_dim=256,
      num_kv_heads=1,
      use_post_attn_norm=True,
      use_post_ffw_norm=True,
      use_qk_norm=True,
      attention_types=gm.nn.config.make_attention_layers_types(
          pattern=gm.nn.config.GEMMA3_ATTENTION_PATTERN,
          num_layers=12,
      ),
      query_pre_attn_norm=gm.nn.config.QueryPreAttentionNormalisation.BY_ONE_OVER_SQRT_HEAD_DIM,
      attn_logits_soft_cap=None,
      sliding_window_size=512,
      transpose_gating_einsum=True,
      local_base_frequency=10_000,
      global_base_frequency=1_000_000,
      max_cache_length=None,
      vision_encoder=None,  # Text-only
  )

  INFO = gm.nn.config.ModelInfo(
      tokenizer_version=3,  # Auto-select the tokenizer in the sampler
  )
目录