Add transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338

.gitattributes

@@ -483,3 +483,12 @@ transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_marsadamw_lr3e_3_b1_0_95_b2_0_99
 transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_marsadamw_lr3e_3_b1_0_95_b2_0_99_eps_1e_15_20260511_072915/exp_data/checkpoint/step-30720/__5_0.distcp filter=lfs diff=lfs merge=lfs -text
 transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_marsadamw_lr3e_3_b1_0_95_b2_0_99_eps_1e_15_20260511_072915/exp_data/checkpoint/step-30720/__6_0.distcp filter=lfs diff=lfs merge=lfs -text
 transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_marsadamw_lr3e_3_b1_0_95_b2_0_99_eps_1e_15_20260511_072915/exp_data/checkpoint/step-30720/__7_0.distcp filter=lfs diff=lfs merge=lfs -text
+transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/.metadata filter=lfs diff=lfs merge=lfs -text
+transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
+transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__1_0.distcp filter=lfs diff=lfs merge=lfs -text
+transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__2_0.distcp filter=lfs diff=lfs merge=lfs -text
+transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__3_0.distcp filter=lfs diff=lfs merge=lfs -text
+transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__4_0.distcp filter=lfs diff=lfs merge=lfs -text
+transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__5_0.distcp filter=lfs diff=lfs merge=lfs -text
+transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__6_0.distcp filter=lfs diff=lfs merge=lfs -text
+transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__7_0.distcp filter=lfs diff=lfs merge=lfs -text

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023-2025 Songlin Yang, Yu Zhang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/README.md

