checkpoints:
  checkpoint_interval: 2000
  checkpoints_path: /scratch/elie/checkpoints
  checkpoints_path_is_shared_file_system: false
  load_lr_scheduler: true
  load_optimizer: true
  resume_checkpoint_path: s3://smollm3/tp-fix-final-pre-training/1704-48n-part1
  save_final_state: true
  save_initial_state: false
data_stages:
- data:
    dataset:
    # All the dataset can be find here https://huggingface.co/collections/HuggingFaceTB/smollm3-pretraining-datasets-685a7353fdc01aecde51b1d9
      dataset_read_path:
      - /scratch/smollm3-data-part1/fineweb-edu # https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu
      - /scratch/smollm3-data-part1/dclm # https://huggingface.co/datasets/mlfoundations/dclm-baseline-1.0
      - /scratch/smollm3-data-part1/pes2o # https://huggingface.co/datasets/allenai/dolmino-mix-1124/viewer/pes2o
      - /scratch/smollm3-data-part1/wiki # https://huggingface.co/datasets/allenai/dolmino-mix-1124/viewer/wiki
      - /scratch/smollm3-data-part1/stackexchange # https://huggingface.co/datasets/allenai/dolmino-mix-1124/viewer/stackexchange
      # we use FineWeb2-HQ for all the languages below except Hindi, Thai, Korean for which we use FineWeb2
      - /scratch/smollm3-data-part1/fw2-fra # https://huggingface.co/datasets/epfml/FineWeb2-HQ/viewer/hun_Latn
      - /scratch/smollm3-data-part1/fw2-spa # https://huggingface.co/datasets/epfml/FineWeb2-HQ/viewer/spa_Latn
      - /scratch/smollm3-data-part1/fw2-deu # https://huggingface.co/datasets/epfml/FineWeb2-HQ/viewer/deu_Latn
      - /scratch/smollm3-data-part1/fw2-ita # https://huggingface.co/datasets/epfml/FineWeb2-HQ/viewer/ita_Latn
      - /scratch/smollm3-data-part1/fw2-por # https://huggingface.co/datasets/epfml/FineWeb2-HQ/viewer/por_Latn
      - /scratch/smollm3-data-part1/fw2-cmn # https://huggingface.co/datasets/epfml/FineWeb2-HQ/viewer/cmn_Hani
      - /scratch/smollm3-data-part1/fw2-rus # https://huggingface.co/datasets/epfml/FineWeb2-HQ/viewer/rus_Cyrl
      - /scratch/smollm3-data-part1/fw2-fas # https://huggingface.co/datasets/epfml/FineWeb2-HQ/viewer/fas_Arab
      - /scratch/smollm3-data-part1/fw2-jpn # https://huggingface.co/datasets/epfml/FineWeb2-HQ/viewer/jpn_Jpan
      - /scratch/smollm3-data-part1/fw2-kor # https://huggingface.co/datasets/HuggingFaceFW/fineweb-2/viewer/kor_Hang
      - /scratch/smollm3-data-part1/fw2-hin # https://huggingface.co/datasets/HuggingFaceFW/fineweb-2/viewer/hin_Deva
      - /scratch/smollm3-data-part1/fw2-tha # https://huggingface.co/datasets/HuggingFaceFW/fineweb-2/viewer/tha_Thai
      - /scratch/smollm3-data-part1/fw2-vie # https://huggingface.co/datasets/epfml/FineWeb2-HQ/viewer/vie_Latn
      - /scratch/smollm3-data-part1/fw2-ell # https://huggingface.co/datasets/epfml/FineWeb2-HQ/viewer/ell_Grek
      - /scratch/smollm3-data-part1/infiwebmath # https://huggingface.co/datasets/HuggingFaceTB/finemath/viewer/infiwebmath-3plus
      - /scratch/smollm3-data-part1/finemath # https://huggingface.co/datasets/HuggingFaceTB/finemath/viewer/finemath-3plus
      - /scratch/smollm3-data-part1/stack-v2-Python # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-Java # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-JavaScript # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-C # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-Cpp # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-C-Sharp # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-PHP # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-TypeScript  # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-Swift # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-SQL # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-Ruby # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-Markdown # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-HTML # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-Rust # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-Go # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/stack-v2-Shell # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/pull-requests # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/kaggle # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/jupyter-scripts # https://huggingface.co/datasets/bigcode/the-stack-v2
      - /scratch/smollm3-data-part1/github-issues # https://huggingface.co/datasets/bigcode/the-stack-v2
      dataset_folder:
      - s3://smollm3/datasets/llama_tokenized-global-chunks/fineweb-edu/fineweb-edu/
      - s3://smollm3/datasets/llama_tokenized-global-chunks/dclm/dclm/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/pes2o/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/wiki/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stackexchange/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-fra/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-spa/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-deu/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-ita/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-por/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-cmn/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-rus/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-fas/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-jpn/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-kor/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-hin/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-tha/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-vie/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/fw2-ell/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/infiwebmath/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/finemath/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-Python/ # we use stack v2 (StarCoder2data) data, just typo in s3 path names
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-Java/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-JavaScript/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-C/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-Cpp/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-C-Sharp/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-PHP/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-TypeScript/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-Swift/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-SQL/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-Ruby/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-Markdown/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-HTML/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-Rust/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-Go/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/stack-edu-Shell/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/pull-requests/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/kaggle/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/jupyter-scripts/
      - s3://smollm3/datasets/llama_tokenized-individual-chunks/github-issues/
      dataset_weights:
      - 0.333
      - 0.37
      - 0.02
      - 0.001
      - 0.004
      - 0.016
      - 0.02
      - 0.022
      - 0.0105
      - 0.01
      - 0.01
      - 0.01
      - 0.003
      - 0.00325
      - 0.00325
      - 0.00325
      - 0.00325
      - 0.00325
      - 0.00225
      - 0.01 
      - 0.017
      - 0.025
      - 0.013
      - 0.013
      - 0.007
      - 0.018
      - 0.006
      - 0.006
      - 0.003
      - 0.001
      - 0.004
      - 0.0008
      - 0.005
      - 0.006
      - 0.0008
      - 0.0005
      - 0.0007
      - 0.006
      - 0.0005
      - 0.0055
      - 0.0032
      pad_samples_to_global_batch_size: false
      return_positions: true
      token_size_in_bytes: 4
      tokenizer_name: meta-llama/Llama-3.2-1B
      use_old_brrr_dataloader: false
      vocab_size: 128256
    num_loading_workers: 0
    seed: 6
  name: stable
  start_training_step: 1
