config_stage_3.yaml

# Directory settings
checkpoint_dir: "/hindi/checkpoints/hindi-llm-v2"
train_dataset_dir:
  # Total: ~82.4B
  # Hindi Text (~43.7B) -> (53%)
  - "/hindi_tokenized/edu_score_4"                    # 14.3B
  - "/hindi_tokenized/edu_score_4"                    # 14.3B
  - "/hindi_tokenized/edu_score_4"                    # 14.3B
  - "/hindi_tokenized/edu_score_5"                    # 0.16B
  - "/hindi_tokenized/edu_score_5"                    # 0.16B
  - "/hindi_tokenized/edu_score_5"                    # 0.16B
  - "/hindi_tokenized/edu_score_5"                    # 0.16B
  - "/hindi_tokenized/edu_score_5"                    # 0.16B
  # Synth (~29.4B) -> (35.6%)
  - "/hindi_tokenized/cosmopedia"                     # 29.4B (eng)
  # Edu English Text (~0.34B) -> (0.4%)
  - "/hindi_tokenized/fineweb_edu/edu_score_5"        # 0.068B (eng)
  - "/hindi_tokenized/fineweb_edu/edu_score_5"        # 0.068B (eng)
  - "/hindi_tokenized/fineweb_edu/edu_score_5"        # 0.068B (eng)
  - "/hindi_tokenized/fineweb_edu/edu_score_5"        # 0.068B (eng)
  - "/hindi_tokenized/fineweb_edu/edu_score_5"        # 0.068B (eng)
  # Reasoning (~6B) -> (7.2%)
  - "/hindi_tokenized/big_reasoning_traces"           # 2.4B (eng)
  - "/hindi_tokenized/big_reasoning_traces"           # 2.4B (eng)
  - "/hindi_tokenized/math_meta_reasoning"            # 1.2B (eng) 
  # Edu Math Text (~3.0B) -> (3.7%)
  - "/hindi_tokenized/finemath/edu_score_5"           # 1.0B (eng)
  - "/hindi_tokenized/finemath/edu_score_5"           # 1.0B (eng)
  - "/hindi_tokenized/finemath/edu_score_5"           # 1.0B (eng)

val_dataset_dir: "/hindi_tokenized/val_stage3"  
dataset_type: "parquet"
cache_dir: "/cache"

# Data loading settings
pin_memory: true
num_workers_for_dataloader: 32
shuffle_dataset: true
mask_eos_token: false
mask_pad_token: false

# Model architecture settings
vocab_size: 49152
num_hidden_layers: 28
num_attention_heads: 16
num_key_value_heads: 8
head_dim: null
hidden_size: 1536
intermediate_size: 3072
max_position_embeddings: 4096
tie_word_embeddings: true
hidden_act: "silu"
output_hidden_states: false
attn_implementation: "flash_attention_2"
use_cache: false
no_rope_layer_interval: null
rope_theta: 50000.0
rope_scale_factor: null
rms_norm_eps: 0.000001

# Training settings
total_batch_size: 2097152
micro_batch_size: 16
eval_micro_batch_size: 8
num_train_epochs: 1
warmup_steps: 0
max_learning_rate: 0.0007
min_learning_rate: 0.0
weight_decay: 0.1
beta1: 0.9
beta2: 0.95
eps: 0.00000001
lr_decay_type: "wsd"
lr_decay_iters_coef: 0.6 
use_neg_sqrt: true
seed: 1337
max_steps: 35000 
max_grad_norm: 1.0

# Precision and optimization settings
torch_compile: false
mat_mul_precision: "highest"
tf32: true
bf16: true
gradient_checkpointing: false
use_liger_kernel: true
static_graph: false

# Hub settings
push_to_hub: false
hub_token: null
hub_model_id: null

# Tokenizer and Reference model
tokenizer_name_or_path: "Polygl0t/LilMoo-v0.2"
reference_model: "HuggingFaceTB/SmolLM2-360M"
continual_pretraining: false
chat_template_path: null

# Checkpoint settings
resume_from_checkpoint: "/hindi/checkpoints/hindi-llm-v2/Stable2"
checkpointing_steps: 5000
begin_new_stage: true
stage_name: "Stable3-Decay"

# Miscellaneous settings
sanity_check: false
sanity_check_num_samples: 100000
wandb_token: null
wandb_id: "LilMoo-v0.2"
wandb_project: "Polyglot"
wandb_desc: "Developing LLMs for low-resource languages"