@@ -0,0 +1,519 @@
+<div align="center">
+
+# 🔥 Flame: Flash Language Modeling Made Easy
+
+[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/fla-org/flame)
+
+</div>
+
+Welcome to 🔥 `flame`, a minimal and efficient framework built on `torchtitan` for language models with blazing efficiency.
+
+**Feature Highlights:**
+
+- 🚀 Minimal, easy-to-use, extensible training framework
+- 🤗 Seamless integration with `fla` and `transformers`
+- 🔄 Zero-cost data preprocessing: online tokenization, dataset shuffling, and multiple datasets support
+- 🔮 4D parallelism (coming soon)
+
+## Setup
+
+To get started, clone the `flame` repository and install the required dependencies:
+
+```bash
+git clone https://github.com/fla-org/flame.git
+cd flame
+pip install .
+```
+
+Install the latest version of fla
+```
+pip uninstall flash-linear-attention && pip install -U --no-use-pep517 git+https://github.com/fla-org/flash-linear-attention
+```
+
+[Important] Install specific version of torchtitan
+```
+pip install git+https://github.com/pytorch/torchtitan.git@0b44d4c
+```
+
+
+## Dataset Preparation
+To download the dataset to your local disk, create a new Python file with the following content and execute it:
+
+```py
+from datasets import load_dataset
+
+# load fineweb-edu with parallel processing
+dataset = load_dataset("HuggingFaceFW/fineweb-edu", name="default", num_proc=64, cache_dir="/your/cache/path")
+
+# or load a subset with roughly 100B tokens, suitable for small- or medium-sized experiments
+dataset = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-100BT", num_proc=64, cache_dir="/your/cache/path")
+```
+
+## Training Recipes
+
+Here's an example of training a 340M FLA Transformer model with a LLaMA-like architecture from scratch on a 100BT subset of the Fineweb-edu corpus ~~in streaming mode~~. (Do not use streaming mode if you are concerned about resuming training.)
+ 
+> [!WARNING]
+> If the dataset is not downloaded beforehand, the streaming mode will attempt to fetch it from a remote server and download it on-the-fly, which can be highly unstable during training due to network issues.  
+> For stable training, ensure the dataset is downloaded locally (see [**Dataset Preparation**](#dataset-preparation)). Otherwise, we assume you are only testing the new corpus.
+
+```sh
+bash train.sh \
+  --job.config_file flame/models/fla.toml \
+  --job.dump_folder exp/transformer-340M-4K-10B/batch1.seqlen65536.context4096.warmup1024.update1.steps20480.lr1e-3.cosine \
+  --model.config configs/transformer_340M.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 1 \
+  --training.seq_len 65536 \
+  --training.context_len 4096 \
+  --training.varlen \
+  --training.gradient_accumulation_steps 1 \
+  --training.steps 20480 \
+  --training.max_norm 1.0 \
+  --training.skip_nan_inf \
+  --training.dataset HuggingFaceFW/fineweb-edu \
+  --training.dataset_name sample-100BT \
+  --training.dataset_split train \
+  --training.num_workers 32 \
+  --training.prefetch_factor 2 \
+  --training.seed 42 \
+  --training.compile \
+  --checkpoint.interval 2048 \
+  --checkpoint.load_step -1 \
+  --checkpoint.keep_latest_k 2 \
+  --metrics.log_freq 1
+```
+
+You can specify the number of GPUs by setting the environment variable `NGPU`, which defaults to 8.  
+**For single-GPU debugging, set `NGPU=1`.**
+
+We provide several [config files](https://github.com/fla-org/flame/tree/main/configs) for different models.
+By default, the learning rate is set to 1e-3 with a cosine scheduler. Other schedulers, such as WSD (wsd), are also supported.
+
+**Key parameters:**
+- `--lr_scheduler.decay_ratio`: The proportion of the steps allocated to the decay phase. The learning rate will remain stable after the warmup period and only start decaying during the last `decay_ratio` portion of the total training steps, which is known as the Warmup-Stable-Decay (WSD) schedule.
+- `--lr_scheduler.warmup_steps`: The number of steps for the learning rate warmup phase.
+- `--training.steps`: Total number of training steps.
+- `--training.batch_size`: Batch size per device, must be 1 if `--training.varlen` is set.
+- `--training.seq_len`: The length of each sequence in the batch, which is concatenated from multiple samples.
+- `--training.context_len`: The max allowed length of a sample. For non-varlen mode, this is equivalent to `seq_len`.
+- `--training.varlen`: Whether to conduct variable-length sequence training.
+- `--training.gradient_accumulation_steps`: Number of gradient accumulation steps.
+
+> [!WARNING]
+> The total number of tokens processed per batch, referred to as `global_batch_size`, is calculated as batch_size × gradient_accumulation_steps × num_gpus.
+> Each step processes `global_batch_size * seq_len` tokens.  
+> Monitor the value of `global_batch_size`, `warmup_steps`, and `steps` carefully when modifying any of the hyperparameters!
+
+For a detailed explanation of all parameters, run:
+
+```sh
+bash train.sh -h
+```
+
+<details>
+<summary>Usage</summary>
+
+```py
+options:
+  -h, --help            show this help message and exit
+  --job.config_file JOB.CONFIG_FILE
+                        Job config file
+  --job.dump_folder JOB.DUMP_FOLDER
+                        Folder to dump job outputs
+  --job.description JOB.DESCRIPTION
+                        Description of the job
+  --job.use_for_integration_test
+                        Add this config to the integration test suite
+  --job.print_args      Print the args to terminal
+  --model.config MODEL.CONFIG
+                        Path to the model config
+  --model.norm_type MODEL.NORM_TYPE
+                        Type of layer normalization to use [layernorm,
+                        np_layernorm, rmsnorm, fused_rmsnorm]
+  --model.tokenizer_path MODEL.TOKENIZER_PATH
+                        Tokenizer path
+  --profiling.enable_profiling
+                        Whether to enable pytorch profiler
+  --profiling.save_traces_folder PROFILING.SAVE_TRACES_FOLDER
+                        Trace files location
+  --profiling.profile_freq PROFILING.PROFILE_FREQ
+                        How often to collect profiler traces, in iterations
+  --profiling.enable_memory_snapshot
+                        Whether to dump memory snapshot
+  --profiling.save_memory_snapshot_folder PROFILING.SAVE_MEMORY_SNAPSHOT_FOLDER
+                        Memeory snapshot files location
+  --optimizer.name OPTIMIZER.NAME
+                        Optimizer to use
+  --optimizer.eps OPTIMIZER.EPS
+                        Epsilon value for the optimizer.
+  --optimizer.fused     Whether the fused implementation(CUDA only) is used.
+  --optimizer.scheduler {wsd,cosine,linear}
+                        Scheduler to use. Currently supported: wsd, cosine,
+                        and linear.
+  --optimizer.lr OPTIMIZER.LR
+                        Learning rate to use
+  --optimizer.min_lr_ratio OPTIMIZER.MIN_LR_RATIO
+                        Min lr ratio for lr scheduler
+  --optimizer.early_step_in_backward
+                        Whether to apply optimizer in the backward. Caution,
+                        optimizer_in_backward is not compatible with gradients
+                        clipping, users should not call
+                        register_post_accumulate_grad_hook after the optimizer
+                        is built.
+  --training.batch_size TRAINING.BATCH_SIZE
+                        Batch size
+  --training.seq_len TRAINING.SEQ_LEN
+                        Sequence length
+  --training.context_len TRAINING.CONTEXT_LEN
+                        Max length allowed for each sequence
+  --training.varlen     Whether to take sequences of variable length as input
+  --training.warmup_steps TRAINING.WARMUP_STEPS
+                        Steps for lr scheduler warmup, normally 1/5 of
+                        --training.steps
+  --training.gradient_accumulation_steps TRAINING.GRADIENT_ACCUMULATION_STEPS
+                        Number of steps to accumulate gradients before
+                        updating parameters
+  --training.steps TRAINING.STEPS
+                        How many train steps to run
+  --training.max_norm TRAINING.MAX_NORM
+                        Max norm for gradient clipping
+  --training.skip_nan_inf
+                        Skip batch updates when NaN or INF gradients are
+                        encountered during training
+  --training.dataset TRAINING.DATASET
+                        Dataset to use, with comma separated values
+  --training.dataset_name TRAINING.DATASET_NAME
+                        The name of the dataset config, with comma separated
+                        values if provided
+  --training.dataset_split TRAINING.DATASET_SPLIT
+                        Dataset split to use, with comma separated values if
+                        provided
+  --training.data_dir TRAINING.DATA_DIR
+                        Data dirs to use, with comma separated values if
+                        provided
+  --training.data_files TRAINING.DATA_FILES
+                        Data files to use, with comma separated values if
+                        provided
+  --training.data_probs TRAINING.DATA_PROBS
+                        Data sampling probabilities, with comma separated
+                        values if provided
+  --training.streaming  Whether to load dataset in streaming mode, used for
+                        huge dataset
+  --training.num_workers TRAINING.NUM_WORKERS
+                        Number of subprocesses to use for data loading. 0
+                        means that the data will be loaded in the main
+                        process.
+  --training.prefetch_factor TRAINING.PREFETCH_FACTOR
+                        Number of batches loaded in advance by each worker.2
+                        means there will be a total of 2 * num_workers batches
+                        prefetched across all workers.
+  --training.data_parallel_replicate_degree TRAINING.DATA_PARALLEL_REPLICATE_DEGREE
+                        The `data_parallel_replicate_degree` argument
+                        specifies the degree of data parallelism for weight
+                        replication. When this value is greater than 1,
+                        weights will be replicated across
+                        `data_parallel_replicate_degree` ranks. If
+                        `data_parallel_shard_degree` is also greater than 1,
+                        the parallelism method used is HSDP (Hybrid Sharded
+                        Data Parallelism). Otherwise, the parallelism method
+                        used is DDP (Distributed Data Parallelism). 1 means
+                        disabled.
+  --training.data_parallel_shard_degree TRAINING.DATA_PARALLEL_SHARD_DEGREE
+                        The `data_parallel_shard_degree` argument specifies
+                        the degree of data parallelism for weight sharding.
+                        When this value is greater than 1, weights will be
+                        sharded across `data_parallel_shard_degree` ranks. If
+                        `data_parallel_replicate_degree` is also greater than
+                        1, the parallelism method used is HSDP (Hybrid Sharded
+                        Data Parallelism). Otherwise, the parallelism method
+                        used is FSDP (Fully Sharded Data Parallelism). -1
+                        means leftover ranks will be used (After
+                        DP_REPLICATE/SP/PP). Note that only
+                        `data_parallel_shard_degree` can be negative. 1 means
+                        disabled.
+  --training.enable_cpu_offload
+                        Whether to apply CPU offloading of parameters,
+                        gradients, and optimizer states in FSDP
+  --training.tensor_parallel_degree TRAINING.TENSOR_PARALLEL_DEGREE
+                        Tensor Parallelism degree. 1 means disabled.
+  --training.disable_loss_parallel
+                        Whether to apply loss parallel when sequence parallel
+                        is enabled
+  --training.mixed_precision_param {bfloat16,float32}
+                        torch dtype to use for parameters when applying mixed
+                        precision via FSDP. This feature only takes effect
+                        when data_parallel_shard_degree > 1
+  --training.mixed_precision_reduce {float32}
+                        torch dtype to use for reductions when applying mixed
+                        precision via FSDP. This feature only takes effect
+                        when data_parallel_shard_degree > 1
+  --training.compile    Whether to compile the model
+  --training.gc_freq TRAINING.GC_FREQ
+                        Python garbage control scheduling interval, in steps
+  --training.seed TRAINING.SEED
+                        Choose the base RNG seed used for training
+  --training.deterministic
+                        Use deterministic algorithms wherever possible, may be
+                        slower
+  --metrics.log_freq METRICS.LOG_FREQ
+                        How often to log metrics to TensorBoard, in iterations
+  --metrics.enable_tensorboard
+                        Whether to log metrics to TensorBoard
+  --metrics.disable_color_printing
+                        Whether to disable color printing in logs
+  --metrics.save_tb_folder METRICS.SAVE_TB_FOLDER
+                        Folder to dump TensorBoard states
+  --metrics.rank_0_only
+                        Whether to save TensorBoard metrics only for rank 0 or
+                        for all ranks. When pipeline_parallel_degree is > 1,
+                        this option uses the 0th rank of the last stage
+                        pipeline group, which is the only stage that computes
+                        loss metrics.
+  --metrics.enable_wandb
+                        Whether to log metrics to Weights & Biases
+  --experimental.enable_async_tensor_parallel
+                        Whether to apply async tensor parallel (currently only
+                        effective when compile is enabled)
+  --experimental.pipeline_parallel_degree EXPERIMENTAL.PIPELINE_PARALLEL_DEGREE
+                        Pipeline Parallelism degree, or number of ranks. 1
+                        means disabled. If using looped schedules, this still
+                        specifies the number of physical ranks, not the number
+                        of stages. Stages per rank are inferred from split
+                        points degree, and schedule.
+  --experimental.pipeline_parallel_split_points EXPERIMENTAL.PIPELINE_PARALLEL_SPLIT_POINTS [EXPERIMENTAL.PIPELINE_PARALLEL_SPLIT_POINTS ...]
+                        Specify comma-separated names of modules to use as the
+                        beginning of a split point. e.g. "layers.0,layers.2"
+                        will cause the model to be split into 3 stages, the
+                        first containing all the layers up to layers.0, the
+                        second containing layers.0 and up to layers.2, the
+                        third containing layers.2 and all the remaining
+                        layers. Note: fully-automated splitting may be enabled
+                        in the future, but currently the split points must be
+                        specified manually.
+  --experimental.pipeline_parallel_schedule EXPERIMENTAL.PIPELINE_PARALLEL_SCHEDULE
+                        Specify the Pipeline Parallel schedule to use. The
+                        supported schedules are: https://github.com/pytorch/py
+                        torch/blob/de4c2a3b4e89d96334dc678d1c3f2ae51a6630a0/to
+                        rch/distributed/pipelining/schedules.py#L2161. The
+                        schedule must be compatible with the split points and
+                        stages_per_rank. Looped schedules (e.g.
+                        Interleaved1F1B) require specifying
+                        pipeline_parallel_degree = number of ranks, and
+                        split_points = number of stages - 1
+  --experimental.pipeline_parallel_schedule_csv EXPERIMENTAL.PIPELINE_PARALLEL_SCHEDULE_CSV
+                        Specify the path to the pipeline parallel schedule csv
+                        file to use. The pipeline_parallel_schedule argument
+                        must be either PipelineScheduleSingle,
+                        PipelineScheduleMulti, or _PipelineScheduleRuntime.
+  --experimental.pipeline_parallel_microbatches EXPERIMENTAL.PIPELINE_PARALLEL_MICROBATCHES
+                        How many microbatches to split the global training
+                        batch into when using pipeline parallelism. The global
+                        training batch size must be evenly divisible by the
+                        number of microbatches. The default value will be the
+                        number of pipeline stages, if unspecified.
+  --experimental.enable_compiled_autograd
+                        Enable CompiledAutograd to compile the backward.
+  --experimental.context_parallel_degree EXPERIMENTAL.CONTEXT_PARALLEL_DEGREE
+                        Context parallelism degree. 1 means disabled.
+  --experimental.context_parallel_rotate_method EXPERIMENTAL.CONTEXT_PARALLEL_ROTATE_METHOD
+                        The collective to use in context parallel SDPA for kv
+                        shards exchange. 'allgather' means to all-gather all
+                        kv shards on ranks after the first sub-SDPA
+                        computation, 'alltoall' means to all-to-all shuffle
+                        the kv shards. The default value is 'allgather'.
+  --checkpoint.enable_checkpoint
+                        Whether to enable checkpoint
+  --checkpoint.folder CHECKPOINT.FOLDER
+                        The folder to store the checkpoints. When
+                        enable_checkpoint is set to true, checkpoints will be
+                        in {--job.dump_folder}/{--checkpoint.folder}.
+  --checkpoint.interval_type CHECKPOINT.INTERVAL_TYPE
+                        Checkpointing interval unit of measurement ['step',
+                        'seconds']
+  --checkpoint.interval CHECKPOINT.INTERVAL
+                        Checkpointing interval, in steps or seconds depending
+                        on --checkpoint.interval_type
+  --checkpoint.model_weights_only
+                        When model_weights_only=True, only model weights will
+                        be saved at the end of training. With this,
+                        checkpoints can be loaded using `torch.load(...,
+                        weights_only=True)` after conversion. When
+                        model_weights_only=False, the full checkpoint will be
+                        saved. A full checkpoint includes model, optimizer and
+                        train_state, which can be used to resume training. The
+                        default value is false.
+  --checkpoint.export_dtype {float16,bfloat16,float32}
+                        Converts to the specified precision when training
+                        completes and model_weights_only=true. Currently
+                        supports float32, float16, and bfloat16. The default
+                        value is float32.
+  --checkpoint.create_seed_checkpoint
+                        Initializes the full model without applying
+                        parallelisms, and then saves it as a seed checkpoint.
+                        Note: requires user to call train.py without
+                        specifying any parallelisms, e.g. NGPU=1. Could be
+                        implemented as a separate script, but this way shares
+                        more code.
+  --checkpoint.async_mode CHECKPOINT.ASYNC_MODE
+                        Which async checkpoint mode to use. Currently there
+                        are 3 different modes. 1. "disabled": synchronized
+                        checkpointing will be used. 2. "async":
+                        torch.distributed.checkpoint.async_save will be used.
+                        1. "async_with_pinned_mem": this option utilizes a
+                        dedicated pinned memory space and creates a separate
+                        process for faster GPU->CPU transfer performance and
+                        eliminating GIL contention. The cost is increased CPU
+                        memory usage. If insufficient CPU memory is available,
+                        performance may degrade due to memory paging. For most
+                        users, "async" should suffice as the performance
+                        overhead is typically small (on the order of tens of
+                        seconds) compared to checkpointing frequency. This
+                        mode can be employed to pursue near-zero checkpointing
+                        times (e.g., < 1 second) given appropriate hardware
+                        support such as ample CPU memory and fast PCIe.
+                        "disabled" is the default mode.
+  --checkpoint.keep_latest_k CHECKPOINT.KEEP_LATEST_K
+                        Keeps only the latest k checkpoints, and purging older
+                        ones. If 0, keep all checkpoints. 0 is the default
+                        value.
+  --checkpoint.load_step CHECKPOINT.LOAD_STEP
+                        Load the checkpoint at the specified step. If -1, load
+                        the latest checkpoint.
+  --float8.enable_float8_linear
+                        If true, swaps `torch.nn.Linear` with `Float8Linear`.
+                        This feature requires you to install 'torchao' which
+                        can be found here: https://github.com/pytorch/ao
+  --float8.enable_fsdp_float8_all_gather
+                        Whether enable float8 all-gather in FSDP
+  --float8.precompute_float8_dynamic_scale_for_fsdp
+                        Whether precompute float8 scales dynamically for FSDP
+  --float8.scaling_type_input {dynamic,delayed}
+                        float8 scaling for input, dynamic (default) or delayed
+  --float8.scaling_type_weight FLOAT8.SCALING_TYPE_WEIGHT
+                        float8 scaling for input, dynamic (default) or delayed
+  --float8.scaling_type_grad_output FLOAT8.SCALING_TYPE_GRAD_OUTPUT
+                        float8 scaling for input, dynamic (default) or delayed
+  --comm.init_timeout_seconds COMM.INIT_TIMEOUT_SECONDS
+                        Timeout for communication operations, during
+                        initialization and first train step.
+  --comm.train_timeout_seconds COMM.TRAIN_TIMEOUT_SECONDS
+                        Timeout for communication operations after the first
+                        train step -- usually a tighter bound than during
+                        initialization.
+  --comm.trace_buf_size COMM.TRACE_BUF_SIZE
+                        Flight recorder ring buffer size, >0 means recording
+                        by default, 0 means disabled
+  --memory_estimation.enabled
+                        Whether to estimate memory usage for FSDP
+  --memory_estimation.disable_fake_mode
+                        Whether to estimate memory under FakeTensorMode
+```
+</details>
+
+### Training with variable-length inputs
+When you set the `--training.varlen` flag, you're enabling a more efficient training method that packs multiple documents together into a single long sequence, eliminating the need for padding. 
+This is particularly useful when your dataset contains documents of varying lengths. 
+Let's break down how `--training.seq_len` and `--training.context_len` work in this mode.
+
+* `--training.seq_len` (Packed Sequence Length): This is the total length of the final sequence fed to the model on one device. Instead of processing one document at a time, the dataloader takes multiple documents (each split to sequences no longer than `context_len`), concatenates them end-to-end, and creates a single long sequence of length `seq_len`.
+* `--training.context_len` (Sample Length): This parameter defines the maximum number of tokens for a single document or sample. If a document from the dataset is longer than `context_len`, it will be truncated. For example, if `--training.context_len` is set to 4,096, a document with 5,000 tokens will be cut down to its first 4,096 tokens, leaving the left tokens as another independent sequence, while a document with 3000 tokens remains unchanged.
+
+### Training with `torch.compile`
+
+Starting from `torch 2.0`, `torch.compile` has been introduced as a new feature to seamlessly accelerate training processes.
+In `flame`, one can simply enable `torch.compile` by adding `--training.compile` flag to your training script.
+
+However, `fla` has integrated numerous fused kernels for acceleration, which may potentially conflict with `torch.compile`.
+We are actively working on resolving these issues to make compilation transparent to users.
+In the meantime, please ensure you are using the latest dependencies.
+
+Specifically, **we recommend using `torch>=2.6` and `triton>=3.0`**.
+
+### Training with multiple datasets
+
+If you wish to train a model with all-round capabilities (e.g., code, math, and multilingual ability), it's necessary to train on multiple datasets.
+`flame` allows training with multiple datasets easily.
+For example, you can specify the following arguments to train on 6 datasets with different proportions:
+
+```sh
+  --training.dataset HuggingFaceFW/fineweb-edu,opencsg/Fineweb-Edu-Chinese-V2.1,OpenCoder-LLM/opc-fineweb-code-corpus,math-ai/AutoMathText,EleutherAI/proof-pile-2,OpenCoder-LLM/opc-fineweb-math-corpus   \
+  --training.data_probs 0.6,0.15,0.15,0.014,0.058,0.028     \
+```
+
+### ~Finalizing training~
+
+> [!NOTE]  
+> We have done this conversion automatically in the training script since our latest updates.
+
+Once training is complete, you may want to convert the distributed checkpoints (DCPs) into the 🤗 format for broader use.
+To facilitate this, we provide a straightforward conversion script:
+
+```sh
+python -m flame.utils.convert_dcp_to_hf --path <path_to_model> --step <step> --config <path_to_config> --tokenizer <path_to_tokenizer>
+```
+After this, your model will be in the 🤗 format, ready to be shared or deployed.
+You can then easily publish your model using the `huggingface_hub` for wider accessibility.
+
+### Continual training
+
+If you wish to build upon a strong pre-trained model (in 🤗 format) and continue training, we also offer a script to convert the 🤗 format model back into DCP format.
+This allows you to seamlessly resume training with `flame`.
+```sh
+python -m flame.utils.convert_hf_to_dcp --model <path_to_hf> --checkpoint <path_to_dcp/checkpoint/step-0>
+```
+Here, `<path_to_dcp>` is the directory where your distributed checkpoints will be stored.
+The checkpoint is intentionally saved at `<step-0>` within the checkpoint folder to ensure it is loadable by `flame` during the initial training step, similar to how a seed checkpoint is handled.
+
+Once the conversion is complete, you can proceed with training using `flame` as usual, continuing from where the pretrained model left off.
+
+## Multi-node training
+
+If you have access to multi-node GPUs, consider leveraging them for optimal performance.
+This process is straightforward and well-documented in the PyTorch [docs](https://pytorch.org/docs/stable/elastic/run.html).
+
+To set up multi-node training:
+* Set the environment variables `MASTER_ADDR=<ip>` and `MASTER_PORT=<port>` before running the training script across all nodes.
+* If you're using a job scheduler like Slurm, it will handle these variables for you.
+
+`torchtitan` provides a [Slurm script](https://github.com/pytorch/torchtitan/blob/main/multinode_trainer.slurm) for multi-node training, which you can use as a reference or starting point.
+
+## Custom models
+
+`flame` supports custom model architectures through seamless integration with the Hugging Face `transformers` library. To add your own model:
+
+1. Create a new model directory under `custom_models/` (see `custom_models/sba` for a complete example)
+2. Implement your model classes and configuration:
+   - Define a config class inheriting from `PretrainedConfig` (see `custom_models/sba/config_sba.py` for an example)
+   - Create model classes inheriting from `PreTrainedModel` (see `custom_models/sba/modeling_sba.py` for an example)
+3. Register your models in `__init__.py`:
+   - Import your model classes and config classes
+   - Register your models with the `AutoModelForCausalLM`, `AutoModel` and `AutoConfig` classes (see `custom_models/sba/__init__.py` for an example)
+4. Create a config file for your custom model, just need to specify the `model_type` to the one you just named for your custom model (example: `configs/sba_340m.json`).
+5. Training is extremely simple, you can just use the `flame.train.py` script to train your custom model.
+
+
+
+
+
+
+
+## Citation
+
+If you find `flame` helpful for your work, please consider citing it.
+
+```bib
+@software{yang2025flame,
+  title  = {Flame: Flash Language Modeling Made Easy},
+  author = {Zhang, Yu and Yang, Songlin},
+  url    = {https://github.com/fla-org/flame},
+  month  = jan,
+  year   = {2025}
+}
+```

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/.metadata

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c8eeb4fb92a760e8ebd173686dafcb1195996857abe30e2234f4d72008e5d7a2
+size 1411174

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__0_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10974149c457b268f3ad996002bd1b8d2a8baec49773c19fbf6e81f464ea1746
+size 2001640303

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__1_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6ddbd0ab3bd1163f7e7822c38c3a15a941cc9ad40d47e205cddae4c4a5f34e6
+size 1995395505

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__2_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23b6ad567d472445d5012f0a9e364b1fd43f8c0ebca15236912a0f7355ee38bb
+size 2001569722

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__3_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d3a1a15c84b55588b149230697334d3a99ab4a0d7ad238791798bbdbe14f80a
+size 1992747645

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__4_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:488ecb9a14890d61c18143655835ce910b20c188d7c3fd670f911d9c545927ac
+size 2217877888

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__5_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:998721db82dda199a9f0eb5fccd0984eed4f1856487fc30e607dde94f035c5c4
+size 2217843999

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__6_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:864911064ddd46130fe60104a1ab977dc785989309cdc75b787c1351d5a9d725
+size 1991234513

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/checkpoint/step-30720/__7_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7cb288436d4b820c78953e25bf2e8927315e16a284b82cfaccefa5b9b711ece9
+size 1996738056

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/config.json

@@ -0,0 +1,34 @@
+{
+  "architectures": [
+    "TransformerForCausalLM"
+  ],
+  "bos_token_id": 1,
+  "dtype": "float32",
+  "elementwise_affine": true,
+  "eos_token_id": 2,
+  "fuse_cross_entropy": true,
+  "fuse_linear_cross_entropy": false,
+  "fuse_norm": true,
+  "fuse_swiglu": true,
+  "hidden_act": "swish",
+  "hidden_ratio": 4,
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": null,
+  "max_position_embeddings": 8192,
+  "model_type": "transformer",
+  "norm_eps": 1e-06,
+  "num_heads": 32,
+  "num_hidden_layers": 24,
+  "num_kv_heads": null,
+  "pad_token_id": 2,
+  "qk_norm": false,
+  "qkv_bias": false,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.6",
+  "use_cache": true,
+  "use_l2warp": false,
+  "vocab_size": 32000,
+  "window_size": null
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/blt_transformer_1000hash.json

@@ -0,0 +1,98 @@
+{
+  "model_type": "blt",
+  "vocab_size": 260,
+  "max_position_embeddings": 4096,
+  "initializer_range": 0.02,
+  "tie_word_embeddings": false,
+  "patch_in_forward": true,
+  "patch_size": 4,
+  "patching_mode": "entropy",
+  "patching_threshold": 1.335442066192627,
+  "patching_batch_size": 1,
+  "max_patch_length": null,
+  "patching_device": "cuda",
+  "realtime_patching": true,
+  "patching_threshold_add": null,
+  "monotonicity": false,
+  "cross_attn_k": 2,
+  "encoder_hash_byte_group_size": [3, 4, 5, 6, 7, 8],
+  "encoder_hash_byte_group_vocab": 1000,
+  "encoder_hash_byte_group_nb_functions": 1,
+  "patcher_config": {
+    "model_type": "blt_patcher",
+    "vocab_size": 260,
+    "hidden_size": 512,
+    "num_hidden_layers": 7,
+    "num_attention_heads": 8,
+    "num_key_value_heads": 8,
+    "max_position_embeddings": 8192,
+    "rms_norm_eps": 1e-5,
+    "dropout": 0.0,
+    "intermediate_size": 1365,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rope_parameters": {"rope_type": "default",
+                        "rope_theta": 500000
+                        }
+  },
+  "encoder_config": {
+    "model_type": "blt_local_encoder",
+    "vocab_size": 260,
+    "hidden_size": 512,
+    "hidden_size_global": 1024,
+    "num_hidden_layers": 1,
+    "num_attention_heads": 8,
+    "num_key_value_heads": 8,
+    "head_dim": 64,
+    "intermediate_size": 1365,
+    "rms_norm_eps": 1e-5,
+    "dropout": 0.0,
+    "max_position_embeddings": 24576,
+    "cross_attn_all_layers": false,
+    "cross_attn_k": 2,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rope_parameters": {"rope_type": "default",
+                        "rope_theta": 500000
+                        }
+  },
+  "decoder_config": {
+    "model_type": "blt_local_decoder",
+    "vocab_size": 260,
+    "hidden_size": 512,
+    "hidden_size_global": 1024,
+    "num_hidden_layers": 9,
+    "num_attention_heads": 8,
+    "num_key_value_heads": 8,
+    "head_dim": 64,
+    "intermediate_size": 1365,
+    "rms_norm_eps": 1e-5,
+    "dropout": 0.0,
+    "max_position_embeddings": 24576,
+    "cross_attn_all_layers": true,
+    "cross_attn_k": 2,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rope_parameters": {"rope_type": "default",
+                        "rope_theta": 500000
+                        }
+  },
+  "global_config": {
+    "model_type": "blt_global_transformer",
+    "hidden_size": 1024,
+    "num_hidden_layers": 25,
+    "num_attention_heads": 8,
+    "num_key_value_heads": 8,
+    "head_dim": 128,
+    "intermediate_size": 2731,
+    "rms_norm_eps": 1e-5,
+    "dropout": 0.0,
+    "max_position_embeddings": 4096,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rope_parameters": {"rope_type": "default",
+                        "rope_theta": 500000
+                        },
+    "encoder_cross_output_size": null
+  }
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/blt_transformer_1_5B.json

@@ -0,0 +1,99 @@
+{
+  "model_type": "blt",
+  "vocab_size": 260,
+  "max_position_embeddings": 4096,
+  "initializer_range": 0.02,
+  "tie_word_embeddings": false,
+  "patch_in_forward": true,
+  "patch_size": 4,
+  "patching_mode": "entropy",
+  "patching_threshold": 1.335442066192627,
+  "patching_batch_size": 1,
+  "max_patch_length": null,
+  "patching_device": "cuda",
+  "realtime_patching": true,
+  "patching_threshold_add": null,
+  "monotonicity": false,
+  "cross_attn_k": 2,
+  "encoder_hash_byte_group_size": [3, 4, 5, 6, 7, 8],
+  "encoder_hash_byte_group_vocab": 500,
+  "encoder_hash_byte_group_nb_functions": 1,
+  "patcher_config": {
+    "model_type": "blt_patcher",
+    "vocab_size": 260,
+    "hidden_size": 768,
+    "num_hidden_layers": 7,
+    "num_attention_heads": 12,
+    "num_key_value_heads": 12,
+    "max_position_embeddings": 8192,
+    "rms_norm_eps": 1e-5,
+    "dropout": 0.0,
+    "intermediate_size": 2048,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rope_parameters": {"rope_type": "default",
+                        "rope_theta": 500000
+                        }
+  },
+  "encoder_config": {
+    "model_type": "blt_local_encoder",
+    "vocab_size": 260,
+    "hidden_size": 1024,
+    "hidden_size_global": 2048,
+    "num_hidden_layers": 1,
+    "num_attention_heads": 16,
+    "num_key_value_heads": 16,
+    "head_dim": 64,
+    "intermediate_size": 2816,
+    "rms_norm_eps": 1e-5,
+    "dropout": 0.0,
+    "max_position_embeddings": 24576,
+    "cross_attn_all_layers": false,
+    "cross_attn_k": 2,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rope_parameters": {"rope_type": "default",
+                        "rope_theta": 500000
+                        }
+  },
+
+  "decoder_config": {
+    "model_type": "blt_local_decoder",
+    "vocab_size": 260,
+    "hidden_size": 1024,
+    "hidden_size_global": 2048,
+    "num_hidden_layers": 9,
+    "num_attention_heads": 16,
+    "num_key_value_heads": 16,
+    "head_dim": 64,
+    "intermediate_size": 2816,
+    "rms_norm_eps": 1e-5,
+    "dropout": 0.0,
+    "max_position_embeddings": 24576,
+    "cross_attn_all_layers": true,
+    "cross_attn_k": 2,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rope_parameters": {"rope_type": "default",
+                        "rope_theta": 500000
+                        }
+  },
+  "global_config": {
+    "model_type": "blt_global_transformer",
+    "hidden_size": 2048,
+    "num_hidden_layers": 25,
+    "num_attention_heads": 16,
+    "num_key_value_heads": 16,
+    "head_dim": 128,
+    "intermediate_size": 5632,
+    "rms_norm_eps": 1e-5,
+    "dropout": 0.0,
+    "max_position_embeddings": 4096,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rope_parameters": {"rope_type": "default",
+                        "rope_theta": 500000
+                        },
+    "encoder_cross_output_size": null
+  }
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/blt_transformer_380M.json