general:
  benchmark_csv_path: null
  consumed_train_samples: null
  ignore_sanity_checks: true
  project: smollm3-3B-final
  run: elie-48n-1704-part1
  seed: 6
  step: null
logging:
  iteration_step_info_interval: 1
  log_level: info
  log_level_replica: info
model:
  ddp_bucket_cap_mb: 50
  dtype: bfloat16
  init_method:
    std: 0.02
  make_vocab_size_divisible_by: 1
  model_config:
    _attn_implementation: flash_attention_2
    _fused_rms_norm: true
    _fused_rotary_emb: true
    _use_doc_masking: true
    _use_qkv_packed: true
    attention_bias: false
    bos_token_id: 128000
    eos_token_id: 128001
    flex_attention_mask: null
    hidden_act: silu
    hidden_size: 2048
    initializer_range: 0.02
    intermediate_size: 11008
    is_qwen2_config: true
    max_position_embeddings: 4096
    moe_config: null
    num_attention_heads: 16
    num_hidden_layers: 36
    num_key_value_heads: 4
    pad_token_id: null
    pretraining_tp: 2
    rms_norm_eps: 1.0e-06
    rope_interleaved: false
    rope_scaling: null
    rope_theta: 50000.0
    sliding_window_size: null
    tie_word_embeddings: true
    use_cache: true
    vocab_size: 128256
    z_loss_coefficient: 1.0e-05
    z_loss_enabled: false
    no_rope_layer: 4
optimizer:
  accumulate_grad_in_fp32: true
  clip_grad: 1.0
  learning_rate_scheduler:
    learning_rate: 0.0002
    lr_decay_starting_step: 4198001
    lr_decay_steps: 522000
    lr_decay_style: linear
    lr_warmup_steps: 2000
    lr_warmup_style: linear
    min_decay_lr: 0
  optimizer_factory:
    adam_beta1: 0.9
    adam_beta2: 0.95
    adam_eps: 1.0e-08
    name: adamW
    torch_adam_is_fused: true
  weight_decay: 0.1
  weight_decay_exclude_named_params:
  - .*token_embedding.*
  zero_stage: 0
parallelism:
  context_parallel_size: 1
  dp: 192
  expert_parallel_size: 1
  moe_layer_recompute: false
  pp: 1
  pp_engine: 1f1b
  recompute_layer: false
  tp: 2
  tp_linear_async_communication: true
  tp_mode: REDUCE_SCATTER
  tp_recompute_allgather: true
profiler: null
s3_upload:
  remove_after_upload: true
  s5cmd_concurrency: 5
  s5cmd_numworkers: 16
  s5cmd_path: /fsx/elie_bakouch/smollm3_training/1004-nn/1004-hope/bin/s5cmd
  upload_s3_path: s3://smollm3/tp-fix-final-pre-training/1704-48n-part1
tokenizer:
  tokenizer_max_length: 4096
  tokenizer_name_or_path: meta-llama/Llama-3.2-1B
  tokenizer_revision: null
metrics_logging:
  log_level: 1
  log_detail_interval: 200
tokens:
  batch_accumulation_per_replica: 1
  limit_test_batches: 0
  limit_val_batches: 0
  micro_batch_size: 3
  sequence_length: 4096
  train_steps: 4720000
  val_check_interval: 100
lighteval:
  slurm_script_dir: "/fsx/elie_bakouch/smollm3_training/1004-nn/nanotron/_final1004/eval_results/launch-config"  # Default path for launch scripts
  logs_path: "/fsx/elie_bakouch/smollm3_training/1004-nn/nanotron/_final1004/eval_results/logs"  # Default path for evaluation logs
  local_checkpoint_dir: "/fsx/elie_bakouch/smollm3_training/1004-nn/nanotron/_final1004/evals-ckpt"  # Default path for temporary checkpoint storage. Will store under {local_checkpoint_dir}/{run_name}/{step}
  nanotron_path: "/fsx/elie_bakouch/smollm3_training/1004-nn/nanotron"
  output_dir: "/fsx/elie_bakouch/smollm3_training/1004-nn/nanotron/_final1004/eval_results/results"
  s3_save_path: "s3://smollm3/eval-prod"
  upload_to_wandb: true
  wandb_project: smollm3-3B-evals
  wandb_entity: huggingface
  parallelism:
    dp: 4
    pp: 1
    tp: 2
    tp_linear_async_communication: true
  batch_size: 8  # Optional batch size for evaluation
  eval_config_override: "/fsx/elie_bakouch/smollm3_training/1004-nn/nanotron/_final1004/smollm3_eval.yaml"  # The evaluation config file to use
  eval_interval: 4000
  eval_interval_file: null
  slurm:
    gpus_per_node: 8
    partition: "hopper-prod"
    hf_cache: "/fsx/elie_bakouch/.cache/huggingface"
    cpus_per_task: 88
    qos: "normal"
    time: "01:59:00"
    reservation: smollm  # Optional reservation name, can be null