From 13ddc8dc583d8b1355909970cb8c27f85b7d3c8b Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Mon, 9 Feb 2026 11:00:39 -0600 Subject: Initial implementation: DAGFormer Phase 1 - olmo_graph.py: Modified OLMo2-1B forward with per-head routing via 256x256 adjacency matrix A - Proportional attribution for post-norm decomposition - All 6 GPU sanity checks pass (baseline diff = 0.000001) - predictor.py: Qwen3-Embedding encoder + MLP decoder + Gumbel-Sigmoid + cascading gate - pipeline.py: End-to-end glue (predictor -> A -> OLMo -> NLL) - trainer.py: Full training loop with DDP, gradient accumulation, eval, checkpointing - dolma.py: Streaming Dolma v1.7 with sequence packing - 43/43 unit tests pass Co-Authored-By: Claude Opus 4.6 --- configs/ablation_lambda.yaml | 0 configs/ablation_rank.yaml | 0 configs/ablation_tau.yaml | 0 configs/phase1_full.yaml | 0 configs/sanity_check.yaml | 50 ++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 50 insertions(+) create mode 100644 configs/ablation_lambda.yaml create mode 100644 configs/ablation_rank.yaml create mode 100644 configs/ablation_tau.yaml create mode 100644 configs/phase1_full.yaml create mode 100644 configs/sanity_check.yaml (limited to 'configs') diff --git a/configs/ablation_lambda.yaml b/configs/ablation_lambda.yaml new file mode 100644 index 0000000..e69de29 diff --git a/configs/ablation_rank.yaml b/configs/ablation_rank.yaml new file mode 100644 index 0000000..e69de29 diff --git a/configs/ablation_tau.yaml b/configs/ablation_tau.yaml new file mode 100644 index 0000000..e69de29 diff --git a/configs/phase1_full.yaml b/configs/phase1_full.yaml new file mode 100644 index 0000000..e69de29 diff --git a/configs/sanity_check.yaml b/configs/sanity_check.yaml new file mode 100644 index 0000000..92d4a24 --- /dev/null +++ b/configs/sanity_check.yaml @@ -0,0 +1,50 @@ +# Sanity check config — verify baseline NLL reproduction and basic training +# Run: python scripts/train.py --config configs/sanity_check.yaml + +# Model +olmo_model_id: "allenai/OLMo-2-0425-1B" +qwen_model_id: "Qwen/Qwen3-Embedding-0.6B" + +# Predictor +predictor_hidden_dim: 1024 +predictor_rank: 32 +cascading_gate_k: 5.0 +input_norm: "none" # use "none" to verify baseline reproduction + +# Data +dataset: "allenai/dolma" +dataset_name: "v1_7" +seq_len: 1024 +batch_size: 4 +micro_batch_size: 2 # gradient accumulation: effective batch=4, micro=2 +qwen_input_prefix: "" + +# Eval +eval_skip: 10000 # reduced for sanity check (1M too slow for streaming) +eval_size: 50 # small eval set for sanity check + +# Training +total_steps: 1000 +lr: 3e-4 +weight_decay: 0.01 +optimizer: "adamw" + +# Schedules +tau_init: 5.0 +tau_final: 0.2 +tau_schedule: "cosine" +lambda_max: 0.0 # no sparsity for sanity check +lambda_warmup_frac: 0.2 + +# Logging +wandb_project: "dagformer" +wandb_run_name: "sanity-check" +log_every: 10 +eval_every: 100 + +# Checkpointing +save_every: 500 +save_dir: "checkpoints/" + +# Hardware +num_gpus: 1 -- cgit v1.2.3