Add 1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741

.gitattributes

@@ -681,3 +681,12 @@ transformer_pp_340m_c4/transformer_pp_340m_c4_valc4_lion_lr1e_4_b1_0_9_b2_0_99_e
 1b_archs_fwe/gla_1b_fwe_adan_lr3e_3_b1_0_9_b2_0_92_b3_0_99_eps_1e_8_20260518_235004/exp_data/checkpoint/step-30720/__5_0.distcp filter=lfs diff=lfs merge=lfs -text
 1b_archs_fwe/gla_1b_fwe_adan_lr3e_3_b1_0_9_b2_0_92_b3_0_99_eps_1e_8_20260518_235004/exp_data/checkpoint/step-30720/__6_0.distcp filter=lfs diff=lfs merge=lfs -text
 1b_archs_fwe/gla_1b_fwe_adan_lr3e_3_b1_0_9_b2_0_92_b3_0_99_eps_1e_8_20260518_235004/exp_data/checkpoint/step-30720/__7_0.distcp filter=lfs diff=lfs merge=lfs -text
+1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/.metadata filter=lfs diff=lfs merge=lfs -text
+1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
+1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__1_0.distcp filter=lfs diff=lfs merge=lfs -text
+1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__2_0.distcp filter=lfs diff=lfs merge=lfs -text
+1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__3_0.distcp filter=lfs diff=lfs merge=lfs -text
+1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__4_0.distcp filter=lfs diff=lfs merge=lfs -text
+1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__5_0.distcp filter=lfs diff=lfs merge=lfs -text
+1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__6_0.distcp filter=lfs diff=lfs merge=lfs -text
+1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__7_0.distcp filter=lfs diff=lfs merge=lfs -text

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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.

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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}
+}
+```

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/.metadata

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

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__0_0.distcp

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

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__1_0.distcp

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

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__2_0.distcp

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

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__3_0.distcp

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

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__4_0.distcp

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

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__5_0.distcp

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

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__6_0.distcp

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

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/checkpoint/step-30720/__7_0.distcp

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

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+  }
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+  }
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+  }
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+    }
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/configs/gla_1B.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": 2048,
+    "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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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": []
+  }
+  

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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": []
+  }
+  

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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
+}

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/exp_data/flame/__init__.py

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

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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)

1b_archs_fwe/transformer_1b_fwe_lion_lr3e_4_b1_0_9_b2_0_99_20260519_061741/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)