@@ -0,0 +1,98 @@
+{
+  "model_type": "blt",
+  "vocab_size": 260,
+  "max_position_embeddings": 4096,
+  "initializer_range": 0.02,
+  "tie_word_embeddings": false,
+  "patch_in_forward": true,
+  "patch_size": 4,
+  "patching_mode": "entropy",
+  "patching_threshold": 1.335442066192627,
+  "patching_batch_size": 1,
+  "max_patch_length": null,
+  "patching_device": "cuda",
+  "realtime_patching": true,
+  "patching_threshold_add": null,
+  "monotonicity": false,
+  "cross_attn_k": 2,
+  "encoder_hash_byte_group_size": [3, 4, 5, 6, 7, 8],
+  "encoder_hash_byte_group_vocab": 500,
+  "encoder_hash_byte_group_nb_functions": 1,
+  "patcher_config": {
+    "model_type": "blt_patcher",
+    "vocab_size": 260,
+    "hidden_size": 512,
+    "num_hidden_layers": 7,
+    "num_attention_heads": 8,
+    "num_key_value_heads": 8,
+    "max_position_embeddings": 8192,
+    "rms_norm_eps": 1e-5,
+    "dropout": 0.0,
+    "intermediate_size": 1365,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rope_parameters": {"rope_type": "default",
+                        "rope_theta": 500000
+                        }
+  },
+  "encoder_config": {
+    "model_type": "blt_local_encoder",
+    "vocab_size": 260,
+    "hidden_size": 512,
+    "hidden_size_global": 1024,
+    "num_hidden_layers": 1,
+    "num_attention_heads": 8,
+    "num_key_value_heads": 8,
+    "head_dim": 64,
+    "intermediate_size": 1365,
+    "rms_norm_eps": 1e-5,
+    "dropout": 0.0,
+    "max_position_embeddings": 24576,
+    "cross_attn_all_layers": false,
+    "cross_attn_k": 2,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rope_parameters": {"rope_type": "default",
+                        "rope_theta": 500000
+                        }
+  },
+  "decoder_config": {
+    "model_type": "blt_local_decoder",
+    "vocab_size": 260,
+    "hidden_size": 512,
+    "hidden_size_global": 1024,
+    "num_hidden_layers": 9,
+    "num_attention_heads": 8,
+    "num_key_value_heads": 8,
+    "head_dim": 64,
+    "intermediate_size": 1365,
+    "rms_norm_eps": 1e-5,
+    "dropout": 0.0,
+    "max_position_embeddings": 24576,
+    "cross_attn_all_layers": true,
+    "cross_attn_k": 2,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rope_parameters": {"rope_type": "default",
+                        "rope_theta": 500000
+                        }
+  },
+  "global_config": {
+    "model_type": "blt_global_transformer",
+    "hidden_size": 1024,
+    "num_hidden_layers": 25,
+    "num_attention_heads": 8,
+    "num_key_value_heads": 8,
+    "head_dim": 128,
+    "intermediate_size": 2731,
+    "rms_norm_eps": 1e-5,
+    "dropout": 0.0,
+    "max_position_embeddings": 4096,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rope_parameters": {"rope_type": "default",
+                        "rope_theta": 500000
+                        },
+    "encoder_cross_output_size": null
+  }
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/delta_net_1B.json

