example_text_generation_config.yaml

settings:
  referencing_keys:
    sample_key: input_ids
    prediction_key: logits
  model_path: ./checkpoints/2024-06-27__14-17-52/eid_2024-06-27__14-17-52-model-num_steps_48-num_tokens_393216.bin
  device: 0
  sequence_length: 1024

text_inference_component:
  component_key: inference_component
  variant_key: text
  config:
    device: ${settings.device}
    model:
      instance_key: checkpointed_model
      pass_type: BY_REFERENCE
    tokenizer:
      instance_key: tokenizer
      pass_type: BY_REFERENCE
    sequence_length: ${settings.sequence_length}
    eod_token: <|endoftext|>
    prompt_template: "{prompt_input}" # "<instruction> Du bist Moody, ein LLM welches Menschen helfen soll. user: {prompt_input}"
    temperature: 1
    # chat: false

checkpointed_model:
  component_key: model
  variant_key: checkpointed
  config: 
    checkpoint_loading:
      component_key: checkpoint_loading
      variant_key: torch
      config:
        device: ${settings.device}
        precision: BF16
    model:
      instance_key: model
      pass_type: BY_REFERENCE
    checkpoint_path: ${settings.model_path}

model:
  component_key: model
  variant_key: gpt2
  config:
    sample_key: ${settings.referencing_keys.sample_key}
    poe_type: NOPE
    sequence_length: ${settings.sequence_length}
    prediction_key: ${settings.referencing_keys.prediction_key}
    vocab_size: 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
    n_layer: 2
    n_head_q: 8
    n_head_kv: 8
    ffn_hidden: 128
    n_embd: 128
    dropout: 0.0
    bias: true # True: bias in Linears, like GPT-2. False: a bit better and faster
    attention_config:
      qkv_transforms:
        - type_hint: RotaryTransform
          config:
            n_embd: ${model.config.n_embd}
            n_head: ${model.config.n_head_q} #it has to be head_q here
            seq_length_dim: -2
    attention_implementation: manual
    activation_type: gelu
    attention_norm:
      component_key: layer_norm
      variant_key: rms_norm
      config:
        ndim: ${model.config.n_embd}
        bias: true
        epsilon: 1e-5
    ffn_norm:
      component_key: layer_norm
      variant_key: rms_norm
      config:
        ndim: ${model.config.n_embd}
        bias: true
        epsilon: 1e-5
    lm_head_norm:
      component_key: layer_norm
      variant_key: rms_norm
      config:
        ndim: ${model.config.n_embd}
        bias: true
        epsilon: 1e-5

tokenizer:
  component_key: tokenizer
  variant_key: pretrained_hf_tokenizer
  config:
    pretrained_model_name_or_path: tokenizer
    padding: false
    truncation: false