Yy/qwen36 upgrade

pyproject.toml

@@ -33,6 +33,10 @@ warn_unused_ignores = false
 module = "torch.distributed.*,torch.optim.*"
 implicit_reexport = true
 
+[[tool.mypy.overrides]]
+module = "causal_conv1d.*,fla.*"
+ignore_missing_imports = true
+
 [[tool.mypy.overrides]]
 module = "fairseq2.data.parquet.*,tests.unit.data.parquet.*,tests.integration.parquet.*"
 ignore_errors = true

recipes/lm/sft/configs/qwen35_0.8b_gsm8k.yaml

@@ -0,0 +1,81 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Qwen 3.5 0.8B GSM8K SFT Fine-tuning Config
+#
+# Validates training recipe integration and loss convergence for the Qwen 3.5
+# model on the GSM8K math reasoning dataset.
+#
+# Usage:
+#   torchrun --standalone --nproc_per_node=8 -m recipes.lm.sft \
+#     --config-file recipes/lm/sft/configs/qwen35_0.8b_gsm8k.yaml \
+#     /path/to/output_dir
+
+model:
+  name: "qwen35_0.8b"
+  dtype: bfloat16
+  config_overrides:
+    pad_idx: 248044
+
+tokenizer:
+  name: "qwen35_0.8b"
+  config_overrides:
+    use_im_end: true
+
+dataset:
+  max_seq_len: 4096
+  max_num_tokens: 8192
+  valid_split: "sft_test"
+  chat_mode: false
+  config_overrides:
+    sources:
+      train:
+      - path: "hg://facebook/fairseq2-lm-gsm8k"
+        split: "sft_train"
+        weight: 1.0
+      sft_test:
+      - path: "hg://facebook/fairseq2-lm-gsm8k"
+        split: "sft_test"
+        weight: 1.0
+
+trainer:
+  data_parallelism: fsdp
+  max_grad_norm: 1.0
+  mixed_precision:
+    mode: static
+    dtype: bfloat16
+
+optimizer:
+  name: adamw
+  config:
+    lr: 2.0e-5
+    betas: [0.9, 0.95]
+    weight_decay: 0.1
+    impl: fused
+
+lr_scheduler:
+  name: cosine_annealing
+  config:
+    final_lr_scale: 0.1
+    num_warmup_steps: 100
+
+regime:
+  num_steps: 100000
+  checkpoint_every_n_steps: 100
+  validate_every_n_steps: 100
+  keep_last_n_checkpoints: 10
+  publish_metrics_every_n_steps: 1
+  save_model_only: false
+
+common:
+  seed: 0
+  metric_recorders:
+    wandb:
+      enabled: true
+      entity: "yunchaoyang1"
+      project: "fairseq2"
+    tensorboard:
+      enabled: false

recipes/lm/train/configs/llama3_2_1b_fineweb_edu_10bt.yaml

@@ -0,0 +1,162 @@
+# LLaMA 3.2 1B Continued Pretraining on FineWeb-Edu 10BT
+#
+# Loads the pretrained LLaMA 3.2 1B checkpoint and continues training
+# on the FineWeb-Edu Sample 10BT educational text dataset.
+#
+# Prerequisites:
+#   1. Convert parquet data to chunked JSONL:
+#      cd /checkpoint/smallomnillm/shared/data/fineweb-edu
+#      python convert_to_jsonl.py --sample 10BT --chunk-format --lightweight \
+#          --num-shards 256 --output-dir jsonl/10BT
+#
+# Usage:
+#   torchrun --standalone --nproc_per_node=8 -m recipes.lm.train \
+#     --config-file recipes/lm/train/configs/llama3_2_1b_fineweb_edu_10bt.yaml \
+#     /path/to/output_dir
+
+model:
+  name: llama3_2_1b
+  dtype: bfloat16
+  mmap: false
+  compile: false
+  compile_options:
+    fullgraph: false
+    dynamic: false
+    mode: default
+    backend: inductor
+    backend_options: null
+
+dataset:
+  name: null
+  family: lm_train
+  config_overrides:
+    sources:
+      - path: /checkpoint/smallomnillm/shared/data/fineweb-edu/jsonl/10BT
+        weight: 1.0
+  max_seq_len: 4096
+  max_num_tokens: 8192
+  prefetch: 4
+  sync_ranks: false
+
+tokenizer:
+  name: llama3
+  path: null
+  family: null
+  config_overrides: null
+
+gang:
+  tensor_parallel_size: 1
+  timeout: 15
+  high_priority: true
+
+trainer:
+  data_parallelism: fsdp
+  fsdp:
+    version: v2
+    granularity: layer
+    hybrid: false
+    reshard_after_forward: true
+    fp32_reduce: true
+  mixed_precision:
+    mode: static
+    dtype: bfloat16
+  grad_accumulation:
+    num_batches: 1
+    no_sync: false
+  activation_checkpointing:
+    mode: layerwise
+    every_nth_layer: 1
+  max_grad_norm: 1.0
+  fp16_loss_scale:
+  - 128.0
+  - 0.0001
+  gc_every_n_steps: 1000
+  grad_check: false
+  anomaly_detection: false
+
+optimizer:
+  name: adamw
+  config:
+    lr: 5.0e-5
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-8
+    weight_decay: 0.1
+    amsgrad: false
+    maximize: false
+    capturable: false
+    differentiable: false
+    impl: fused
+    groups: []
+
+lr_scheduler:
+  name: cosine_annealing
+  config:
+    cycle_len: null
+    num_warmup_steps: 500
+    cycle_mul: 1.0
+    lr_mul: 1.0
+    start_lr: 1.0e-30
+    final_lr: null
+    final_lr_scale: 0.01
+
+regime:
+  num_steps: 76000
+  num_data_epochs: null
+  validate_at_start: false
+  validate_after_n_steps: 0
+  validate_every_n_steps: 4000
+  validate_after_n_data_epochs: 0
+  validate_every_n_data_epochs: null
+  score_metric: null
+  checkpoint_after_n_steps: 0
+  checkpoint_every_n_steps: 2000
+  checkpoint_after_n_data_epochs: 0
+  checkpoint_every_n_data_epochs: null
+  save_model_only: all_but_last
+  export_hugging_face: true
+  keep_last_n_checkpoints: 3
+  keep_best_n_checkpoints: null
+  keep_checkpoint_every_n_steps: 10000
+  publish_metrics_after_n_steps: 0
+  publish_metrics_every_n_steps: 10
+  publish_metrics_after_n_data_epochs: 0
+  publish_metrics_every_n_data_epochs: null
+
+common:
+  torch:
+    num_threads: null
+    allow_tf32: true
+    fp16_reduced_precision: true
+    bf16_reduced_precision: true
+    default_sdpa: torch
+    compiled_region_activation_memory_budget: 0.9
+  metric_recorders:
+    tensorboard:
+      enabled: true
+    wandb:
+      enabled: true
+      entity: yunchaoyang1
+      project: llama3_2_1b_fineweb_edu_10bt
+      run_id: persistent
+      run_name: null
+      group: null
+      job_type: null
+      resume_mode: null
+  profilers:
+    torch:
+      enabled: false
+      skip_n_steps: 4
+      wait_n_steps: 0
+      num_warmup_steps: 1
+      num_active_steps: 4
+      repeat: 1
+  assets:
+    extra_paths: []
+    prev_checkpoint_dir: null
+  seed: 2
+  debug: false
+  cluster: auto
+  no_sweep_dir: false
+  sweep_format: ws_{world_size}.{hash}