@@ -0,0 +1,29 @@
+{
+    "attn": null,
+    "attn_mode": "chunk",
+    "bos_token_id": 1,
+    "conv_size": 4,
+    "eos_token_id": 2,
+    "expand_k": 1,
+    "expand_v": 1,
+    "fuse_cross_entropy": true,
+    "fuse_norm": true,
+    "hidden_act": "swish",
+    "hidden_ratio": 4,
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": null,
+    "model_type": "delta_net",
+    "norm_eps": 1e-06,
+    "num_heads": 16,
+    "num_hidden_layers": 24,
+    "pad_token_id": 2,
+    "qk_activation": "silu",
+    "qk_norm": "l2",
+    "tie_word_embeddings": false,
+    "use_beta": true,
+    "use_cache": true,
+    "use_gate": false,
+    "use_output_norm": true,
+    "use_short_conv": true
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/delta_net_340M.json

@@ -0,0 +1,26 @@
+{
+    "attn_mode": "chunk",
+    "bos_token_id": 1,
+    "conv_size": 4,
+    "eos_token_id": 2,
+    "expand_k": 1,
+    "expand_v": 1,
+    "fuse_cross_entropy": true,
+    "hidden_act": "swish",
+    "hidden_ratio": 4,
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "intermediate_size": null,
+    "model_type": "delta_net",
+    "norm_eps": 1e-06,
+    "num_heads": 8,
+    "num_hidden_layers": 24,
+    "qk_activation": "silu",
+    "qk_norm": "l2",
+    "tie_word_embeddings": false,
+    "use_beta": true,
+    "use_cache": true,
+    "use_gate": false,
+    "use_output_norm": true,
+    "use_short_conv": true
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/gated_deltanet_1B.json

@@ -0,0 +1,22 @@
+{
+    "attn_mode": "chunk",
+    "bos_token_id": 1,
+    "conv_size": 4,
+    "eos_token_id": 2,
+    "expand_v": 2,
+    "fuse_cross_entropy": true,
+    "head_dim": 256,
+    "hidden_act": "swish",
+    "hidden_ratio": 4,
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": null,
+    "model_type": "gated_deltanet",
+    "norm_eps": 1e-06,
+    "num_heads": 6,
+    "num_hidden_layers": 21,
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "use_gate": true,
+    "use_short_conv": true
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/gated_deltanet_340M.json

@@ -0,0 +1,22 @@
+{
+    "attn_mode": "chunk",
+    "bos_token_id": 1,
+    "conv_size": 4,
+    "eos_token_id": 2,
+    "expand_v": 2,
+    "fuse_cross_entropy": true,
+    "head_dim": 256,
+    "hidden_act": "swish",
+    "hidden_ratio": 4,
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "intermediate_size": null,
+    "model_type": "gated_deltanet",
+    "norm_eps": 1e-06,
+    "num_heads": 6,
+    "num_hidden_layers": 21,
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "use_gate": true,
+    "use_short_conv": true
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/gated_deltanet_h_340M.json

@@ -0,0 +1,28 @@
+{
+    "model_type": "gated_deltanet",
+    "attn_mode": "chunk",
+    "hidden_size": 1024,
+    "num_hidden_layers": 21,
+    "head_dim": 256,
+    "num_heads": 6,
+    "expand_v": 2,
+    "hidden_ratio": 4,
+    "use_gate": true,
+    "use_short_conv": true,
+    "conv_size": 4,
+    "vocab_size": 32000,
+    "hidden_act": "swish",
+    "norm_eps": 1e-06,
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "fuse_cross_entropy": true,
+    "initializer_range": 0.02,
+    "attn": {
+        "layers": [3, 7, 11, 15, 19],
+        "num_heads": 8,
+        "num_kv_heads": 1,
+        "window_size": 2048,
+        "rope_theta": 100000.0,
+        "qkv_bias": false
+    }
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/gla_340M.json

@@ -0,0 +1,24 @@
+{
+  "attn_mode": "chunk",
+  "bos_token_id": 1,
+  "clamp_min": null,
+  "eos_token_id": 2,
+  "expand_k": 0.5,
+  "expand_v": 1,
+  "fuse_cross_entropy": true,
+  "fuse_norm": true,
+  "hidden_act": "swish",
+  "hidden_ratio": 4,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": null,
+  "model_type": "gla",
+  "num_heads": 4,
+  "num_hidden_layers": 24,
+  "norm_eps": 1e-06,
+  "tie_word_embeddings": false,
+  "use_cache": true,
+  "use_gk": true,
+  "use_gv": false,
+  "vocab_size": 32000
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/gla_7B.json

@@ -0,0 +1,25 @@
+{
+    "attn": null,
+    "attn_mode": "chunk",
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "expand_k": 0.5,
+    "expand_v": 1,
+    "fuse_cross_entropy": true,
+    "fuse_norm": true,
+    "hidden_act": "swish",
+    "hidden_ratio": 4,
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 11008,
+    "model_type": "gla",
+    "norm_eps": 1e-06,
+    "num_heads": 16,
+    "num_hidden_layers": 32,
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "use_gk": true,
+    "use_gv": false,
+    "use_output_gate": true,
+    "use_short_conv": false
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/gsa_340M.json

@@ -0,0 +1,29 @@
+{
+    "bos_token_id": 1,
+    "conv_size": 4,
+    "eos_token_id": 2,
+    "expand_k": 1,
+    "expand_v": 1,
+    "elementwise_affine": false,
+    "feature_map": "swish",
+    "fuse_cross_entropy": true,
+    "fuse_norm": true,
+    "gate_logit_normalizer": 4,
+    "hidden_act": "swish",
+    "hidden_ratio": 4,
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "intermediate_size": null,
+    "model_type": "gsa",
+    "num_heads": 4,
+    "num_hidden_layers": 24,
+    "num_slots": 64,
+    "norm_eps": 1e-06,
+    "share_conv_kernel": true,
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "use_norm": true,
+    "use_output_gate": true,
+    "use_rope": false,
+    "use_short_conv": false
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/mergenet_340M.json

@@ -0,0 +1,34 @@
+{
+    "model_type": "mergenet",
+    "vocab_size": 260,
+    "hidden_size": 1024,
+    "num_local_layers": 6,
+    "local_depth": 4,
+    "num_latent_layers": 12,
+    "num_heads": 16,
+    "num_kv_heads": 16,
+    "intermediate_size": 4096,
+    "hidden_act": "swish",
+    "max_position_embeddings": 8192,
+    "lambda_local": 4.0,
+    "dtem_window_size": 8,
+    "dtem_t": 1,
+    "dtem_feat_dim": null,
+    "use_softkmax": false,
+    "grid_bias_gamma": 1.0,
+    "W_infer": null,
+    "qkv_bias": true,
+    "qk_norm": false,
+    "rope_theta": 10000.0,
+    "norm_eps": 1e-6,
+    "initializer_range": 0.02,
+    "use_cache": true,
+    "pad_token_id": 0,
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "tie_word_embeddings": false,
+    "phase": "phase2",
+    "drop_rate": 0.0,
+    "attn_drop_rate": 0.0,
+    "drop_path_rate": 0.1
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/mergenet_64M.json

@@ -0,0 +1,34 @@
+{
+    "model_type": "mergenet",
+    "vocab_size": 32000,
+    "hidden_size": 512,
+    "num_local_layers": 4,
+    "local_depth": 4,
+    "num_latent_layers": 8,
+    "num_heads": 8,
+    "num_kv_heads": 8,
+    "intermediate_size": 2048,
+    "hidden_act": "swish",
+    "max_position_embeddings": 4096,
+    "lambda_local": 4.0,
+    "dtem_window_size": 8,
+    "dtem_t": 1,
+    "dtem_feat_dim": null,
+    "use_softkmax": false,
+    "grid_bias_gamma": 1.0,
+    "W_infer": null,
+    "qkv_bias": true,
+    "qk_norm": false,
+    "rope_theta": 10000.0,
+    "norm_eps": 1e-6,
+    "initializer_range": 0.02,
+    "use_cache": true,
+    "pad_token_id": 0,
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "tie_word_embeddings": false,
+    "phase": "phase2",
+    "drop_rate": 0.0,
+    "attn_drop_rate": 0.0,
+    "drop_path_rate": 0.1
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/qwen3_next_1B.json

@@ -0,0 +1,44 @@
+{
+    "model_type": "qwen3_next",
+    "vocab_size": 151936,
+    "hidden_size": 2048,
+    "intermediate_size": 5632,
+    "num_hidden_layers": 48,
+    "num_attention_heads": 16,
+    "num_key_value_heads": 2,
+    "head_dim": 256,
+    "hidden_act": "silu",
+    "max_position_embeddings": 32768,
+    "initializer_range": 0.02,
+    "rms_norm_eps": 1e-6,
+    "use_cache": true,
+    "tie_word_embeddings": false,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "rope_parameters": {
+      "rope_type": "default",
+      "factor": 1.0
+    },
+    "partial_rotary_factor": 0.25,
+    "layer_types": [
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention"
+    ],
+    "linear_conv_kernel_dim": 4,
+    "linear_key_head_dim": 128,
+    "linear_value_head_dim": 128,
+    "linear_num_key_heads": 16,
+    "linear_num_value_heads": 32,
+    "decoder_sparse_step": 1,
+    "moe_intermediate_size": 512,
+    "shared_expert_intermediate_size": 512,
+    "num_experts_per_tok": 10,
+    "num_experts": 512,
+    "norm_topk_prob": true,
+    "output_router_logits": false,
+    "router_aux_loss_coef": 0.001,
+    "mlp_only_layers": []
+  }
+  

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/qwen3_next_350M.json

@@ -0,0 +1,44 @@
+{
+    "model_type": "qwen3_next",
+    "vocab_size": 32000,
+    "hidden_size": 2048,
+    "intermediate_size": 5632,
+    "num_hidden_layers": 26,
+    "num_attention_heads": 16,
+    "num_key_value_heads": 2,
+    "head_dim": 256,
+    "hidden_act": "silu",
+    "max_position_embeddings": 32768,
+    "initializer_range": 0.02,
+    "rms_norm_eps": 1e-6,
+    "use_cache": true,
+    "tie_word_embeddings": false,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "rope_parameters": {
+      "rope_type": "default",
+      "factor": 1.0
+    },
+    "partial_rotary_factor": 0.25,
+    "layer_types": [
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention"
+    ],
+    "linear_conv_kernel_dim": 4,
+    "linear_key_head_dim": 128,
+    "linear_value_head_dim": 128,
+    "linear_num_key_heads": 16,
+    "linear_num_value_heads": 32,
+    "decoder_sparse_step": 1,
+    "moe_intermediate_size": 512,
+    "shared_expert_intermediate_size": 512,
+    "num_experts_per_tok": 10,
+    "num_experts": 512,
+    "norm_topk_prob": true,
+    "output_router_logits": false,
+    "router_aux_loss_coef": 0.001,
+    "mlp_only_layers": []
+  }
+  

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/transformer_1B.json

@@ -0,0 +1,22 @@
+{
+    "bos_token_id": 1,
+    "elementwise_affine": true,
+    "eos_token_id": 2,
+    "fuse_cross_entropy": true,
+    "fuse_norm": true,
+    "fuse_swiglu": true,
+    "hidden_act": "swish",
+    "hidden_ratio": 4,
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": null,
+    "max_position_embeddings": 8192,
+    "model_type": "transformer",
+    "norm_eps": 1e-06,
+    "num_heads": 32,
+    "num_hidden_layers": 24,
+    "num_kv_heads": null,
+    "pad_token_id": 2,
+    "rope_theta": 10000.0,
+    "tie_word_embeddings": false
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/transformer_340M.json

@@ -0,0 +1,18 @@
+{
+    "attention_bias": false,
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "fuse_cross_entropy": true,
+    "fuse_norm": true,
+    "hidden_act": "swish",
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "max_position_embeddings": 8192,
+    "model_type": "transformer",
+    "num_heads": 16,
+    "num_hidden_layers": 24,
+    "norm_eps": 1e-06,
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "vocab_size": 32000
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/configs/transformer_7B.json

@@ -0,0 +1,21 @@
+{
+    "attention_bias": false,
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "fuse_cross_entropy": true,
+    "fuse_norm": true,
+    "hidden_act": "swish",
+    "hidden_ratio": 4,
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 14336,
+    "model_type": "transformer",
+    "norm_eps": 1e-06,
+    "num_heads": 32,
+    "num_hidden_layers": 32,
+    "num_kv_heads": 8,
+    "rope_theta": 10000.0,
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "window_size": null
+}

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/flame/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/flame/c4_test.py

@@ -0,0 +1,603 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+import os
+import time
+import json
+
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+
+from tqdm import tqdm
+from loguru import logger
+
+import transformers
+
+transformers.logging.set_verbosity_error()
+
+import wandb
+
+from utils.argparse import parse_args
+from utils.setup import getting_svd_cnt, set_seed, setup_model, saving_model_weight, load_model_weight
+from utils.optimizer_factory import setup_optimization
+from utils.eval import evaluate_model
+from utils.dataloader import setup_dataset
+from utils.modeling_llama import LlamaForCausalLM
+from utils.fake_quantization import QLinear
+from utils.quantization import QScaleLinear
+
+
+def main(args):
+    import torch
+    ############ Setup random seed ############
+    set_seed(args)
+
+    ############ Setup DDP environment ############
+    assert "LOCAL_RANK" in os.environ, "torchrun should set LOCAL_RANK"
+    global_rank = int(os.environ["RANK"])
+    local_rank = int(os.environ["LOCAL_RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+    torch.cuda.set_device(local_rank)
+
+    logger.info(f"Global rank {global_rank}, local rank {local_rank}, device: {torch.cuda.current_device()}")
+    dist.init_process_group(backend="nccl", rank=global_rank, world_size=world_size)
+
+    logger.info("Process group initialized")
+    device = f"cuda:{local_rank}"
+
+    if global_rank != 0:
+        logger.remove()  # turn off logger
+
+    logger.info(f"Using dist with rank {global_rank} (only rank 0 will log)")
+    logger.info("*" * 40)
+    logger.info(f"Starting training with the arguments")
+    for k, v in vars(args).items():
+        logger.info(f"{k:30} {v}")
+    logger.info("*" * 40)
+
+    ############ Initialize wandb without config (it is passed later) ############
+    if (not args.unset_wandb) and global_rank == 0:
+        if args.entity is None:
+            os.environ['WANDB_MODE'] = 'offline'
+        # Set wandb directory for offline mode
+        wandb_dir = getattr(args, 'wandb_dir', None) if getattr(args, 'wandb_dir', None) is not None else args.save_dir
+        if getattr(args, 'wandb_dir', None) is not None:
+            logger.info(f"Wandb directory set to: {wandb_dir}")
+        wandb.init(project=args.project, name=args.name, entity=args.entity, dir=wandb_dir)
+
+    ############ Setup training data ############
+    if args.total_batch_size is not None:
+        if args.gradient_accumulation is None:
+            assert args.total_batch_size % world_size == 0, "total_batch_size must be divisible by world_size"
+            args.gradient_accumulation = args.total_batch_size // (args.batch_size * world_size)
+            assert args.gradient_accumulation > 0, "gradient_accumulation must be greater than 0"
+
+    assert (
+        args.gradient_accumulation * args.batch_size * world_size == args.total_batch_size
+    ), "gradient_accumulation * batch_size * world_size must be equal to total_batch_size"
+
+    dataloader, tokenizer = setup_dataset(args, global_rank, world_size)
+
+    ############ Initialize model ############
+    model_config, model = setup_model(args)
+    # Ensure model has generation_config (fix for transformers version compatibility)
+    if model.generation_config is None:
+        from transformers import GenerationConfig
+        model.generation_config = GenerationConfig()
+    model.generation_config.pad_token_id = tokenizer.pad_token_id
+
+    ############ Resuming from checkpoints ############
+    global_step = 0
+    update_step = 0
+    beginning_step = 0
+    tokens_seen = 0
+    tokens_seen_before = 0
+
+    # identifying checkpointing
+    if args.continue_from is not None and os.path.exists(args.continue_from):
+        # searching the latest checkpoints
+        checkpoint_path_list = os.listdir(args.continue_from)
+        checkpoint_path_list = [int(x.split("_")[-1]) for x in checkpoint_path_list if x.startswith("model_")]
+        if len(checkpoint_path_list) > 0:
+            logger.info("Find Checkpoints", checkpoint_path_list)
+            beginning_step = max(checkpoint_path_list)
+            if args.resume_step is not None:
+                beginning_step = args.resume_step
+            args.continue_from = os.path.join(args.continue_from, f"model_{beginning_step}")
+            logger.info("Continue from", args.continue_from)
+        else:
+            logger.warning(f"Did not find any checkpoints in {args.continue_from}")
+            args.continue_from = None
+
+    # resuming from checkpointing
+    if args.continue_from is not None:
+        logger.info("*" * 40)
+        logger.info(f"Loading model from {args.continue_from}")
+        checkpoint_path = os.path.join(args.continue_from, "pytorch_model.bin")
+        if os.path.exists(checkpoint_path):
+            load_model_weight(model, checkpoint_path, args)
+            logger.info(f"Model successfully loaded (strict=False policy)")
+        else:
+            # Try safetensors format
+            checkpoint_path = os.path.join(args.continue_from, "model.safetensors")
+            if os.path.exists(checkpoint_path):
+                from safetensors import safe_open
+                tensors = {}
+                with safe_open(checkpoint_path, framework="pt", device=0) as f:
+                    for k in f.keys():
+                        tensors[k] = f.get_tensor(k)
+                        print(k, tensors[k].shape)
+                ret = model.load_state_dict(tensors, strict=False)
+                logger.info(f"Model successfully loaded from safetensors (strict=False policy)", ret)
+            else:
+                logger.warning(f"No model checkpoint found in {args.continue_from}")
+
+        if os.path.exists(os.path.join(args.continue_from, "training_state.json")):
+            logger.info(
+                f"Loading training state like global_step, update_step, and tokens_seen from {args.continue_from}"
+            )
+            with open(os.path.join(args.continue_from, "training_state.json")) as f:
+                _old_state = json.load(f)
+            global_step = _old_state["global_step"]
+            update_step = _old_state["update_step"]
+            tokens_seen = _old_state["tokens_seen"]
+            tokens_seen_before = _old_state["tokens_seen_before"]
+            logger.info(f"global_step       : {global_step}")
+            logger.info(f"update_step       : {update_step}")
+            logger.info(f"tokens_seen       : {tokens_seen}")
+            logger.info(f"tokens_seen_before: {tokens_seen_before}")
+            logger.info(f"Will train for {args.num_training_steps - update_step} update steps")
+        else:
+            logger.warning(f"Did not find training state in {args.continue_from}, global step will start from zero")
+        logger.info("*" * 40)
+
+    ############ Setup model ############
+    if args.dtype in ["bf16", "bfloat16"]:
+        model = model.to(dtype=torch.bfloat16)
+    model = model.to(device=device)
+
+    for _, module in model.named_modules():
+        if isinstance(module, QScaleLinear):
+            weight_device = module.weight.device
+            module.weight.scales = module.weight.scales.to(device=weight_device)
+            module.weight.zeros = module.weight.zeros.to(device=weight_device)
+
+    n_total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = [p for p in model.parameters() if p.requires_grad]
+    trainable_params_int8 = [p for p in model.parameters() if hasattr(p, "group_size")]
+
+    ############ Initialize wandb ############
+    run_config = dict(vars(args))
+    run_config.update(
+        {
+            "max_lr": run_config.pop("lr"),  # rename lr to max_lr to avoid conflicts with scheduler
+            "total_params_M": n_total_params / 1_000_000,
+            "dataset": "c4",
+            "model": model_config.to_dict(),
+            "world_size": world_size,
+            "device": str(device),
+        }
+    )
+
+    if global_rank == 0:
+        if not args.unset_wandb:
+            wandb.config.update(run_config, allow_val_change=True)
+            wandb.save(os.path.abspath(__file__), policy="now")  # save current script
+        # fix tqdm visual length to 80 so that the progress bar
+        # doesn't jump around when changing from external display to laptop
+        pbar = tqdm(total=args.num_training_steps - update_step, desc="Update steps", ncols=80)
+
+    ############ Initialize optimization ############
+    if "galore" in args.optimizer.lower():
+        # make parameters with "rank" to a single group, if param_name has "mlp" or "attn"
+        lowrank_params = []
+        target_modules_list = ["attn", "mlp"]
+        for module_name, module in model.named_modules():
+            if not (isinstance(module, nn.Linear) or isinstance(module, QScaleLinear) or isinstance(module, QLinear)):
+                continue
+            if not any(target_key in module_name for target_key in target_modules_list):
+                continue
+            logger.info(f"Adding {module_name} to GaLore parameters")
+            lowrank_params.append(module.weight)
+
+        id_lowrank_params = [id(p) for p in lowrank_params]
+        # make parameters without "rank" to another group
+        regular_params = [p for p in model.parameters() if id(p) not in id_lowrank_params]
+        # then call low rank optimizer
+        param_groups = [
+            {"params": regular_params},
+            {
+                "params": lowrank_params,
+                "rank": args.rank,
+                "update_proj_gap": args.update_proj_gap,
+                "scale": args.galore_scale,
+                "proj_type": args.proj_type,
+                "quant": args.proj_quant,
+                "quant_n_bit": args.proj_bits,
+                "quant_group_size": args.proj_group_size,
+                "cos_threshold": args.cos_threshold,
+                "gamma_proj": args.gamma_proj,
+                "queue_size": args.queue_size,
+            },
+        ]
+    elif "apollo" in args.optimizer.lower():
+        # make parameters with "rank" to a single group, if param_name has "mlp" or "attn"
+        lowrank_params = []
+        target_modules_list = ["attn", "mlp"]
+        for module_name, module in model.named_modules():
+            if not (isinstance(module, nn.Linear) or isinstance(module, QScaleLinear) or isinstance(module, QLinear)):
+                continue
+            if not any(target_key in module_name for target_key in target_modules_list):
+                continue
+            logger.info(f"Adding {module_name} to APOLLO parameters")
+            lowrank_params.append(module.weight)
+
+        id_lowrank_params = [id(p) for p in lowrank_params]
+        # make parameters without "rank" to another group
+        regular_params = [p for p in model.parameters() if id(p) not in id_lowrank_params]
+        # then call low rank optimizer
+        param_groups = [
+            {"params": regular_params},
+            {
+                "params": lowrank_params,
+                "rank": args.rank,
+                "update_proj_gap": args.update_proj_gap,
+                "scale": args.apollo_scale,
+                "proj_type": args.proj_type,
+                "proj": args.proj,
+                "scale_type": args.scale_type,
+            },
+        ]
+    elif "conda" in args.optimizer.lower():
+        # make parameters with "rank" to a single group, if param_name has "mlp" or "attn"
+        lowrank_params = []
+        target_modules_list = ["attn", "mlp"]
+        for module_name, module in model.named_modules():
+            if not (isinstance(module, nn.Linear) or isinstance(module, QScaleLinear) or isinstance(module, QLinear)):
+                continue
+            if not any(target_key in module_name for target_key in target_modules_list):
+                continue
+            logger.info(f"Adding {module_name} to conda parameters")
+            lowrank_params.append(module.weight)
+
+        id_lowrank_params = [id(p) for p in lowrank_params]
+        # make parameters without "rank" to another group
+        regular_params = [p for p in model.parameters() if id(p) not in id_lowrank_params]
+        # then call low rank optimizer
+        param_groups = [
+            {"params": regular_params},
+            {
+                "params": lowrank_params,
+                "rank": args.rank,
+                "update_proj_gap": args.update_proj_gap,
+                "scale": args.apollo_scale,
+                "proj_type": args.proj_type,
+                "proj": args.proj,
+                "scale_type": args.scale_type,
+            },
+        ]
+    else:
+        param_groups = None
+        id_lowrank_params = None
+
+    # print params and trainable params
+    logger.info(f"\n{model}\n")
+    logger.info(f"Total params: {sum(p.numel() for p in model.parameters()) / 1_000_000:.2f}M")
+
+    if args.simulation:
+        num_train_params = sum(p.numel() for p in trainable_params)
+    else:
+        num_train_params = sum(p.numel() for p in trainable_params) + sum(p.numel() for p in trainable_params_int8)
+
+    logger.info(f"Trainable params: {num_train_params / 1_000_000:.2f}M")
+    if "q_galore" in args.optimizer.lower():
+        logger.info(
+            f"Trainable params with Q-GaLore enabled: {sum(p.numel() for p in trainable_params_int8) / 1_000_000:.2f}M"
+        )
+    elif "galore" in args.optimizer.lower():
+        logger.info(f"Total params with GaLore enabled: {sum(p.numel() for p in lowrank_params) / 1_000_000:.2f}M")
+    elif "q_apollo" in args.optimizer.lower():
+        logger.info(
+            f"Trainable params with Q-APOLLO enabled: {sum(p.numel() for p in trainable_params_int8) / 1_000_000:.2f}M"
+        )
+    elif "apollo" in args.optimizer.lower():
+        logger.info(f"Total params with APOLLO enabled: {sum(p.numel() for p in lowrank_params) / 1_000_000:.2f}M")
+
+    logger.info(f"Saving model to {args.save_dir} every {args.save_every} update steps")
+
+    model, optimizer, scheduler, layer_wise_flag = setup_optimization(
+        args, model, trainable_params, param_groups, id_lowrank_params, model_config
+    )
+
+    if layer_wise_flag:
+        # will pass optimizer_dict and scheduler_dict out instead of optimizer and scheduler
+        optimizer_dict = optimizer
+        scheduler_dict = scheduler
+
+    # Bug-3 fix: wrap with DDP *before* torch.compile per PyTorch recommendation.
+    # This ensures gradient reduction hooks are correctly installed on the DDP module,
+    # and the compiled graph captures the full DDP+model forward pass.
+    # (Issue-5: optimizer.load_state_dict is called after both DDP and compile below.)
+    if not args.single_gpu:
+        model: LlamaForCausalLM = torch.nn.parallel.DistributedDataParallel(
+            model,
+            device_ids=[local_rank],
+            output_device=local_rank,
+            broadcast_buffers=False,
+        )
+
+    # compile the model (after DDP so the compiled graph includes DDP reduction)
+    if args.compile:
+        print("Compiling the model... (takes a ~minute)")
+        unoptimized_model = model
+
+        # Configure TorchDynamo to suppress errors and fall back to eager mode
+        import torch._dynamo
+        torch._dynamo.config.suppress_errors = args.dynamo_suppress_errors
+        torch._dynamo.config.verbose = False
+        # Set cache size limit to prevent memory issues during long training
+        torch._dynamo.config.cache_size_limit = args.dynamo_cache_limit
+
+        model = torch.compile(model) # requires PyTorch 2.0
+
+    # resume optimizer
+    if args.restore_optimizer and args.continue_from is not None:
+        logger.info("Restoring optimizer and scheduler from the checkpoint")
+        _optimizer_dir = args.continue_from
+        optimizer_checkpoint = torch.load(os.path.join(_optimizer_dir, "optimizer.pt"), map_location="cpu")
+        optimizer.load_state_dict(optimizer_checkpoint["optimizer"])
+        scheduler.load_state_dict(optimizer_checkpoint["scheduler"])
+        update_step = optimizer_checkpoint["update_step"]
+        beginning_step = update_step
+        global_step = optimizer_checkpoint["global_step"]
+        logger.info(f"Optimizer and scheduler restored from {_optimizer_dir}")
+
+    # ##############################
+    # TRAINING LOOP
+    # we use iterable dataset, so we may never go through all the data
+    # ##############################
+    # global steps and others are defined above
+    pad_idx = tokenizer.pad_token_id
+    update_time = time.time()
+    local_step = 0  # when continue_from is used, local_step != global_step
+    total_svd_count = 0
+
+    dataloader_iter = iter(dataloader)
+
+    # Issue-4 fix: accumulate loss across micro-batches so logged loss is the true
+    # gradient-accumulation average, not just the last micro-batch.
+    accumulated_loss = 0.0
+
+    # Skip data if resuming from checkpoint
+    if update_step != 0:
+        skip_batches = args.gradient_accumulation * update_step
+        logger.info(f"Skipping {skip_batches} batches to resume from update step {update_step}")
+        skipped = 0
+        for _ in range(skip_batches):
+            # Issue-6 fix: handle StopIteration during skip so all ranks stay aligned
+            try:
+                next(dataloader_iter)
+            except StopIteration:
+                logger.warning(
+                    f"Dataset exhausted during skip at batch {skipped}/{skip_batches}; "
+                    f"restarting iterator to keep ranks aligned."
+                )
+                dataloader_iter = iter(dataloader)
+                next(dataloader_iter)
+            skipped += 1
+        logger.info(f"Skipped {skipped} batches successfully")
+
+    while update_step <= args.num_training_steps:
+        try:
+            batch = next(dataloader_iter)
+        except StopIteration:
+            logger.info(f"Dataset completed one epoch. Starting new epoch with reshuffled data.")
+            dataloader_iter = iter(dataloader)
+            batch = next(dataloader_iter)
+
+        global_step += 1
+        local_step += 1
+
+        if update_step >= args.num_training_steps:
+            logger.info(f"Reached max number of update steps ({args.num_training_steps}). Stopping training.")
+            logger.info(f"Rank {global_rank} stopping training.")
+            break
+
+        # forward & backward
+        batch = {k: v.to(device) for k, v in batch.items()}
+        labels = batch["input_ids"].clone()
+        labels[labels == pad_idx] = -100
+        tokens_seen += (batch["input_ids"] != pad_idx).sum().item() * world_size
+
+        loss = model(**batch, labels=labels).loss
+
+        scaled_loss = loss / args.gradient_accumulation
+        scaled_loss.backward()
+        accumulated_loss += loss.item()  # Issue-4: accumulate before the continue
+
+        if global_step % args.gradient_accumulation != 0:
+            continue
+
+        # The below code is only executed during the update step
+        # Issue-4: compute average loss over all micro-batches in this accumulation window
+        avg_loss = accumulated_loss / args.gradient_accumulation
+        accumulated_loss = 0.0  # reset for next accumulation window
+        # add grad clipping: TODO: add gradient clipping of int8 weight
+        if args.grad_clipping != 0.0:
+            torch.nn.utils.clip_grad_norm_(trainable_params, args.grad_clipping)
+        # Periodic memory cleanup to prevent symbolic tensor issues during long training
+        if global_step % args.memory_cleanup_frequency == 0:
+            torch.cuda.empty_cache()
+            # Clear TorchDynamo cache to prevent memory accumulation
+            if args.compile:
+                import torch._dynamo
+                torch._dynamo.reset()
+
+        if global_rank == 0:
+            pbar.update(1)
+        if not layer_wise_flag:  # layer-wise updation is done during backward; requires gradient_accumulation equals 1
+            optimizer.step()
+            scheduler.step()
+            optimizer.zero_grad()
+
+        update_step += 1
+        update_time = time.time() - update_time
+
+        # save checkpoint by save_every
+        if local_step > args.gradient_accumulation and update_step % args.save_every == 0 and global_rank == 0:
+            current_model_directory = f"{args.save_dir}/model_{update_step}"
+            logger.info(f"Saving model and optimizer to {current_model_directory}, update step {update_step}")
+            os.makedirs(args.save_dir, exist_ok=True)
+            # Bug-1 fix: unwrap DDP/compiled model for saving; works in both single-GPU and multi-GPU modes
+            unwrapped_model = model.module if hasattr(model, 'module') else model
+            unwrapped_model.save_pretrained(current_model_directory, max_shard_size="500GB", from_pt=True)
+            saving_model_weight(unwrapped_model, f"{current_model_directory}/pytorch_model.bin", args)
+
+            optimizer_checkpoint = {
+                "optimizer": optimizer.state_dict(),
+                "scheduler": scheduler.state_dict(),
+                "update_step": update_step,
+                "global_step": global_step,
+                "config": run_config,
+                "wandb": wandb.run.dir if not args.unset_wandb else None,
+                "dtype": args.dtype,
+            }
+            torch.save(optimizer_checkpoint, f"{current_model_directory}/optimizer.pt")
+
+            training_state_checkpoint = {
+                "global_step": global_step,
+                "update_step": update_step,
+                "tokens_seen": tokens_seen,
+                "tokens_seen_before": tokens_seen_before,
+                "update_time": update_time,
+            }
+            with open(f"{current_model_directory}/training_state.json", "w") as f:
+                json.dump(training_state_checkpoint, f, indent=4)
+
+            # save wandb related info
+            if not args.unset_wandb:
+                wandb_info = {
+                    "wandb_id": wandb.run.id,
+                }
+                with open(f"{args.save_dir}/wandb.json", "w") as f:
+                    json.dump(wandb_info, f, indent=4)
+
+        # evaluation
+        if update_step % args.eval_every == 0:
+            logger.info(f"Performing evaluation at step {update_step}")
+            total_loss, evaluated_on_tokens, perplexity = evaluate_model(
+                model, tokenizer, pad_idx, global_rank, world_size, device, args
+            )
+
+            if global_rank == 0:
+                if not args.unset_wandb:
+                    wandb.log(
+                        {
+                            "eval_loss": total_loss,
+                            "eval_perplexity": perplexity,
+                            "eval_tokens": evaluated_on_tokens,
+                        },
+                        step=update_step,
+                    )
+            logger.info(f"Eval loss at step {update_step}: {total_loss}, Eval perplexity: {perplexity}")
+
+        if not layer_wise_flag:
+            lr = optimizer.param_groups[0]["lr"]
+        else:
+            lr = list(optimizer_dict.values())[0].param_groups[0]["lr"]
+        tokens_in_update = tokens_seen - tokens_seen_before
+        tokens_seen_before = tokens_seen
+        batches_in_update = args.gradient_accumulation * world_size
+        if not layer_wise_flag:
+            total_svd_count = getting_svd_cnt(optimizer)
+        else:
+            total_svd_count = 0
+
+        if global_rank == 0:
+            if not args.unset_wandb:
+                wandb.log(
+                    {
+                        "loss": avg_loss,
+                        "lr": lr,
+                        "update_step": update_step,
+                        "tokens_seen": tokens_seen,
+                        "total_svd_count": total_svd_count,
+                        "throughput_tokens": tokens_in_update / update_time,
+                        "throughput_examples": args.total_batch_size / update_time,
+                        "throughput_batches": batches_in_update / update_time,
+                    },
+                    step=update_step,
+                )
+        update_time = time.time()
+
+    # ##############################
+    # END of training loop
+    # ##############################
+    logger.info("Training finished")
+    if global_rank == 0:
+        pbar.close()
+
+    current_model_directory = f"{args.save_dir}/model_{update_step}"
+    if global_rank == 0 and not os.path.exists(current_model_directory):
+        logger.info(f"Saving model and optimizer to {current_model_directory}, update step {update_step}")
+        os.makedirs(args.save_dir, exist_ok=True)
+        # Bug-1 fix: unwrap DDP/compiled model for saving; works in both single-GPU and multi-GPU modes
+        unwrapped_model = model.module if hasattr(model, 'module') else model
+        unwrapped_model.save_pretrained(current_model_directory, max_shard_size="500GB", from_pt=True)
+        saving_model_weight(unwrapped_model, f"{current_model_directory}/pytorch_model.bin", args)
+
+        optimizer_checkpoint = {
+            "optimizer": optimizer.state_dict(),
+            "scheduler": scheduler.state_dict(),
+            "update_step": update_step,
+            "global_step": global_step,
+            "config": run_config,
+            "wandb": wandb.run.dir if not args.unset_wandb else None,
+            "dtype": args.dtype,
+        }
+        torch.save(optimizer_checkpoint, f"{current_model_directory}/optimizer.pt")
+
+        training_state_checkpoint = {
+            "global_step": global_step,
+            "update_step": update_step,
+            "tokens_seen": tokens_seen,
+            "tokens_seen_before": tokens_seen_before,
+            "update_time": update_time,
+        }
+        with open(f"{current_model_directory}/training_state.json", "w") as f:
+            json.dump(training_state_checkpoint, f, indent=4)
+
+    # Final evaluation
+    logger.info("Running final evaluation")
+    model.eval()
+    del loss, optimizer, scheduler
+    import gc
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    total_loss, evaluated_on_tokens, perplexity = evaluate_model(model, tokenizer, pad_idx, global_rank, world_size, device, args)
+
+    if global_rank == 0:
+        if not args.unset_wandb:
+            wandb.log(
+                {
+                    "final_eval_loss": total_loss,
+                    "final_eval_perplexity": perplexity,
+                    "final_eval_tokens": evaluated_on_tokens,
+                },
+                step=update_step,
+            )
+        logger.info(f"Final eval loss: {total_loss}, Final eval perplexity: {perplexity}")
+
+    logger.info("Script finished successfully")
+    print(f"Rank {global_rank} finished successfully")
+
+
+if __name__ == "__main__":
+    print("Starting script")
+    args = parse_args(None)
+    main(args)

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/flame/components/checkpoint.py

@@ -0,0 +1,59 @@
+# 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.
+
+from dataclasses import dataclass, field
+from datetime import timedelta
+from io import BytesIO
+from typing import Any, Dict, List
+
+import torch
+from torch.distributed.checkpoint.stateful import Stateful
+
+
+@dataclass
+class TrainState(Stateful):
+    step: int = 0
+    skipped_step: int = 0
+    token: int = 0
+    elapsed: timedelta = timedelta(0)
+    global_avg_losses: List[float] = field(default_factory=list)
+    global_max_losses: List[float] = field(default_factory=list)
+    log_steps: List[int] = field(default_factory=list)
+
+    def state_dict(self) -> Dict[str, Any]:
+        # Only checkpoint global_avg_losses and global_max_losses per log frequency
+        # to avoid sync overhead in every iteration.
+        global_avg_losses_bytes = BytesIO()
+        torch.save(self.global_avg_losses, global_avg_losses_bytes)
+        global_max_losses_bytes = BytesIO()
+        torch.save(self.global_max_losses, global_max_losses_bytes)
+        log_steps_bytes = BytesIO()
+        torch.save(self.log_steps, log_steps_bytes)
+        return {
+            "step": torch.tensor(self.step, dtype=torch.int32),
+            "skipped_step": torch.tensor(self.skipped_step, dtype=torch.int32),
+            "token": torch.tensor(self.token, dtype=torch.int64),
+            "elapsed": self.elapsed,
+            "global_avg_losses": global_avg_losses_bytes,
+            "global_max_losses": global_max_losses_bytes,
+            "log_steps": log_steps_bytes,
+        }
+
+    def load_state_dict(self, state_dict) -> None:
+        self.step = state_dict["step"].item()
+        self.skipped_step = state_dict.get("skipped_step", 0).item()
+        self.token = state_dict["token"].item()
+        self.elapsed = state_dict["elapsed"]
+        state_dict["global_avg_losses"].seek(0)
+        self.global_avg_losses = torch.load(
+            state_dict["global_avg_losses"], weights_only=False
+        )
+        state_dict["global_max_losses"].seek(0)
+        self.global_max_losses = torch.load(
+            state_dict["global_max_losses"], weights_only=False
+        )
+        state_dict["log_steps"].seek(0)
+        self.log_steps = torch.load(state_dict["log_steps"], weights_only=False)

transformer_pp_1b_c4/transformer_pp_1b_c4_valc4_soap_pdim2048_pfreq10_lr3e_3_b1_0_9_b2_0_95_eps_1e_15_20260508_191338/exp_data/flame/config_manager.py

@@ -0,0 +1,981 @@
+# 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.
+
+import argparse
+import sys
+from collections import defaultdict
+from typing import Tuple
+
+import torch
+
+try:
+    import tomllib
+except ModuleNotFoundError:
+    import tomli as tomllib
+
+from torchtitan.tools.logging import logger
+
+TORCH_DTYPE_MAP = {
+    "float16": torch.float16,
+    "float32": torch.float32,
+    "bfloat16": torch.bfloat16,
+}
+
+
+def string_list(raw_arg):
+    """Comma-separated string list argument."""
+    return [s.strip() for s in raw_arg.split(",") if s.strip()]
+
+
+def check_string_list_argument(args_dict: dict[str, any], fullargname: str):
+    section, name = fullargname.split(".")
+    # Split string list which are still raw strings.
+    if (
+        section in args_dict
+        and name in args_dict[section]
+        and isinstance(args_dict[section][name], str)
+    ):
+        sec = args_dict[section]
+        sec[name] = string_list(sec[name])
+
+
+class JobConfig:
+    """
+    A helper class to manage the train configuration.
+    Semantics:
+    - Default config is loaded from a toml file. If no toml file is provided,
+    then the default config is loaded from argparse defaults.
+    - if toml file has missing keys, they are filled with argparse defaults.
+    - if additional explicit cmd args are provided in addition to the toml
+    file, they will override the toml config and the argparse defaults
+
+    precedence order: cmdline > toml > argparse default
+
+    Arg parsing semantics:
+
+    Each argument starts with <prefix>_ which is the section name in the toml file
+    followed by name of the option in the toml file. For ex,
+    model.name translates to:
+        [model]
+        name
+    in the toml file
+    """
+
+    def __init__(self):
+        self.args_dict = None
+        # main parser
+        self.parser = argparse.ArgumentParser(description="torchtitan arg parser.")
+
+        self.parser.add_argument(
+            "--job.config_file",
+            type=str,
+            default=None,
+            help="Job config file",
+        )
+
+        # job level configs
+        self.parser.add_argument(
+            "--job.dump_folder",
+            type=str,
+            default="./torchtitan/outputs",
+            help="Folder to dump job outputs",
+        )
+        self.parser.add_argument(
+            "--job.description",
+            type=str,
+            default="default job",
+            help="Description of the job",
+        )
+        self.parser.add_argument(
+            "--job.use_for_integration_test",
+            action="store_true",
+            help="Add this config to the integration test suite",
+        )
+        self.parser.add_argument(
+            "--job.print_args",
+            action="store_true",
+            help="Print the args to terminal",
+        )
+
+        # model configs
+        self.parser.add_argument(
+            "--model.name",
+            type=str,
+            default="fla",
+            help="Which model to train",
+        )
+        self.parser.add_argument(
+            "--model.config",
+            type=str,
+            default="fla-hub/transformer-1.3B-100B",
+            help="Path to the model config",
+        )
+        self.parser.add_argument(
+            "--model.tokenizer_path",
+            type=str,
+            default="fla-hub/transformer-1.3B-100B",
+            help="Tokenizer path",
+        )
+        self.parser.add_argument(
+            "--model.converters",
+            type=string_list,
+            nargs="+",
+            default=[],
+            help="""
+                Comma separated list of converters to apply to the model.
+                For instance, the `float8` converter swaps `torch.nn.Linear`
+                with `Float8Linear`. This feature requires you to install 'torchao'
+                which can be found here: https://github.com/pytorch/ao
+            """,
+        )
+        self.parser.add_argument(
+            "--model.print_after_conversion",
+            action="store_true",
+            help="""
+            If true, model definition will be printed to stdout after all model
+            converters have been applied.
+            """,
+        )
+
+        # profiling configs
+        self.parser.add_argument(
+            "--profiling.enable_profiling",
+            action="store_true",
+            help="Whether to enable pytorch profiler",
+        )
+        self.parser.add_argument(
+            "--profiling.save_traces_folder",
+            type=str,
+            default="profile_traces",
+            help="Trace files location",
+        )
+        self.parser.add_argument(
+            "--profiling.profile_freq",
+            type=int,
+            default=10,
+            help="How often to collect profiler traces, in iterations",
+        )
+        self.parser.add_argument(
+            "--profiling.enable_memory_snapshot",
+            action="store_true",
+            help="Whether to dump memory snapshot",
+        )
+        self.parser.add_argument(
+            "--profiling.save_memory_snapshot_folder",
+            type=str,
+            default="memory_snapshot",
+            help="Memeory snapshot files location",
+        )
+
+        # optimizer configs
+        self.parser.add_argument(
+            "--optimizer.name", type=str, default="AdamW", help="Optimizer to use"
+        )
+        self.parser.add_argument(
+            "--optimizer.eps",
+            type=float,
+            default=1e-8,
+            help="Epsilon value for the optimizer.",
+        )
+        self.parser.add_argument(
+            "--optimizer.lr", type=float, default=8e-4, help="Learning rate to use"
+        )
+        self.parser.add_argument(
+            "--optimizer.beta1", type=float, default=0.9,
+            help="Exponential moving average hyperparameters to use"
+        )
+        self.parser.add_argument(
+            "--optimizer.beta2", type=float, default=0.95,
+            help="Exponential moving average hyperparameters to use"
+        )
+        self.parser.add_argument(
+            "--optimizer.weight_decay", type=float, default=0.1,
+            help="Weight decay to use"
+        )
+        self.parser.add_argument(
+            "--optimizer.implementation",
+            type=str,
+            default="fused",
+            choices=["for-loop", "foreach", "fused"],
+            help="""
+            Specify which optimizer implementation to use:
+            - 'fused': Use fused implementation (CUDA only) for best performance.
+            - 'foreach': Use some horizontal fusion of tensors for better performance.
+            - 'for-loop': Use the default implementation for the optimizer (slowest).
+            - more info: https://pytorch.org/docs/stable/optim.html
+            """,
+        )
+        self.parser.add_argument(
+            "--optimizer.early_step_in_backward",
+            action="store_true",
+            help="""
+            Whether to apply optimizer in the backward. Caution, optimizer_in_backward
+            is not compatible with gradients clipping, users should not call
+            register_post_accumulate_grad_hook after the optimizer is built.""",
+        )
+
+        # lr scheduler configs
+        self.parser.add_argument(
+            "--lr_scheduler.warmup_steps",
+            type=int,
+            default=200,
+            help="Steps for lr scheduler warmup, normally 1/5 of --training.steps",
+        )
+        self.parser.add_argument(
+            "--lr_scheduler.decay_ratio",
+            type=float,
+            default=None,
+            help="""
+            Controls the proportion of the training steps allocated to the learning rate decay phase.
+
+            If `None`, the learning rate will begin decaying immediately after the warmup period.
+            Otherwise, the learning rate will remain stable after the warmup period and
+            only start decaying during the last `decay_ratio` portion of the total training steps.
+
+            This is known as the Warmup-Stable-Decay (WSD) schedule, as described in https://arxiv.org/abs/2404.06395.
+            """,
+        )
+        self.parser.add_argument(
+            "--lr_scheduler.decay_type",
+            type=str,
+            default="linear",
+            choices=["linear", "sqrt", "cosine"],
+            help="""
+            Learning rate decay type to use during training:
+            - 'linear': linearly decays learning rate from initial to final value
+            - 'sqrt': decays learning rate following a 1 minus square root curve
+            - 'cosine': smoothly decays learning rate following a cosine curve
+            """,
+        )
+        self.parser.add_argument(
+            "--lr_scheduler.lr_min",
+            type=float,
+            default=0.0,
+            help="""
+            Min lr ratio for lr scheduler.
+
+            If provided, the range of decay factor is scaled from 1 to `lr_min`
+            to ensure the learning rate does not drop below `optimizer.lr * lr_scheduler.lr_min`.
+            """,
+        )
+
+        # training configs
+        self.parser.add_argument(
+            "--training.batch_size", type=int, default=8, help="Batch size"
+        )
+        self.parser.add_argument(
+            "--training.seq_len", type=int, default=2048, help="Sequence length"
+        )
+        self.parser.add_argument(
+            "--training.context_len",
+            type=int,
+            default=2048,
+            help="Max length allowed for each sequence",
+        )
+        self.parser.add_argument(
+            "--training.varlen",
+            action="store_true",
+            help="Whether to take sequences of variable length as input",
+        )
+        self.parser.add_argument(
+            "--training.gradient_accumulation_steps",
+            type=int,
+            default=1,
+            help="Number of steps to accumulate gradients before updating parameters",
+        )
+        self.parser.add_argument(
+            "--training.steps",
+            type=int,
+            default=10000,
+            help="How many train steps to run",
+        )
+        self.parser.add_argument(
+            "--training.max_norm",
+            type=float,
+            default=1.0,
+            help="Max norm for gradient clipping",
+        )
+        self.parser.add_argument(
+            "--training.skip_nan_inf",
+            action="store_true",
+            help="Skip batch updates when NaN or INF gradients are encountered during training",
+        )
+        self.parser.add_argument(
+            "--training.dataset",
+            default="HuggingFaceFW/fineweb-edu",
+            help="Dataset to use, with comma separated values",
+        )
+        self.parser.add_argument(
+            "--training.dataset_name",
+            default=None,
+            help="The name of the dataset config, with comma separated values if provided",
+        )
+        self.parser.add_argument(
+            "--training.dataset_split",
+            default=None,
+            help="Dataset split to use, with comma separated values if provided",
+        )
+        self.parser.add_argument(
+            "--training.data_dir",
+            default=None,
+            help="Data dirs to use, with comma separated values if provided",
+        )
+        self.parser.add_argument(
+            "--training.data_files",
+            default=None,
+            help="Data files to use, with comma separated values if provided",
+        )
+        self.parser.add_argument(
+            "--training.data_probs",
+            default=None,
+            help="Data sampling probabilities, with comma separated values if provided",
+        )
+        self.parser.add_argument(
+            "--training.streaming",
+            action="store_true",
+            help="Whether to load dataset in streaming mode, used for huge dataset",
+        )
+        self.parser.add_argument(
+            "--training.num_workers",
+            type=int,
+            default=32,
+            help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
+        )
+        self.parser.add_argument(
+            "--training.prefetch_factor",
+            type=int,
+            default=2,
+            help="Number of batches loaded in advance by each worker."
+            "2 means there will be a total of 2 * num_workers batches prefetched across all workers.",
+        )
+        self.parser.add_argument(
+            "--training.data_parallel_replicate_degree",
+            type=int,
+            default=1,
+            help="""
+            The `data_parallel_replicate_degree` argument specifies the degree of
+            data parallelism for weight replication. When this value is greater
+            than 1, weights will be replicated across `data_parallel_replicate_degree`
+            ranks. If `data_parallel_shard_degree` is also greater than 1, the parallelism
+            method used is HSDP (Hybrid Sharded Data Parallelism). Otherwise, the
+            parallelism method used is DDP (Distributed Data Parallelism).
+            1 means disabled.""",
+        )
+        self.parser.add_argument(
+            "--training.data_parallel_shard_degree",
+            type=int,
+            default=-1,
+            help="""
+            The `data_parallel_shard_degree` argument specifies the degree of data
+            parallelism for weight sharding. When this value is greater than 1, weights
+            will be sharded across `data_parallel_shard_degree` ranks. If
+            `data_parallel_replicate_degree` is also greater than 1, the parallelism
+            method used is HSDP (Hybrid Sharded Data Parallelism).  Otherwise, the
+            parallelism method used is FSDP (Fully Sharded Data Parallelism).
+
+            -1 means leftover ranks will be used (After DP_REPLICATE/SP/PP). Note that
+            only `data_parallel_shard_degree` can be negative. 1 means disabled.""",
+        )
+        self.parser.add_argument(
+            "--training.enable_cpu_offload",
+            action="store_true",
+            help="""
+            Whether to apply CPU offloading of parameters, gradients, and optimizer states in FSDP""",
+        )
+        self.parser.add_argument(
+            "--training.tensor_parallel_degree",
+            type=int,
+            default=1,
+            help="Tensor Parallelism degree. 1 means disabled.",
+        )
+        self.parser.add_argument(
+            "--training.disable_loss_parallel",
+            action="store_true",
+            help="Whether to apply loss parallel when sequence parallel is enabled",
+        )
+        self.parser.add_argument(
+            "--training.fsdp_reshard_after_forward",
+            type=str,
+            default="default",
+            choices=["default", "always", "never"],
+            help="""
+            `reshard_after_forward` specifies the policy for applying `reshard_after_forward`
+            within an FSDP setup. `reshard_after_forward` controls parameter behavior after forward,
+            trading off memory and communication. See torch's `fully_shard` API for more documentation
+            on `reshard_after_forward`.
+            The supported policies include "default", "always" and "never":
+            - "default" applies default resharding behavior, implementing "smart defaults" for known optimal
+              scenarios.
+            - "always" will enable `reshard_after_forward` for all forward passes.
+            - "never" will disable `reshard_after_forward` for all forward passes.
+            """,
+        )
+        self.parser.add_argument(
+            "--training.mixed_precision_param",
+            type=str,
+            default="bfloat16",
+            choices=["bfloat16", "float32"],
+            help="""
+                torch dtype to use for parameters when applying mixed precision via fully_shard or torch.autocast.
+                This feature takes effect via fully_shard when data_parallel_shard_degree > 1 or
+                context_parallel_degree > 1; it takes effect via torch.autocast when data_replicate_degree >= 1
+                and no other parallelism is enabled, i.e. under DDP or single-device training.
+            """,
+        )
+        self.parser.add_argument(
+            "--training.mixed_precision_reduce",
+            type=str,
+            default="float32",
+            choices=["float32"],
+            help="""
+                torch dtype to use for reductions when applying mixed precision via FSDP.
+                This feature only takes effect when data_parallel_shard_degree > 1
+            """,
+        )
+        self.parser.add_argument(
+            "--training.compile",
+            action="store_true",
+            help="Whether to compile the model",
+        )
+        self.parser.add_argument(
+            "--training.gc_freq",
+            type=int,
+            default=50,
+            help="Python garbage control scheduling interval, in steps",
+        )
+        self.parser.add_argument(
+            "--training.seed",
+            type=int,
+            default=42,
+            help="Choose the base RNG seed used for training",
+        )
+        self.parser.add_argument(
+            "--training.deterministic",
+            action="store_true",
+            help="Use deterministic algorithms wherever possible, may be slower",
+        )
+        # ------ jinxin ------ #
+        self.parser.add_argument(
+            "--training.val_times",
+            type=int,
+            default=0,
+            help="Number of times to evaluate val PPL during training. 0 means no intermediate eval. "
+                 "e.g. 10 means evaluate every (total_steps // 10) steps.",
+        )
+        self.parser.add_argument(
+            "--training.val_data_dir",
+            type=str,
+            default=None,
+            help="Path to the validation data directory containing parquet files. "
+                 "If None, defaults to 'data/wiki_val/' relative to cwd.",
+        )
+        # metrics configs
+        self.parser.add_argument(
+            "--metrics.log_freq",
+            type=int,
+            default=10,
+            help="How often to log metrics to TensorBoard, in iterations",
+        )
+        self.parser.add_argument(
+            "--metrics.enable_tensorboard",
+            action="store_true",
+            help="Whether to log metrics to TensorBoard",
+        )
+        self.parser.add_argument(
+            "--metrics.disable_color_printing",
+            action="store_true",
+            help="Whether to disable color printing in logs",
+        )
+        self.parser.add_argument(
+            "--metrics.save_tb_folder",
+            type=str,
+            default="tb",
+            help="Folder to dump TensorBoard states",
+        )
+        self.parser.add_argument(
+            "--metrics.save_for_all_ranks",
+            action="store_true",
+            default=False,
+            help="""
+                Whether to save TensorBoard/Wandb metrics only for rank 0 or for all ranks.
+                When this option is False and pipeline_parallel_degree is > 1, the metrics
+                component uses the 0th rank of the last stage pipeline group, which is the
+                only stage that computes loss metrics.
+            """,
+        )
+        self.parser.add_argument(
+            "--metrics.enable_wandb",
+            action="store_true",
+            help="Whether to log metrics to Weights & Biases",
+        )
+        self.parser.add_argument(
+            "--no-metrics.enable_wandb",
+            dest="metrics.enable_wandb",
+            action="store_false",
+            help="Disable Weights & Biases logging (e.g. to avoid disk quota on ~/.cache/wandb)",
+        )
+
+        self.parser.add_argument(
+            "--experimental.enable_async_tensor_parallel",
+            action="store_true",
+            help="Whether to apply async tensor parallel (currently only effective when compile is enabled)",
+        )
+        self.parser.add_argument(
+            "--experimental.pipeline_parallel_degree",
+            type=int,
+            default=1,
+            help="""
+                Pipeline Parallelism degree, or number of ranks. 1 means disabled.
+                If using looped schedules, this still specifies the number of physical ranks, not the number
+                of stages.  Stages per rank are inferred from split points degree, and schedule.""",
+        )
+        self.parser.add_argument(
+            "--experimental.pipeline_parallel_split_points",
+            type=string_list,
+            nargs="+",
+            default=[],
+            help="""
+                Specify comma-separated names of modules to use as the beginning of a split point.
+
+                e.g. "layers.0,layers.2" will cause the model to be split into 3 stages,
+                the first containing all the layers up to layers.0,
+                the second containing layers.0 and up to layers.2,
+                the third containing layers.2 and all the remaining layers.
+
+                Note: fully-automated splitting may be enabled in the future,
+                but currently the split points must be specified manually.""",
+        )
+        self.parser.add_argument(
+            "--experimental.pipeline_parallel_schedule",
+            type=str,
+            default="1F1B",
+            help="""
+                Specify the Pipeline Parallel schedule to use. The supported schedules are:
+                https://github.com/pytorch/pytorch/blob/de4c2a3b4e89d96334dc678d1c3f2ae51a6630a0/torch/distributed/pipelining/schedules.py#L2161.
+                The schedule must be compatible with the split points and stages_per_rank.
+
+                Looped schedules (e.g. Interleaved1F1B) require specifying pipeline_parallel_degree = number of ranks,
+                and split_points = number of stages - 1
+                """,
+        )
+        self.parser.add_argument(
+            "--experimental.pipeline_parallel_schedule_csv",
+            type=str,
+            default="",
+            help="""
+                Specify the path to the pipeline parallel schedule csv file to use.
+                The pipeline_parallel_schedule argument must be either
+                PipelineScheduleSingle, PipelineScheduleMulti, or _PipelineScheduleRuntime.
+            """,
+        )
+
+        self.parser.add_argument(
+            "--experimental.pipeline_parallel_microbatches",
+            type=int,
+            default=None,
+            help="""
+                How many microbatches to split the global training batch into when using pipeline parallelism.
+
+                The global training batch size must be evenly divisible by the number of microbatches.
+
+                The default value will be the number of pipeline stages, if unspecified.
+            """,
+        )
+        self.parser.add_argument(
+            "--experimental.enable_compiled_autograd",
+            action="store_true",
+            help="Enable CompiledAutograd to compile the backward.",
+        )
+        self.parser.add_argument(
+            "--experimental.context_parallel_degree",
+            type=int,
+            default=1,
+            help="Context parallelism degree. 1 means disabled.",
+        )
+        self.parser.add_argument(
+            "--experimental.context_parallel_rotate_method",
+            type=str,
+            default="allgather",
+            help="""
+                The collective to use in context parallel SDPA for kv shards exchange.
+
+                'allgather' means to all-gather all kv shards on ranks after the first sub-SDPA computation,
+
+                'alltoall' means to all-to-all shuffle the kv shards.
+
+                The default value is 'allgather'.
+            """,
+        )
+        # I'm not particularly fond of this. Users can choose to write their own wrapper
+        # module and import TorchTitan training loop and execute it, which look cleaner.
+        # One reason to provide this option is to allow users to use the existing run script.
+        # While the script is pretty trivial now, we may add more logic when integrating
+        # with TorchFT.
+        # This option is subject to change and may be deleted in the future.
+        self.parser.add_argument(
+            "--experimental.custom_model_path",
+            type=str,
+            default="",
+            help="""
+                The --custom_model_path option allows to specify a custom path to a model module
+                that is not natively implemented within TorchTitan.
+                Acceptable values are the file system path to the module (e.g., my_models/model_x)
+                dotted import module  (e.g., some_package.model_x).
+            """,
+        )
+        # checkpointing configs
+        self.parser.add_argument(
+            "--checkpoint.enable_checkpoint",
+            action="store_true",
+            help="Whether to enable checkpoint",
+        )
+        self.parser.add_argument(
+            "--checkpoint.folder",
+            type=str,
+            default="checkpoint",
+            help="""
+                The folder to store the checkpoints.
+                When enable_checkpoint is set to true, checkpoints will be in {--job.dump_folder}/{--checkpoint.folder}.
+            """,
+        )
+        self.parser.add_argument(
+            "--checkpoint.initial_load_path", type=str, default=None,
+            help="""
+                This option specifies the path to the initial checkpoint to load, which is
+                particularly useful for resuming training from a previous run with a
+                different output path or when loading a checkpoint from a pre-trained model.
+                If the checkpoint folder for the current run is not empty,
+                located at {--job.dump_folder}/{--checkpoint.folder}, this option will be ignored.
+                This feature allows users to load an initial checkpoint from a different folder and
+                continue training, saving new checkpoints to the specified folder without affecting
+                the existing ones.
+            
+                Note that the path should contain the full path to the checkpoint folder,
+                including the step number, if any; for example,
+                "//pre_train/checkpoints/llama3/llama3_8b/step_10000".
+                """
+        )
+        self.parser.add_argument(
+            "--checkpoint.initial_load_model_weights_only",
+            dest='checkpoint.initial_load_model_weights_only', action="store_true", default=True,
+            help="""
+                This option specifies if only the model weights should be loaded during the initial
+                checkpoint load. The option is only used when `initial_load_path` is specified, and
+                only applies to a model_weights_only checkpoint. Loading a periodic checkpoint 
+                may lead to unexpected behavior if this option is set to True.
+                If False, the checkpoint at `initial_load_path` is treated as a standard training
+                checkpoint, including optimizer and training states.
+                The default setting for this option is True. Note that you will have to use
+                `--checkpoint.no_initial_load_model_weights_only` to override the default setting.
+            """
+        )
+        self.parser.add_argument(
+            "--checkpoint.no_initial_load_model_weights_only",
+            dest='checkpoint.initial_load_model_weights_only', action="store_false",
+        )
+        self.parser.add_argument(
+            "--checkpoint.interval",
+            type=int,
+            default=500,
+            help="Checkpointing interval in steps.",
+        )
+        self.parser.add_argument(
+            "--checkpoint.last_save_model_weights_only",
+            action="store_true",
+            help="""
+                When last_save_model_weights_only=True, only model weights will be saved at the end of training,
+                the last save.  With this, checkpoints can be loaded using `torch.load(..., weights_only=True)`
+                after conversion.  When last_save_model_weights_only=False, the full checkpoint will be saved.
+                A full checkpoint includes model, optimizer and train_state, which can be used to resume training.
+                The default value is false.
+            """,
+        )
+        self.parser.add_argument(
+            "--checkpoint.export_dtype",
+            type=str,
+            default="float32",
+            choices=["float16", "bfloat16", "float32"],
+            help="""
+                Converts to the specified precision when training completes and model_weights_only=true.
+                Currently supports float32, float16, and bfloat16.
+                The default value is float32.
+            """,
+        )
+        self.parser.add_argument(
+            "--checkpoint.create_seed_checkpoint",
+            action="store_true",
+            help="""
+                Initializes the full model without applying parallelisms, and then saves it as a seed checkpoint.
+                Note: requires user to call train.py without specifying any parallelisms, e.g. NGPU=1.
+                Could be implemented as a separate script, but this way shares more code.
+            """,
+        )
+        self.parser.add_argument(
+            "--checkpoint.async_mode",
+            type=str,
+            default="disabled",
+            help="""
+                Which async checkpoint mode to use. Currently there are 3 different modes.
+                1. "disabled": synchronized checkpointing will be used.
+                2. "async": torch.distributed.checkpoint.async_save will be used.
+                3. "async_with_pinned_mem": this option utilizes a dedicated pinned memory
+                   space and creates a separate process for faster GPU->CPU transfer
+                   performance and eliminating GIL contention. The cost is increased CPU
+                   memory usage. If insufficient CPU memory is available, performance may
+                   degrade due to memory paging. For most users, "async" should suffice as
+                   the performance overhead is typically small (on the order of tens of
+                   seconds) compared to checkpointing frequency. This mode can be employed
+                   to pursue near-zero checkpointing times (e.g., < 1 second) given
+                   appropriate hardware support such as ample CPU memory and fast PCIe.
+
+                "disabled" is the default mode.
+            """,
+        )
+        self.parser.add_argument(
+            "--checkpoint.keep_latest_k",
+            type=int,
+            default=0,
+            help="""
+                Keeps only the latest k checkpoints, and purging older ones. If 0, keep all checkpoints.
+                0 is the default value. k cannot be 1 as the last one may be in the process of being
+                saved. As a result, the metadata of the last one may not be ready yet.
+            """,
+        )
+        self.parser.add_argument(
+            "--checkpoint.load_step",
+            type=int,
+            default=-1,
+            help="Load the checkpoint at the specified step. If -1, load the latest checkpoint.",
+        )
+        self.parser.add_argument(
+            "--checkpoint.exclude_from_loading",
+            type=string_list,
+            nargs="*",
+            default=[],
+            help="""
+                Exclude specific keys from being loaded from the checkpoint.
+                Provide a comma-separated list of keys to exclude, e.g. 'optimizer,lr_scheduler,dataloader'.
+                This will load the model only, excluding the specified keys.
+            """,
+        )
+        # activation checkpointing configs
+        self.parser.add_argument(
+            "--activation_checkpoint.mode",
+            type=str,
+            default="selective",
+            help="Type of activation checkpointing to use ['none', 'full', 'selective']",
+        )
+        self.parser.add_argument(
+            "--activation_checkpoint.selective_ac_option",
+            type=str,
+            default="2",  # 2 = checkpoint every other layer
+            help="""
+                Selective activation checkpointing options ['int', 'op'].
+                'int' (e.g., 2) for every nth layer, or 'op' for op level ac.
+            """,
+        )
+
+        self.parser.add_argument(
+            "--activation_offload.mode",
+            type=str,
+            default="none",
+            help="""
+                if we are using activation offload or not. Options are ['none', 'full'].
+            """,
+        )
+
+        # float8 configs
+        self.parser.add_argument(
+            "--float8.enable_fsdp_float8_all_gather",
+            action="store_true",
+            help="Whether enable float8 all-gather in FSDP, recommended for tensorwise scaling",
+        )
+        self.parser.add_argument(
+            "--float8.precompute_float8_dynamic_scale_for_fsdp",
+            action="store_true",
+            help="Whether precompute float8 scales dynamically for FSDP, recommended for tensorwise scaling",
+        )
+        self.parser.add_argument(
+            "--float8.force_recompute_fp8_weight_in_bwd",
+            action="store_true",
+            help="""
+            Whether to force the recomputation of FP8 weights during backward pass.
+            When using FSDP with tensorwise scaling, it is recommended to enable
+            `force_recompute_fp8_weight_in_bwd` to prevent saving unsharded FP8 weights
+            for backward computation.
+            """,
+        )
+        self.parser.add_argument(
+            "--float8.recipe_name",
+            type=str,
+            default=None,
+            choices=["tensorwise", "rowwise", "rowwise_with_gw_hp"],
+            help="""
+            If specified, creates float8 config from recipe name, valid choices are
+            `tensorwise`, `rowwise` and `rowwise_with_gw_hp`.
+            """,
+        )
+
+        # communications library settings
+        self.parser.add_argument(
+            "--comm.init_timeout_seconds",
+            type=int,
+            default=300,
+            help="Timeout for communication operations, during initialization and first train step.",
+        )
+        self.parser.add_argument(
+            "--comm.train_timeout_seconds",
+            type=int,
+            default=100,
+            help=(
+                "Timeout for communication operations after the first train step -- "
+                "usually a tighter bound than during initialization."
+            ),
+        )
+        self.parser.add_argument(
+            "--comm.trace_buf_size",
+            type=int,
+            default=20000,
+            help="Flight recorder ring buffer size, >0 means recording by default, 0 means disabled",
+        )
+
+        # memory estimation settings
+        self.parser.add_argument(
+            "--memory_estimation.enabled",
+            help="Whether to estimate memory usage for FSDP",
+            action="store_true",
+        )
+
+        self.parser.add_argument(
+            "--memory_estimation.disable_fake_mode",
+            help="Whether to estimate memory under FakeTensorMode",
+            action="store_true",
+        )
+
+        self.parser.add_argument(
+            "--fault_tolerance.enable",
+            action="store_true",
+            help="""
+                Enable TorchFT integration. When TorchFT is enabled, HSDP will be used.
+                And --fault_tolerance.data_parallel_replicate_degree should be 1 and
+                --fault_tolerance.group_size will be used to control the maximum
+                replicate group size as the replicate group size is dynamic.
+
+                Note that this is still an experimental feature.
+            """,
+        )
+
+        self.parser.add_argument(
+            "--fault_tolerance.replica_id",
+            type=int,
+            default=0,
+            help="The TorchFT replica ID of this run.",
+        )
+
+        self.parser.add_argument(
+            "--fault_tolerance.group_size",
+            type=int,
+            default=0,
+            help="""
+                The number of TorchFT replicate groups. This number will be used for
+                dataloader to split the dataset across the replicate groups and FSDP
+                dimension
+            """,
+        )
+
+        self.parser.add_argument(
+            "--fault_tolerance.min_replica_size",
+            type=int,
+            default=1,
+            help="The minimum number of FT replica for each step.",
+        )
+
+    def to_dict(self):
+        return self.args_dict
+
+    def parse_args(self, args_list: list = sys.argv[1:]):
+        args, cmd_args = self.parse_args_from_command_line(args_list)
+        config_file = getattr(args, "job.config_file", None)
+        # build up a two level dict
+        args_dict = self._args_to_two_level_dict(args)
+        if config_file is not None:
+            try:
+                with open(config_file, "rb") as f:
+                    for k, v in tomllib.load(f).items():
+                        # to prevent overwrite of non-specified keys
+                        args_dict[k] |= v
+            except (FileNotFoundError, tomllib.TOMLDecodeError) as e:
+                logger.exception(
+                    f"Error while loading the configuration file: {config_file}"
+                )
+                logger.exception(f"Error details: {str(e)}")
+                raise e
+
+        # Checking string-list arguments are properly split into a list
+        # if split-points came from 'args' (from cmd line) it would have already been parsed into a list by that parser
+        string_list_argnames = self._get_string_list_argument_names()
+        for n in string_list_argnames:
+            check_string_list_argument(args_dict, n)
+
+        # override args dict with cmd_args
+        cmd_args_dict = self._args_to_two_level_dict(cmd_args)
+        for section, section_args in cmd_args_dict.items():
+            for k, v in section_args.items():
+                args_dict[section][k] = v
+
+        self.args_dict = args_dict
+
+        for k, v in args_dict.items():
+            class_type = type(k.title(), (), v)
+            setattr(self, k, class_type())
+        self._validate_config()
+
+    def _args_to_two_level_dict(self, args: argparse.Namespace) -> defaultdict:
+        args_dict = defaultdict(defaultdict)
+        for k, v in vars(args).items():
+            first_level_key, second_level_key = k.split(".", 1)
+            args_dict[first_level_key][second_level_key] = v
+        return args_dict
+
+    def _validate_config(self) -> None:
+        # TODO: Add more mandatory validations
+        assert self.model.config
+        assert self.model.tokenizer_path
+
+    def _get_string_list_argument_names(self) -> list[str]:
+        """Get the parser argument names of type `string_list`."""
+        string_list_args = [
+            v.dest for v in self.parser._actions if v.type is string_list
+        ]
+        return string_list_args
+
+    def parse_args_from_command_line(
+        self, args_list
+    ) -> Tuple[argparse.Namespace, argparse.Namespace]:
+        """
+        Parse command line arguments and return the parsed args and the command line only args
+        """
+        args = self.parser.parse_args(args_list)
+        string_list_argnames = set(self._get_string_list_argument_names())
+
+        # aux parser to parse the command line only args, with no defaults from main parser
+        aux_parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS)
+        for arg, val in vars(args).items():
+            if isinstance(val, bool):
+                aux_parser.add_argument(
+                    "--" + arg, action="store_true" if val else "store_false"
+                )
+            elif arg in string_list_argnames:
+                # without this special case, type inference breaks here,
+                # since the inferred type is just 'list' and it ends up flattening
+                # e.g. from ["layers.0", "layers.1"] into ["l", "a", "y", "e", "r", "s", ".0", ...]
+                aux_parser.add_argument("--" + arg, type=string_list)
+            else:
+                aux_parser.add_argument("--" + arg, type=type(val))
+
+        cmd_args, _ = aux_parser.parse_known_args(args_list)
+
+        return args, cmd_args