# GammaNet Configuration for Pretraining
# Gamma Space Model blocks in the TaoTrain causal LM shell

# ============================================================================
# Model Architecture - GammaNet (Gamma Space Model)
# ============================================================================
model:
  architecture_type: gamma_net
  vocab_size: 8192
  hidden_dim: 768
  num_layers: 10
  num_heads: 8
  max_seq_length: 1024

  # Shared LM shell feed-forward / embedding controls
  hidden_dim_ff: 2048
  dropout: 0.02
  use_factorized_embedding: false
  d_embed_rank: 96

  # GammaSpaceBlock-specific settings
  gamma_hidden_dim: 256
  gamma_dt_min: 1e-3
  gamma_dt_max: 1e-1
  gamma_dt_init: 1e-2
  gamma_discretization: bilinear
  gamma_prenorm: true
  gamma_residual_scale: 1.0
  gamma_activation: gelu
  gamma_gate: true
  gamma_use_D: true
  gamma_kernel_mode: auto
  gamma_kernel_threshold: 64
  gamma_use_output_linear: true
  gamma_gate_bias: 2.0
  gamma_input_gate: true
  gamma_input_gate_bias: 2.0
  gamma_layer_scale_init: 0.1

  # Weight initialization standard deviation for TaoTrain LM shell
  init_std: 0.02

# ============================================================================
# Dataset Configuration - Local JSONL
# ============================================================================
dataset:
  local: true
  jsonl_path: /home/student/Data/TaoData/pretrain.jsonl
  text_field: text
  max_samples: 1000000
  samples_per_chunk: 1000

  # Tokenizer configuration
  tokenizer_type: sentencepiece
  tokenizer_path: tokenizer/tokenizer.model
  tokenizer_threads: 4

# ============================================================================
# Training Hyperparameters
# ============================================================================
batch_size: 32
num_epochs: 2
gradient_accumulation_steps: 8
max_grad_norm: 1.0

# ============================================================================
# Optimizer - Hybrid Muon + AdamW
# ============================================================================
optimizer:
  optimizer_type: hybrid_muon_adamw
  learning_rate: 5e-3
  adamw_lr: 5e-4
  weight_decay: 0.01
  betas: [0.9, 0.999]
  eps: 1e-8

# ============================================================================
# Learning Rate Scheduler - 3-Phase Cosine with Warmup
# ============================================================================
scheduler:
  scheduler_type: cosineWarmup
  warmup_steps: 300
  warmup_ratio: 0.0
  steady_ratio: 0.05
  min_lr_ratio: 0.1
  num_cycles: 0.5

# ============================================================================
# Data Type and Device
# ============================================================================
dtype: bfloat16
device: cuda

# ============================================================================
# Checkpointing and Validation
# ============================================================================
checkpoint_dir: checkpoints/pretrain_gamma
save_every_steps: 81920
save_best_model: true
keep_last_n_checkpoints: 3

eval_every_steps: 8192
eval_samples: 8000

# ============================================================================
# Logging
# ============================================================================
log_every_steps: 50
aim_repo: .aim

# ============================================================================
# Miscellaneous
# ============================================================================
seed: 42
num_workers: 0
pin_memory: true