###############################
# llama-3-attenuated.yaml #
###############################

# Use: mergekit-yaml --clone-tensors ./llama-3-attenuated.yaml ./llama-3-attenuated
# See: https://github.com/arcee-ai/mergekit/issues/198 for discussion/reasoning behind this idea.

# ---

# The scale factor to use, eg: solve x^2 = 1/2 --> x = 1/sqrt(2) ≈ 0.7071067812
const_tag: &scale_factor 0.7071067812 # 1/sqrt(2)

# The filter parameters of a scaled block. 
attenuate-env: &attenuated_env
  parameters:
    scale:
      - filter: q_proj
        value: *scale_factor
      - filter: k_proj
        value: *scale_factor
      - value: 1.0

# ---

slices:
  - sources:
    - model: kuotient/Meta-Llama-3-8B-Instruct
      layer_range: [0, 4] 
  - sources:
    - model: kuotient/Meta-Llama-3-8B-Instruct
      layer_range: [4, 8] 
      <<: *attenuated_env
      
  - sources:
    - model: kuotient/Meta-Llama-3-8B-Instruct
      layer_range: [4, 12] 
      <<: *attenuated_env
  - sources:
    - model: kuotient/Meta-Llama-3-8B-Instruct
      layer_range: [8, 16] 
      <<: *attenuated_env
  - sources:
    - model: kuotient/Meta-Llama-3-8B-Instruct
      layer_range: [12, 20] 
      <<: *attenuated_env
  - sources:
    - model: kuotient/Meta-Llama-3-8B-Instruct
      layer_range: [16, 24] 
      <<: *attenuated_env
  - sources:
    - model: kuotient/Meta-Llama-3-8B-Instruct
      layer_range: [20, 28] 
      <<: *attenuated_env
      
  - sources:
    - model: kuotient/Meta-Llama-3-8B-Instruct
      layer_range: [24, 28]
      <<: *attenuated_env
  - sources:
    - model: kuotient/Meta-Llama-3-8B-Instruct
      layer_range: [28, 32]

merge_method: passthrough
dtype: bfloat16