export WANDB_NAME="lact-8k-rope-swa2048-nh4-fw02-id-rank32-momentum-140M-8K-B-bs32-lr1e-3-cosine-warmup1024-steps4-init0.02.1gpu"
export HF_DATASETS_CACHE="/home/uceeppm/Data"

bash train.sh \
  --job.config_file flame/models/fla.toml \
  --job.dump_folder exp/lact/lact-swa2048-rope-fw02-id-rank32-init0.5-gain0.5-nh4-momentum-140M-8K-B/batch1.seqlen32768.bs32.warmup1024.update1.steps4.lr1e-3.cosine.1gpu \
  --model.config configs/140M_lact_swiglu_nh4_fwlow_rank_momentum.json \
  --model.tokenizer_path fla-hub/transformer-1.3B-100B \
  --optimizer.name AdamW \
  --optimizer.eps 1e-15 \
  --optimizer.lr 1e-3 \
  --lr_scheduler.warmup_steps 1024 \
  --lr_scheduler.lr_min 0.1 \
  --lr_scheduler.decay_type cosine \
  --training.batch_size 4 \
  --training.seq_len 8192 \
  --training.context_len 8192 \
  --training.gradient_accumulation_steps 1 \
  --training.steps 40 \
  --training.max_norm 1.0 \
  --training.skip_nan_inf \
  --training.dataset Salesforce/wikitext \
  --training.dataset_name wikitext-2-raw-v1 \
  --training.dataset_split train \
  --training.num_workers 2 \
  --training.prefetch_factor 1 \
  --training.seed 42 \
  --training.compile \
  --metrics.log_freq 10 \
  --profiling.profile_freq 2000 \