iclr-debug-decoder-135b-20260916/decoder-mstok/config.yaml

task: mstok-next-concept
experiment: iclr-debug-decoder-135b-20260916-decoder-mstok
experiment_dir: /home/ubuntu/mstok-results/iclr-debug-decoder-135b-20260916/decoder-mstok
dataset:
  train_source: /home/ubuntu/data/full_owt/train_gpt2.bin
  validate_source: /home/ubuntu/data/full_owt/valid_gpt2.bin
  pad_token_id: 50257
pre_tokenizer:
  name: hf
  tokenizer: hf
  model_id: gpt2
  special_tokens:
    eos_token: <|endoftext|>
    pad_token: <|pad|>
codec:
  n_layers: 6
  context_length: 256
  embed_dim: 384
  in_vocab_size: 50304
  vqvae_vocab_size:
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  compression_factor: 1
  dropout: 0.1
  pre_quant_groupnorm: 4
  pre_quant_dropout: 0.2
  vector_quantizer_config:
    clss: multiscale_residual_vector_quantizer
    decay: 0.99
    epsilon: 1.0e-05
    commitment_cost: 0.25
    learned_l1_sampling: true
    learned_all_sampling: false
    quant_resi:
      enabled: true
      ratio: 0.5
      share_mode: 0
      learnable_ratio: false
    levels:
      use_manual_levels: true
      manual_levels:
      - 1
      - 4
      - 9
      - 16
      - 25
      - 36
      - 49
      - 64
      - 81
      - 100
      - 121
      - 144
      - 169
      - 196
      - 225
      - 256
    aux:
      fine_drop:
        prob: 0.5
        min_keep: 1
generator:
  n_layer: 12
  n_head: 12
  bias: true
  dropout: 0.1
  n_embd: 768
  context_length: 256
  vocab_size:
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  - 16384
  attn_pattern: block_diagonal
  use_positional_encoding: true
  use_level_encoding: true
  prefix_len: 0
  use_rope: true
  rope_base: 10000
  use_qk_norm: true
  post_upsample_conv:
    enabled: true
    kernel_size: 3
  use_relu2: true
  use_flex_attention: true
  shared_output_head: true
  shared_head_per_level_bias: false
  shared_head_per_level_scale: false
  shared_head_adapter_rank: 0
objectives:
  codec_weight: 1.0
  ncp_weight: 1.0
  soft_assignment_temperature: 1.0
  prediction_temperature: 1.0
  mstok_weight: 0.25
  residual_weight: 1.0
  reconstruction_weight: 1.0
  teacher_ema_decay: 0.999
  teacher_ema_warmup_steps: 1000
  mstok_warmup_steps: 500
optimization:
  codec_lr: 0.001
  codec_min_lr: 0.0001
  codec_warmup_iters: 0
  codec_lr_decay_iters: 171662
  generator_lr: 0.0005
  generator_min_lr: 1.0e-05
  generator_warmup_iters: 300
  generator_lr_decay_iters: 171662
  beta_1: 0.9
  codec_beta_2: 0.99
  generator_beta_2: 0.99
  weight_decay: 0.1
  max_grad_norm: 1.0
  level_loss_alpha: 1.0
  grad_accumulation_steps: 12
  corruption:
    mode: per_level
    per_level_probs:
    - 0.85
    - 0.8321428571
    - 0.8142857143
    - 0.7964285714
    - 0.7785714286
    - 0.7607142857
    - 0.7428571429
    - 0.725
    - 0.7071428571
    - 0.6892857143
    - 0.6714285714
    - 0.6535714286
    - 0.6357142857
    - 0.6178571429
    - 0.6
    skip_level0: true
training:
  log_interval: 10
  eval_interval: 1000
  checkpoint_interval: 1000
  eval_batch_size: 4
  val_iters: 25
  keep_last: 3
  seed: 55
  codec_initialization_seed: 42
  generator_initialization_seed: 55
  total_iters: 171662
  batch_size: 32
  expected_world_size: 8
  expected_global_batch_size: 3072
  resume_checkpoint: null
  milestone_steps:
  - 1272
  - 6358
  - 17167
  - 34333
  - 68665
  - 102997
  - 137330
  - 171662
  generation_eval_steps:
  - 17167
  - 34333
  - 68665
  - 102997
  - 137330
  - 171662
torch_compile:
  enable: true
  scope: loss
  mode: default
  dynamic: false
  fullgraph: true
  backend: inductor
mixed_precision:
  enable: true
wandb:
  entity: mstok
  project: iclr-debug
  group: full-owt-ctx256-16sq-135b-v1
  enable: true
  id: iclr-debug-decoder-135b-20260916-decoder-mstok
  resume: allow
  gradients_and_params:
    enable: false
    log: all
    log_freq: 1000
export:
  vqvae_config_template: /home/ubuntu/mstok-runs/iclr-debug-joint-20260916/config/repro-ctx256/vqvae.yaml
  ncp_config_template: /home/ubuntu/mstok-runs/iclr-debug-joint-20260916/config/repro-ctx256/ncp-sharedhead.yaml
semantic:
  enabled: true
  mode: eostok
  weight: 0.5
  warmup_steps: 500
  teacher_id: FacebookAI/roberta-base
  teacher_revision: e2da8e2f811d1448a5b465c236feacd80ffbac7b
  tokenizer_revision: 607a30d783dfa663caf39e06633721c8d4cfcd7e
  teacher_dim: 768
  feature_layer: 6
  projector_dim: 2048
  projector_seed: 56
  alignment_site: decoder
iclr_debug:
  stage: decoder-mstok
  target_positions: 135000000000
  lr_horizon_positions: 135000000000
  stop_step: 171662
  pilot: false
  tokenizer_checkpoint: null
  tokenizer_steps: null
  budget_unit: input positions including padding; log non-padding tokens separately