Add files using upload-large-folder tool

README.md

@@ -0,0 +1,512 @@
+<div align="center">
+
+# 🔥 Flame: Flash Linear Attention Made Easy
+# This is a fork for the paper:
+# Softpick: No Attention Sink, No Massive Activations with Rectified Softmax
+
+</div>
+
+## Instructions for Softpick Attention
+
+This fork can only work on an older commit of torchtitan and flame, so the setup looks like this:
+
+```bash
+git clone https://github.com/zaydzuhri/flame.git
+cd flame
+git checkout softpick-attention
+git submodule update --init --recursive --remote
+cd 3rdparty/torchtitan
+git checkout 4f532e0
+cd ../../
+
+pip install .
+pip install flash-attn --no-build-isolation
+```
+The flash-linear-attention submodule has been changed to link to our fork: https://github.com/zaydzuhri/flash-linear-attention/tree/softpick-attention
+So no need to manually clone it.
+
+Then prepare the fineweb-edu 100B sample the same way as described in the flame repo guide below.
+
+These are the training commands used in the paper:
+```bash
+NGPU=8 bash train.sh   --job.config_file flame/models/fla.toml   --job.dump_folder exp/vanilla.340M.batch16.seqlen4096.context4096.warmup1000.update1.steps100000.lr3e-4.cosine   --model.config configs/vanilla_transformer_340M.json   --model.tokenizer_path fla-hub/transformer-1.3B-100B   --optimizer.name AdamW   --optimizer.eps 1e-15   --optimizer.lr 3e-4   --lr_scheduler.warmup_steps 1000   --lr_scheduler.lr_min 0.1   --lr_scheduler.decay_type cosine   --training.batch_size 16   --training.seq_len 4096   --training.context_len 4096   --training.gradient_accumulation_steps 1   --training.steps 100000   --training.max_norm 1.0   --training.skip_nan_inf  --training.dataset ~/.cache/HuggingFaceFW___fineweb-edu/sample-100BT  --training.dataset_split train   --training.num_workers 32   --training.prefetch_factor 2   --training.seed 79   --training.compile   --checkpoint.interval 10000   --checkpoint.load_step -1   --metrics.log_freq 5 --checkpoint.hf_upload_enabled   --checkpoint.hf_repo_base_name "zaydzuhri/vanilla-340M-4096-batch16-steps100000" --comm.init_timeout_seconds 600 --comm.train_timeout_seconds 300
+
+NGPU=8 bash train.sh   --job.config_file flame/models/fla.toml   --job.dump_folder exp/softpick.340M.batch16.seqlen4096.context4096.warmup1000.update1.steps100000.lr3e-4.cosine   --model.config configs/softpick_transformer_340M.json   --model.tokenizer_path fla-hub/transformer-1.3B-100B   --optimizer.name AdamW   --optimizer.eps 1e-15   --optimizer.lr 3e-4   --lr_scheduler.warmup_steps 1000   --lr_scheduler.lr_min 0.1   --lr_scheduler.decay_type cosine   --training.batch_size 16   --training.seq_len 4096   --training.context_len 4096   --training.gradient_accumulation_steps 1   --training.steps 100000   --training.max_norm 1.0   --training.skip_nan_inf  --training.dataset ~/.cache/HuggingFaceFW___fineweb-edu/sample-100BT  --training.dataset_split train   --training.num_workers 32   --training.prefetch_factor 2   --training.seed 79   --training.compile   --checkpoint.interval 10000   --checkpoint.load_step -1   --metrics.log_freq 5 --checkpoint.hf_upload_enabled   --checkpoint.hf_repo_base_name "zaydzuhri/softpick-340M-4096-batch16-steps100000" --comm.init_timeout_seconds 600 --comm.train_timeout_seconds 300
+```
+
+And the same for the extra experiments in the appendix:
+```bash
+NGPU=8 bash train.sh   --job.config_file flame/models/fla.toml   --job.dump_folder exp/rectified.340M.batch16.seqlen4096.context4096.warmup1000.update1.steps100000.lr3e-4.cosine   --model.config configs/rectified_transformer_340M.json   --model.tokenizer_path fla-hub/transformer-1.3B-100B   --optimizer.name AdamW   --optimizer.eps 1e-15   --optimizer.lr 3e-4   --lr_scheduler.warmup_steps 1000   --lr_scheduler.lr_min 0.1   --lr_scheduler.decay_type cosine   --training.batch_size 16   --training.seq_len 4096   --training.context_len 4096   --training.gradient_accumulation_steps 1   --training.steps 100000   --training.max_norm 1.0   --training.skip_nan_inf  --training.dataset ~/.cache/HuggingFaceFW___fineweb-edu/sample-100BT  --training.dataset_split train   --training.num_workers 32   --training.prefetch_factor 2   --training.seed 79   --training.compile   --checkpoint.interval 10000   --checkpoint.load_step -1   --metrics.log_freq 5 --checkpoint.hf_upload_enabled   --checkpoint.hf_repo_base_name "zaydzuhri/rectified-340M-4096-batch16-steps100000" --comm.init_timeout_seconds 600 --comm.train_timeout_seconds 300
+
+NGPU=8 bash train.sh   --job.config_file flame/models/fla.toml   --job.dump_folder exp/softpick.scaled.340M.batch16.seqlen4096.context4096.warmup1000.update1.steps100000.lr3e-4.cosine   --model.config configs/softpick_scaled_transformer_340M.json   --model.tokenizer_path fla-hub/transformer-1.3B-100B   --optimizer.name AdamW   --optimizer.eps 1e-15   --optimizer.lr 3e-4   --lr_scheduler.warmup_steps 1000   --lr_scheduler.lr_min 0.1   --lr_scheduler.decay_type cosine   --training.batch_size 16   --training.seq_len 4096   --training.context_len 4096   --training.gradient_accumulation_steps 1   --training.steps 100000   --training.max_norm 1.0   --training.skip_nan_inf  --training.dataset ~/.cache/HuggingFaceFW___fineweb-edu/sample-100BT  --training.dataset_split train   --training.num_workers 32   --training.prefetch_factor 2   --training.seed 79   --training.compile   --checkpoint.interval 10000   --checkpoint.load_step -1   --metrics.log_freq 5 --checkpoint.hf_upload_enabled   --checkpoint.hf_repo_base_name "zaydzuhri/softpick-scaled-340M-4096-batch16-steps100000" --comm.init_timeout_seconds 600 --comm.train_timeout_seconds 300
+```
+
+Feel free to DM @zmkzmkz on X for any questions regarding the paper or this code!
+
+## Flame
+
+Welcome to 🔥 `flame`, a minimal and efficient framework built on `torchtitan` for training Flash Linear Attention (FLA) models (and more broadly, arbitrary autoregressive 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 .
+```
+
+`flame` manages minimal dependencies, only including `fla` and `torchtitan` as submodules.
+After installation, initialize and update the submodules:
+```sh
+git submodule update --init --recursive
+```
+
+## 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.
+
+> [!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.lr3e-4.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 3e-4 \
+  --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.streaming \
+  --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 3e-4 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 `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.

config.json

@@ -0,0 +1,33 @@
+{
+  "architectures": [
+    "TransformerForCausalLM"
+  ],
+  "attention_bias": false,
+  "attn_impl": "parallel_softpick_attn",
+  "bos_token_id": 1,
+  "elementwise_affine": true,
+  "eos_token_id": 2,
+  "fuse_cross_entropy": false,
+  "fuse_norm": false,
+  "fuse_swiglu": true,
+  "hidden_act": "swish",
+  "hidden_ratio": 4,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": null,
+  "max_position_embeddings": 4096,
+  "model_type": "transformer",
+  "norm_eps": 1e-06,
+  "num_heads": 12,
+  "num_hidden_layers": 14,
+  "num_kv_heads": null,
+  "qk_norm": false,
+  "qkv_bias": false,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.51.3",
+  "use_cache": true,
+  "vocab_size": 32000,
+  "window_size": null
+}

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.006,
+    "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
+}

configs/delta_net_340M.json

@@ -0,0 +1,27 @@
+{
+    "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.006,
+    "intermediate_size": null,
+    "model_type": "delta_net",
+    "norm_eps": 1e-06,
+    "norm_first": false,
+    "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
+}

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.006,
+  "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
+}

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.006,
+    "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
+}

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.006,
+    "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
+}

configs/hgrn2_340M.json

@@ -0,0 +1,20 @@
+{
+    "attn_mode": "chunk",
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "expand_ratio": 128,
+    "fuse_cross_entropy": true,
+    "fuse_norm": true,
+    "hidden_act": "swish",
+    "hidden_ratio": 4,
+    "hidden_size": 1024,
+    "initializer_range": 0.006,
+    "intermediate_size": null,
+    "model_type": "hgrn2",
+    "num_heads": 8,
+    "num_hidden_layers": 24,
+    "norm_eps": 1e-06,
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "vocab_size": 32000
+}

configs/rectified_transformer_340M.json

@@ -0,0 +1,19 @@
+{
+    "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.006,
+    "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,
+    "attn_impl": "parallel_rectified_attn"
+}

configs/scaled_softpick_transformer_120M.json

@@ -0,0 +1,19 @@
+{
+    "attention_bias": false,
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "fuse_cross_entropy": true,
+    "fuse_norm": false,
+    "hidden_act": "swish",
+    "hidden_size": 768,
+    "initializer_range": 0.02,
+    "max_position_embeddings": 4096,
+    "model_type": "transformer",
+    "num_heads": 12,
+    "num_hidden_layers": 14,
+    "norm_eps": 1e-06,
+    "tie_word_embeddings": true,
+    "use_cache": true,
+    "vocab_size": 32000,
+    "attn_impl": "parallel_scaled_softpick_attn"
+}

configs/scaled_vanilla_transformer_120M.json

@@ -0,0 +1,19 @@
+{
+    "attention_bias": false,
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "fuse_cross_entropy": true,
+    "fuse_norm": false,
+    "hidden_act": "swish",
+    "hidden_size": 768,
+    "initializer_range": 0.02,
+    "max_position_embeddings": 4096,
+    "model_type": "transformer",
+    "num_heads": 12,
+    "num_hidden_layers": 14,
+    "norm_eps": 1e-06,
+    "tie_word_embeddings": true,
+    "use_cache": true,
+    "vocab_size": 32000,
+    "attn_impl": "parallel_scaled_attn"
+}

configs/scaled_vanilla_transformer_340M.json

@@ -0,0 +1,19 @@
+{
+    "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.006,
+    "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,
+    "attn_impl": "parallel_scaled_attn"
+}

configs/softpick_transformer_1B.json

@@ -0,0 +1,23 @@
+{
+    "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.006,
+    "intermediate_size": null,
+    "max_position_embeddings": 8192,
+    "model_type": "transformer",
+    "norm_eps": 1e-06,
+    "num_heads": 32,
+    "num_hidden_layers": 32,
+    "num_kv_heads": null,
+    "pad_token_id": 2,
+    "rope_theta": 10000.0,
+    "tie_word_embeddings": false,
+    "attn_impl": "parallel_softpick_attn"
+}

download_checkpoint.py

@@ -0,0 +1,35 @@
+import argparse
+import os
+from huggingface_hub import HfApi, HfFolder, snapshot_download
+
+def main(args):
+    api = HfApi()
+    token = HfFolder.get_token()
+    experiment_checkpoint_folder = os.path.join(args.experiment_checkpoint_folder, "checkpoint")
+    os.makedirs(
+        experiment_checkpoint_folder,
+        exist_ok=True
+    )
+
+    snapshot_download(
+        repo_id=args.repo_id, 
+        token=token,
+        local_dir=experiment_checkpoint_folder,
+    )
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Download a checkpoint from Hugging Face Hub.")
+    parser.add_argument(
+        "--repo_id",
+        type=str,
+        required=True,
+        help="The repository ID on Hugging Face Hub.",
+    )
+    parser.add_argument(
+        "--experiment_checkpoint_folder",
+        type=str,
+        required=True,
+        help="The local directory to save the downloaded checkpoint.",
+    )
+    args = parser.parse_args()
+    main(args)

fla/layers/__init__.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from .abc import ABCAttention
+from .attn import Attention
+from .based import BasedLinearAttention
+from .bitattn import BitAttention
+from .delta_net import DeltaNet
+from .forgetting_attn import ForgettingAttention
+from .gated_deltanet import GatedDeltaNet
+from .gated_deltaproduct import GatedDeltaProduct
+from .gla import GatedLinearAttention
+from .gsa import GatedSlotAttention
+from .hgrn import HGRNAttention
+from .hgrn2 import HGRN2Attention
+from .lightnet import LightNetAttention
+from .linear_attn import LinearAttention
+from .multiscale_retention import MultiScaleRetention
+from .nsa import NativeSparseAttention
+from .rebased import ReBasedLinearAttention
+from .rwkv6 import RWKV6Attention
+from .rwkv7 import RWKV7Attention
+
+__all__ = [
+    'ABCAttention',
+    'Attention',
+    'BasedLinearAttention',
+    'BitAttention',
+    'DeltaNet',
+    'ForgettingAttention',
+    'GatedDeltaNet',
+    'GatedDeltaProduct',
+    'GatedLinearAttention',
+    'GatedSlotAttention',
+    'HGRNAttention',
+    'HGRN2Attention',
+    'LightNetAttention',
+    'LinearAttention',
+    'MultiScaleRetention',
+    'NativeSparseAttention',
+    'ReBasedLinearAttention',
+    'RWKV6Attention',
+    'RWKV7Attention',
+]

fla/layers/abc.py

@@ -0,0 +1,218 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from fla.modules import FusedRMSNormGated, RMSNorm, RotaryEmbedding, ShortConvolution
+from fla.modules.activations import swiglu, swish
+from fla.ops.abc.chunk import chunk_abc
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+
+class ABCAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        expand_k: float = 0.5,
+        expand_v: float = 1.0,
+        num_heads: int = 4,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        num_slots: Optional[int] = None,
+        elementwise_affine: Optional[bool] = True,
+        norm_eps: float = 1e-5,
+        gate_low_rank_dim: int = 16,
+        gate_logit_normalizer: int = 16,
+        use_rope: bool = True,
+        use_input_gate: bool = False,
+        use_output_gate: bool = True,
+        use_norm: bool = True,
+        clamp_min: Optional[float] = -32,
+        clamp_max: Optional[float] = 32,
+        layer_idx: Optional[int] = None,
+        **kwargs
+    ) -> ABCAttention:
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.num_heads = num_heads
+        self.key_dim = int(self.hidden_size * self.expand_k)
+        self.value_dim = int(self.hidden_size * self.expand_v)
+        self.head_k_dim = self.key_dim // self.num_heads
+        self.head_v_dim = self.value_dim // self.num_heads
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.gate_low_rank_dim = gate_low_rank_dim
+        self.gate_logit_normalizer = gate_logit_normalizer
+
+        self.use_rope = use_rope
+        self.use_input_gate = use_input_gate
+        self.use_output_gate = use_output_gate
+        self.use_norm = use_norm
+
+        if num_slots is None:
+            num_slots = self.head_k_dim
+        self.num_slots = num_slots
+
+        self.norm_eps = norm_eps
+
+        self.clamp_min = clamp_min
+        self.clamp_max = clamp_max
+        self.layer_idx = layer_idx
+
+        if layer_idx is None:
+            warnings.warn(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.q_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.value_dim, bias=False)
+
+        if use_output_gate:
+            self.g_proj = nn.Linear(self.hidden_size, self.value_dim, bias=False)
+        self.s_proj = nn.Linear(self.hidden_size, self.num_heads * self.num_slots, bias=False)
+        self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu')
+            self.k_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu')
+            self.v_conv1d = ShortConvolution(self.value_dim, conv_size, activation='silu')
+
+        if self.use_norm:
+            if self.use_output_gate:
+                self.g_norm = FusedRMSNormGated(
+                    hidden_size=self.head_v_dim,
+                    elementwise_affine=elementwise_affine,
+                    eps=norm_eps
+                )
+            else:
+                self.g_norm = RMSNorm(
+                    hidden_size=self.head_v_dim,
+                    elementwise_affine=elementwise_affine,
+                    eps=norm_eps
+                )
+
+        if self.use_rope:
+            self.rotary = RotaryEmbedding(self.head_k_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        last_state = None
+        if past_key_values is not None and len(past_key_values) > self.layer_idx:
+            last_state = past_key_values[self.layer_idx]
+
+        cu_seqlens = kwargs.get('cu_seqlens', None)
+        if cu_seqlens is not None:
+            raise NotImplementedError("Training with cu_seqlens is not supported yet for ABCAttention")
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens
+            )
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+
+        if self.use_input_gate:
+            q, k, v = map(lambda x: swish(x), (q, k, v))
+        # dealing with left-padding
+        if attention_mask is not None:
+            v = v.mul_(attention_mask[:, -v.shape[-2]:, None])
+
+        q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim)
+        if self.use_rope:
+            seqlen_offset = 0
+            if past_key_values is not None:
+                seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            q, k = self.rotary(q, k, seqlen_offset=seqlen_offset)
+
+        s = rearrange(self.s_proj(hidden_states), '... (h m) -> ... h m', m=self.num_slots)
+        s = s.clamp_(self.clamp_min, self.clamp_max)
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        o, recurrent_state = chunk_abc(
+            q=q,
+            k=k,
+            v=v,
+            s=s,
+            initial_state=recurrent_state,
+            output_final_state=use_cache,
+            head_first=False
+        )
+        if past_key_values is not None:
+            past_key_values.update(
+                recurrent_state=recurrent_state,
+                conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+                layer_idx=self.layer_idx,
+                offset=q.shape[1]
+            )
+
+        if self.use_norm and not self.use_output_gate:
+            o = self.g_norm(o)
+        elif self.use_output_gate:
+            g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim)
+            o = self.g_norm(o, g) if self.use_norm else swiglu(g, o)
+        o = rearrange(o, '... h d -> ... (h d)')
+        o = self.o_proj(o)
+
+        return o, None, past_key_values
+
+    def state_size(self, seq_len: int = 2048):
+        return 2 * self.num_slots * self.hidden_size

fla/layers/attn.py

@@ -0,0 +1,490 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from einops import rearrange
+from transformers.utils import logging
+
+from fla.modules import RMSNorm, RotaryEmbedding
+from fla.ops import parallel_attn, parallel_rectified_attn, parallel_softpick_attn, naive_attn, naive_rectified_attn, naive_softpick_attn
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+try:
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
+except ImportError:
+    warnings.warn(
+        "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`",
+        category=ImportWarning
+    )
+    flash_attn_func = None
+
+logger = logging.get_logger(__name__)
+
+
+class Attention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 32,
+        num_kv_heads: Optional[int] = None,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        window_size: Optional[int] = None,
+        rope_theta: Optional[float] = 10000.,
+        max_position_embeddings: Optional[int] = None,
+        layer_idx: int = None,
+        attn_impl: str = "flash_attn",
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        if num_kv_heads is None:
+            self.num_kv_heads = self.num_heads
+        else:
+            self.num_kv_heads = num_kv_heads
+        self.num_kv_groups = num_heads // self.num_kv_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.qkv_bias = qkv_bias
+        self.qk_norm = qk_norm
+
+        self.window_size = window_size
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_idx = layer_idx
+        self.attn_impl = attn_impl
+
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+
+        if "scaled" in self.attn_impl:
+            self.s = nn.Parameter(torch.empty(self.num_heads, 1))
+            self.register_buffer("logn", torch.log(torch.arange(2, self.max_position_embeddings*4+2, dtype=self.s.dtype)[:, None, None]))
+
+        if qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta)
+
+    def reset_parameters(self):
+        if "scaled" in self.attn_impl:
+            nn.init.constant_(self.s, 0.3)
+            self.logn.copy_(torch.log(torch.arange(2, self.max_position_embeddings*4+2, dtype=self.s.dtype)[:, None, None]))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+
+        q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim)
+        k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+
+        if self.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+
+        # equivalent to cu_seqlens in `flash_attn`
+        cu_seqlens = kwargs.get('cu_seqlens', None)
+
+        seqlen_offset, max_seqlen = 0, q_len
+        if past_key_values is not None:
+            seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            max_seqlen = q.shape[1] + seqlen_offset
+
+            if attention_mask is not None:
+                # to deliminate the offsets of padding tokens
+                seqlen_offset = seqlen_offset + attention_mask.sum(-1) - attention_mask.shape[-1]
+                max_seqlen = q.shape[1] + max(seqlen_offset)
+
+        if self.max_position_embeddings is not None:
+            max_seqlen = max(max_seqlen, self.max_position_embeddings)
+        q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens)
+
+        if past_key_values is not None:
+            cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0
+            k_cached, v_cached = past_key_values.update(
+                attn_state=(k.flatten(-2, -1), v.flatten(-2, -1)),
+                layer_idx=self.layer_idx,
+                offset=q_len,
+                cache_kwargs=dict(window_size=self.window_size)
+            )['attn_state']
+            if cache_has_content:
+                k, v = k_cached, v_cached
+                k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+                v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+
+        # if flash_attn_func is None:
+        #     raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first")
+
+        if "scaled" in self.attn_impl:
+            k_len = k.shape[1]
+            q = q * self.s.to(q.dtype) * self.logn[k_len-q_len:k_len].to(q.dtype)
+
+        # Contains at least one padding token in the sequence
+        if self.attn_impl == "flash_attn":
+            if attention_mask is not None:
+                q, k, v, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(q, k, v, attention_mask, q_len)
+                cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+                max_seqlen_q, max_seqlen_k = max_seq_lens
+                o = flash_attn_varlen_func(
+                    q, k, v,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_k=max_seqlen_k,
+                    causal=True,
+                    window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0)
+                )
+                o = pad_input(o, indices_q, batch_size, q_len)
+            elif cu_seqlens is not None:
+                o = flash_attn_varlen_func(
+                    q.squeeze(0), k.squeeze(0), v.squeeze(0),
+                    cu_seqlens_q=cu_seqlens,
+                    cu_seqlens_k=cu_seqlens,
+                    max_seqlen_q=max_seqlen,
+                    max_seqlen_k=max_seqlen,
+                    causal=True,
+                    window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0)
+                ).unsqueeze(0)
+            else:
+                o = flash_attn_func(
+                    q, k, v,
+                    causal=True,
+                    window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0)
+                )
+        elif self.attn_impl == "parallel_attn":
+            o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "parallel_scaled_attn":
+            o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "parallel_rectified_attn":
+            o = parallel_rectified_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "parallel_softpick_attn":
+            o = parallel_softpick_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "parallel_scaled_softpick_attn":
+            o = parallel_softpick_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "naive_attn":
+            o, attentions = naive_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "naive_scaled_attn":
+            o, attentions = naive_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "naive_rectified_attn":
+            o, attentions = naive_rectified_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "naive_softpick_attn":
+            o, attentions = naive_softpick_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "naive_scaled_softpick_attn":
+            o, attentions = naive_softpick_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        else:
+            raise ValueError(f"Unknown attention implementation: {self.attn_impl}")
+
+        o = o.reshape(batch_size, q_len, -1)
+        o = self.o_proj(o)
+
+        if not output_attentions or "parallel" in self.attn_impl or "flash" in self.attn_impl:
+            attentions = None
+
+        return o, attentions, past_key_values
+
+    def _upad_input(self, q, k, v, attention_mask, q_len):
+        batch_size, seq_len, num_key_value_heads, head_dim = k.shape
+        cache_mask = attention_mask[:, -seq_len:]
+        seqlens = cache_mask.sum(-1, dtype=torch.int32)
+        indices_k = torch.nonzero(cache_mask.flatten(), as_tuple=False).flatten()
+        max_seqlen_k = seqlens.max().item()
+        cu_seqlens_k = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))
+
+        k = index_first_axis(k.reshape(batch_size * seq_len, num_key_value_heads, head_dim), indices_k)
+        v = index_first_axis(v.reshape(batch_size * seq_len, num_key_value_heads, head_dim), indices_k)
+        if q_len == seq_len:
+            q = index_first_axis(q.reshape(batch_size * seq_len, self.num_heads, head_dim), indices_k)
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_q = max_seqlen_k
+            indices_q = indices_k
+        elif q_len == 1:
+            max_seqlen_q = 1
+            # There is a memcpy here, that is very bad.
+            cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=q.device)
+            indices_q = cu_seqlens_q[:-1]
+            q = q.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -q_len:]
+            q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, attention_mask)
+
+        return q, k, v, indices_q, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k)
+
+class StochasticSoftpickAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 32,
+        num_kv_heads: Optional[int] = None,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        window_size: Optional[int] = None,
+        rope_theta: Optional[float] = 10000.,
+        max_position_embeddings: Optional[int] = None,
+        layer_idx: int = None,
+        attn_impl: str = "flash_attn",
+        stochastic_p: float = 0.5,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        if num_kv_heads is None:
+            self.num_kv_heads = self.num_heads
+        else:
+            self.num_kv_heads = num_kv_heads
+        self.num_kv_groups = num_heads // self.num_kv_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.qkv_bias = qkv_bias
+        self.qk_norm = qk_norm
+
+        self.window_size = window_size
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_idx = layer_idx
+        self.attn_impl = attn_impl
+        self.stochastic_value = stochastic_p
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+
+        if "scaled" in self.attn_impl:
+            self.s = nn.Parameter(torch.empty(self.num_heads, 1))
+            self.register_buffer("logn", torch.log(torch.arange(2, self.max_position_embeddings*4+2, dtype=self.s.dtype)[:, None, None]))
+
+        if qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta)
+
+    def reset_parameters(self):
+        if "scaled" in self.attn_impl:
+            nn.init.constant_(self.s, 0.3)
+            self.logn.copy_(torch.log(torch.arange(2, self.max_position_embeddings*4+2, dtype=self.s.dtype)[:, None, None]))
+
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+
+        q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim)
+        k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+
+        if self.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+
+        # equivalent to cu_seqlens in `flash_attn`
+        cu_seqlens = kwargs.get('cu_seqlens', None)
+
+        seqlen_offset, max_seqlen = 0, q_len
+        if past_key_values is not None:
+            seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            max_seqlen = q.shape[1] + seqlen_offset
+
+            if attention_mask is not None:
+                # to deliminate the offsets of padding tokens
+                seqlen_offset = seqlen_offset + attention_mask.sum(-1) - attention_mask.shape[-1]
+                max_seqlen = q.shape[1] + max(seqlen_offset)
+
+        if self.max_position_embeddings is not None:
+            max_seqlen = max(max_seqlen, self.max_position_embeddings)
+        q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens)
+
+        if past_key_values is not None:
+            cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0
+            k_cached, v_cached = past_key_values.update(
+                attn_state=(k.flatten(-2, -1), v.flatten(-2, -1)),
+                layer_idx=self.layer_idx,
+                offset=q_len,
+                cache_kwargs=dict(window_size=self.window_size)
+            )['attn_state']
+            if cache_has_content:
+                k, v = k_cached, v_cached
+                k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+                v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+
+        # if flash_attn_func is None:
+        #     raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first")
+
+        if "scaled" in self.attn_impl:
+            k_len = k.shape[1]
+            q = q * self.s.to(q.dtype) * self.logn[k_len-q_len:k_len].to(q.dtype)
+
+        # Contains at least one padding token in the sequence
+
+        p = torch.rand(1, device=q.device)
+        stochastic_p = torch.tensor(self.stochastic_value, dtype=torch.float32, device=q.device)
+        cond = torch.where(p < stochastic_p, torch.tensor(1, dtype=torch.bool, device=q.device), torch.tensor(0, dtype=torch.bool, device=q.device))
+        if self.attn_impl == "flash_attn":
+            if attention_mask is not None:
+                q, k, v, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(q, k, v, attention_mask, q_len)
+                cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+                max_seqlen_q, max_seqlen_k = max_seq_lens
+                o = flash_attn_varlen_func(
+                    q, k, v,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_k=max_seqlen_k,
+                    causal=True,
+                    window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0)
+                )
+                o = pad_input(o, indices_q, batch_size, q_len)
+            elif cu_seqlens is not None:
+                o = flash_attn_varlen_func(
+                    q.squeeze(0), k.squeeze(0), v.squeeze(0),
+                    cu_seqlens_q=cu_seqlens,
+                    cu_seqlens_k=cu_seqlens,
+                    max_seqlen_q=max_seqlen,
+                    max_seqlen_k=max_seqlen,
+                    causal=True,
+                    window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0)
+                ).unsqueeze(0)
+            else:
+                o = flash_attn_func(
+                    q, k, v,
+                    causal=True,
+                    window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0)
+                )
+        
+        elif self.attn_impl == "parallel_attn":
+            if cond:
+                o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+            else:
+                o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "parallel_scaled_attn":
+            if cond:
+                o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+            else:
+                o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "parallel_rectified_attn":
+            if cond:
+                o = parallel_rectified_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+            else:
+                o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "parallel_softpick_attn":
+            if cond:
+                o = parallel_softpick_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+            else:
+                o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "parallel_scaled_softpick_attn":
+            if cond:
+                o = parallel_softpick_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+            else:
+                o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "naive_attn":
+            if cond:
+                o, attentions = naive_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+            else:
+                o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "naive_scaled_attn":
+            if cond:
+                o, attentions = naive_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+            else:
+                o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "naive_rectified_attn":
+            if cond:
+                o, attentions = naive_rectified_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+            else:
+                o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "naive_softpick_attn":
+            if cond:
+                o, attentions = naive_softpick_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+            else:
+                o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        elif self.attn_impl == "naive_scaled_softpick_attn":
+            if cond:
+                o, attentions = naive_softpick_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+            else:
+                o = parallel_attn(q, k, v, scale=self.head_dim**-0.5, cu_seqlens=cu_seqlens)
+        else:
+            raise ValueError(f"Unknown attention implementation: {self.attn_impl}")
+
+        o = o.reshape(batch_size, q_len, -1)
+        o = self.o_proj(o)
+
+        if not output_attentions or "parallel" in self.attn_impl or "flash" in self.attn_impl:
+            attentions = None
+
+        return o, attentions, past_key_values
+
+    def _upad_input(self, q, k, v, attention_mask, q_len):
+        batch_size, seq_len, num_key_value_heads, head_dim = k.shape
+        cache_mask = attention_mask[:, -seq_len:]
+        seqlens = cache_mask.sum(-1, dtype=torch.int32)
+        indices_k = torch.nonzero(cache_mask.flatten(), as_tuple=False).flatten()
+        max_seqlen_k = seqlens.max().item()
+        cu_seqlens_k = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))
+
+        k = index_first_axis(k.reshape(batch_size * seq_len, num_key_value_heads, head_dim), indices_k)
+        v = index_first_axis(v.reshape(batch_size * seq_len, num_key_value_heads, head_dim), indices_k)
+        if q_len == seq_len:
+            q = index_first_axis(q.reshape(batch_size * seq_len, self.num_heads, head_dim), indices_k)
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_q = max_seqlen_k
+            indices_q = indices_k
+        elif q_len == 1:
+            max_seqlen_q = 1
+            # There is a memcpy here, that is very bad.
+            cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=q.device)
+            indices_q = cu_seqlens_q[:-1]
+            q = q.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -q_len:]
+            q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, attention_mask)
+
+        return q, k, v, indices_q, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k)

fla/layers/based.py

@@ -0,0 +1,96 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+"""
+Linear attention in Based.
+https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py
+"""
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from fla.modules.feature_map import TaylorFeatureMap
+from fla.ops.based import parallel_based
+from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn
+
+
+class BasedLinearAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        feature_dim: int = 16,
+        num_key_value_heads: int = 12,
+        num_heads: int = 12,
+        feature_name: str = "taylor_exp",
+        eps: float = 1e-12,
+        causal: bool = True,
+        mode: str = "parallel",
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.mode = mode
+        self.feature_name = feature_name
+        self.feature_dim = feature_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.num_heads = num_heads
+        self.head_dim = self.hidden_size // self.num_key_value_heads
+        assert self.hidden_size % self.head_dim == 0
+        self.causal = causal
+
+        self.q_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+        self.dropout = nn.Identity()
+        self.feature_map = TaylorFeatureMap(feature_dim)
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor, **kwargs):
+        mode = self.mode
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+        q, k, v = map(lambda x: rearrange(x, "... (h d) -> ... h d", d=self.head_dim), [q, k, v])
+        if mode == "fused_chunk":
+            q, k = self.feature_map(q), self.feature_map(k)
+            o, _ = fused_chunk_linear_attn(q, k, v, normalize=True, scale=1, head_first=False)
+        elif mode == 'chunk':
+            q, k = self.feature_map(q), self.feature_map(k)
+            o, _ = chunk_linear_attn(q, k, v, normalize=True, scale=1, head_first=False)
+        elif mode == 'parallel':
+            assert q.shape[-1] <= 128
+            o = parallel_based(q, k, v, scale=1, use_norm=True, head_first=False)
+        o = rearrange(o, 'b t h d -> b t (h d)')
+        o = self.o_proj(o)
+        o = self.dropout(o)
+        return o
+
+    # https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py#L119
+
+    def forward_reference(self, hidden_states: torch.Tensor, filters: torch.Tensor = None, *args, **kwargs):
+        """
+        x (torch.Tensor): tensor of shape (b, d, t)
+        y (torch.Tensor): tensor of shape (b, d, t)
+        """
+        # hidden_states = hidden_states.transpose(1, 2)
+        b, t, _ = hidden_states.size()
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+
+        q = q.view(b, t, self.num_heads, self.feature_dim).transpose(1, 2)
+        k = k.view(b, t, self.num_key_value_heads, self.feature_dim).transpose(1, 2)
+        v = v.view(b, t, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        # Linear attention
+        q, k = self.feature_map(q), self.feature_map(k)
+        q, k, v = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1)
+
+        # Compute attention
+        if self.causal:
+            y = ((q * (k * v).cumsum(2)).sum(-1) / ((q * k.cumsum(2)).sum(-1) + self.eps))
+        else:
+            y = ((q * (k * v).sum(2, True)).sum(-1) / ((q * k.sum(2, True)).sum(-1) + self.eps))
+        y = rearrange(y, 'b h t d -> b t (h d)')
+        y = self.o_proj(y.to(hidden_states.dtype))
+        y = self.dropout(y)
+        return y.to(hidden_states.dtype)

fla/layers/bitattn.py

@@ -0,0 +1,192 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from einops import rearrange
+from transformers.utils import logging
+
+from fla.modules import RotaryEmbedding
+from fla.modules.fused_bitlinear import FusedBitLinear
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+try:
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
+except ImportError:
+    warnings.warn(
+        "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`",
+        category=ImportWarning
+    )
+    flash_attn_func = None
+
+logger = logging.get_logger(__name__)
+
+
+class BitAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 32,
+        num_kv_heads: Optional[int] = None,
+        window_size: Optional[int] = None,
+        rope_theta: Optional[float] = 10000.,
+        max_position_embeddings: Optional[int] = None,
+        norm_eps: float = 1e-5,
+        layer_idx: int = None
+    ):
+        super().__init__()
+
+        self.num_heads = num_heads
+        if num_kv_heads is None:
+            self.num_kv_heads = self.num_heads
+        else:
+            self.num_kv_heads = num_kv_heads
+        self.num_kv_groups = num_heads // self.num_kv_heads
+        self.hidden_size = hidden_size
+        self.head_dim = self.hidden_size // self.num_heads
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.window_size = window_size
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_idx = layer_idx
+
+        self.q_proj = FusedBitLinear(self.hidden_size, self.hidden_size, bias=False)
+        self.k_proj = FusedBitLinear(self.hidden_size, self.kv_dim, bias=False)
+        self.v_proj = FusedBitLinear(self.hidden_size, self.kv_dim, bias=False)
+        self.o_proj = FusedBitLinear(self.hidden_size, self.hidden_size, bias=False)
+
+        self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+
+        # equivalent to cu_seqlens in `flash_attn`
+        cu_seqlens = kwargs.get('cu_seqlens', None)
+
+        seqlen_offset, max_seqlen = 0, q_len
+        if past_key_values is not None:
+            seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            max_seqlen = q.shape[1] + seqlen_offset
+
+            if attention_mask is not None:
+                # to deliminate the offsets of padding tokens
+                seqlen_offset = seqlen_offset + attention_mask.sum(-1) - attention_mask.shape[-1]
+                max_seqlen = q.shape[1] + max(seqlen_offset)
+
+        if self.max_position_embeddings is not None:
+            max_seqlen = max(max_seqlen, self.max_position_embeddings)
+        q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens)
+
+        if past_key_values is not None:
+            cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0
+            k_cached, v_cached = past_key_values.update(
+                attn_state=(k.flatten(-2, -1), v.flatten(-2, -1)),
+                layer_idx=self.layer_idx,
+                offset=q_len,
+                cache_kwargs=dict(window_size=self.window_size)
+            )['attn_state']
+            if cache_has_content:
+                k, v = k_cached, v_cached
+                k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+                v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+
+        if flash_attn_func is None:
+            raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first")
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            q, k, v, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(q, k, v, attention_mask, q_len)
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_q, max_seqlen_k = max_seq_lens
+            o = flash_attn_varlen_func(
+                q, k, v,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0)
+            )
+            o = pad_input(o, indices_q, batch_size, q_len)
+        elif cu_seqlens is not None:
+            o = flash_attn_varlen_func(
+                q.squeeze(0), k.squeeze(0), v.squeeze(0),
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0)
+            ).unsqueeze(0)
+        else:
+            o = flash_attn_func(
+                q, k, v,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0)
+            )
+        o = o.reshape(batch_size, q_len, -1)
+        o = self.o_proj(o)
+
+        if not output_attentions:
+            attentions = None
+
+        return o, attentions, past_key_values
+
+    def _upad_input(self, q, k, v, attention_mask, q_len):
+        batch_size, seq_len, num_key_value_heads, head_dim = k.shape
+        cache_mask = attention_mask[:, -seq_len:]
+        seqlens = cache_mask.sum(-1, dtype=torch.int32)
+        indices_k = torch.nonzero(cache_mask.flatten(), as_tuple=False).flatten()
+        max_seqlen_k = seqlens.max().item()
+        cu_seqlens_k = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))
+
+        k = index_first_axis(k.reshape(batch_size * seq_len, num_key_value_heads, head_dim), indices_k)
+        v = index_first_axis(v.reshape(batch_size * seq_len, num_key_value_heads, head_dim), indices_k)
+        if q_len == seq_len:
+            q = index_first_axis(q.reshape(batch_size * seq_len, self.num_heads, head_dim), indices_k)
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_q = max_seqlen_k
+            indices_q = indices_k
+        elif q_len == 1:
+            max_seqlen_q = 1
+            # There is a memcpy here, that is very bad.
+            cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=q.device)
+            indices_q = cu_seqlens_q[:-1]
+            q = q.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -q_len:]
+            q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, attention_mask)
+
+        return q, k, v, indices_q, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k)

fla/layers/forgetting_attn.py

@@ -0,0 +1,109 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from einops import rearrange
+from transformers.utils import logging
+
+from fla.modules import GroupNorm
+from fla.ops.forgetting_attn.parallel import parallel_forgetting_attn
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+
+logger = logging.get_logger(__name__)
+
+
+class ForgettingAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 32,
+        num_kv_heads: Optional[int] = None,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        window_size: Optional[int] = None,
+        use_output_gate: bool = False,
+        layer_idx: int = None
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        if num_kv_heads is None:
+            self.num_kv_heads = self.num_heads
+        else:
+            self.num_kv_heads = num_kv_heads
+        self.num_kv_groups = num_heads // self.num_kv_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.qkv_bias = qkv_bias
+        self.qk_norm = qk_norm
+
+        self.window_size = window_size
+        self.use_output_gate = use_output_gate
+        self.layer_idx = layer_idx
+
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.f_proj = nn.Linear(self.hidden_size, self.num_heads, bias=True)
+
+        if use_output_gate:
+            self.g_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+
+        if qk_norm:
+            self.q_norm = GroupNorm(
+                num_groups=self.num_heads,
+                hidden_size=self.hidden_size,
+                is_rms_norm=True,
+            )
+            self.k_norm = GroupNorm(
+                num_groups=self.num_kv_heads,
+                hidden_size=self.kv_dim,
+                is_rms_norm=True,
+            )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        cu_seqlens = kwargs.get('cu_seqlens', None)
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+        f = F.logsigmoid(self.f_proj(hidden_states).float())
+        if self.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+
+        q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim)
+        k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+
+        o = parallel_forgetting_attn(q, k, v, f, cu_seqlens=cu_seqlens)
+        o = rearrange(o, '... h d -> ... (h d)')
+        if self.use_output_gate:
+            o = self.g_proj(hidden_states).sigmoid() * o
+        o = self.o_proj(o)
+
+        return o, None, past_key_values

fla/layers/gated_deltaproduct.py

@@ -0,0 +1,351 @@
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from fla.modules import FusedRMSNormSwishGate, RMSNorm, ShortConvolution
+from fla.ops.delta_rule import chunk_delta_rule
+from fla.ops.gated_delta_rule import chunk_gated_delta_rule
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+def elu_p1(x):
+    return (F.elu(x, 1.0, False) + 1.0).to(x)
+
+
+def sum_norm(x):
+    return (x / x.sum(-1, keepdim=True)).to(x)
+
+
+def interleave_multiple_sequences(*sequences):
+    """
+    Interleave multiple sequences together.
+    For example, with sequences [A1, A2], [B1, B2], [C1, C2],
+    returns [A1, B1, C1, A2, B2, C2]
+    """
+    if isinstance(sequences[0], (list, tuple)):
+        sequences = sequences[0]
+
+    if len(sequences) == 1:
+        return sequences[0]
+
+    # All sequences should have the same shape
+    assert all(s.shape == sequences[0].shape for s in sequences)
+
+    # Get the original shape
+    batch_size, seq_len, *rest = sequences[0].shape
+
+    # Stack sequences along a new dimension
+    stacked = torch.stack(sequences, dim=2)
+
+    # Reshape to interleave
+    reshaped = stacked.view(batch_size, seq_len * len(sequences), *rest)
+
+    return reshaped
+
+
+class GatedDeltaProduct(nn.Module):
+    """
+    Generalized version of GatedDoubleDeltaNet that supports arbitrary number of householder transformations.
+    """
+
+    def __init__(
+            self,
+            hidden_size: int = 2048,
+            expand_v: float = 2,
+            head_dim: int = 256,
+            num_heads: int = 6,
+            num_householder: int = 2,  # New parameter for number of householder transformations
+            mode: str = "chunk",
+            use_gate: bool = True,
+            use_forget_gate: bool = True,  # when true Gated DeltaProduct, when false DeltaProduct
+            use_short_conv: bool = True,
+            conv_size: int = 4,
+            conv_bias: bool = False,
+            layer_idx: int | None = None,
+            norm_eps: float = 1e-5,
+            allow_neg_eigval: bool = False,  # when true (Gated) DeltaProduct [-1, 1], when false (Gated) DeltaProduct [0, 1]
+            **kwargs,
+    ) -> None:
+        super().__init__()
+
+        self.mode = mode
+        self.hidden_size = hidden_size
+        self.expand_v = expand_v
+        self.use_gate = use_gate
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.num_householder = num_householder
+        self.allow_neg_eigval = allow_neg_eigval
+        self.use_forget_gate = use_forget_gate
+        self.key_dim = self.num_heads * self.head_dim
+        self.value_dim = int(self.key_dim * self.expand_v)
+        self.head_qk_dim = head_dim
+        self.head_v_dim = int(head_dim * self.expand_v)
+        self.layer_idx = layer_idx
+        self.silu = nn.SiLU()
+        assert mode in ["chunk", "fused_recurrent"], f"Not supported mode `{mode}`."
+        # Create multiple projection layers for each householder transformation
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+
+        self.k_projs = nn.ModuleList(
+            [
+                nn.Linear(hidden_size, self.key_dim, bias=False)
+                for _ in range(num_householder)
+            ]
+        )
+        self.v_projs = nn.ModuleList(
+            [
+                nn.Linear(hidden_size, self.value_dim, bias=False)
+                for _ in range(num_householder)
+            ]
+        )
+        self.b_projs = nn.ModuleList(
+            [
+                nn.Linear(hidden_size, self.num_heads, bias=False)
+                for _ in range(num_householder)
+            ]
+        )
+        if use_short_conv:
+            self.q_conv1ds = nn.ModuleList(
+                [
+                    ShortConvolution(
+                        hidden_size=self.key_dim,
+                        kernel_size=conv_size,
+                        activation="silu",
+                    )
+                    for _ in range(num_householder)
+                ]
+            )
+            self.k_conv1ds = nn.ModuleList(
+                [
+                    ShortConvolution(
+                        hidden_size=self.key_dim,
+                        kernel_size=conv_size,
+                        activation="silu",
+                    )
+                    for _ in range(num_householder)
+                ]
+            )
+            self.v_conv1ds = nn.ModuleList(
+                [
+                    ShortConvolution(
+                        hidden_size=self.value_dim,
+                        kernel_size=conv_size,
+                        activation="silu",
+                    )
+                    for _ in range(num_householder)
+                ]
+            )
+
+        if self.use_forget_gate:
+            self.a_proj = nn.Linear(hidden_size, self.num_heads, bias=False)
+            A = torch.empty(self.num_heads, dtype=torch.float32).uniform_(0, 16)
+            A_log = torch.log(A)
+            self.A_log = nn.Parameter(A_log)
+            self.A_log._no_weight_decay = True
+
+            # Initialize dt parameters
+            dt_min = 0.001
+            dt_max = 0.1
+            dt_init_floor = 1e-4
+            dt = torch.exp(
+                torch.rand(self.num_heads) * (math.log(dt_max) - math.log(dt_min))
+                + math.log(dt_min)
+            )
+            dt = torch.clamp(dt, min=dt_init_floor)
+            inv_dt = dt + torch.log(-torch.expm1(-dt))
+            self.dt_bias = nn.Parameter(inv_dt)
+            self.dt_bias._no_weight_decay = True
+
+        if use_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+            self.o_norm = FusedRMSNormSwishGate(self.head_v_dim, eps=norm_eps)
+        else:
+            self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+        self.k_id = torch.nn.Identity()
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            past_key_values: Optional[Cache] = None,
+            use_cache: Optional[bool] = False,
+            output_attentions: Optional[bool] = False,
+            **kwargs: Unpack[Dict],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding)."
+            )
+
+        mode = (
+            "chunk"  # 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+        )
+        if self.training:
+            assert mode == "chunk", "Only chunk mode is supported in training."
+
+        last_state = None
+        if past_key_values is not None and len(past_key_values) > self.layer_idx:
+            last_state = past_key_values[self.layer_idx]
+
+        # Process each householder transformation
+        ks, vs, betas = [], [], []
+        conv_states = []
+
+        for i in range(self.num_householder):
+            if self.use_short_conv:
+                conv_state_q, conv_state_k, conv_state_v = None, None, None
+                if last_state is not None:
+                    conv_state_q, conv_state_k, conv_state_v = last_state["conv_state"][
+                        i
+                    ]
+                conv_mask = (
+                    attention_mask[:, -hidden_states.shape[1]:]
+                    if attention_mask is not None
+                    else None
+                )
+
+                k, conv_state_k = self.k_conv1ds[i](
+                    x=self.k_projs[i](hidden_states),
+                    mask=conv_mask,
+                    cache=conv_state_k,
+                    output_final_state=use_cache,
+                )
+                v, conv_state_v = self.v_conv1ds[i](
+                    x=self.v_projs[i](hidden_states),
+                    mask=conv_mask,
+                    cache=conv_state_v,
+                    output_final_state=use_cache,
+                )
+                conv_states.append((conv_state_q, conv_state_k, conv_state_v))
+            else:
+                k = self.silu(self.k_projs[i](hidden_states))
+                v = self.silu(self.v_projs[i](hidden_states))
+
+            ks.append(k)
+            vs.append(v)
+
+            beta = self.b_projs[i](
+                hidden_states
+            ).sigmoid()  # bs, sequence_length, num_heads
+            if attention_mask is not None:
+                beta = beta.mul(attention_mask[:, -hidden_states.shape[1]:, None])
+            if self.allow_neg_eigval:
+                beta = beta * 2
+            betas.append(beta)
+
+        if self.use_short_conv:
+            q, conv_state_q = self.q_conv1ds[0](
+                x=self.q_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_q,
+                output_final_state=use_cache,
+            )
+        else:
+            q = self.silu(self.q_proj(hidden_states))
+        q = interleave_multiple_sequences(
+            [torch.zeros_like(q)] * (self.num_householder - 1) + [q]
+        )
+        # Interleave all sequences
+        k = interleave_multiple_sequences(ks)
+        v = interleave_multiple_sequences(vs)
+        beta = interleave_multiple_sequences(betas)
+
+        q, k, v = (
+            rearrange(x, "b t (h d) -> b t h d", h=self.num_heads) for x in (q, k, v)
+        )
+
+        recurrent_state = (
+            last_state["recurrent_state"] if last_state is not None else None
+        )
+        offsets = kwargs.get("offsets")
+
+        if mode == "chunk":
+            if self.use_forget_gate:
+                g = -self.A_log.float().exp() * F.softplus(
+                    self.a_proj(hidden_states).float() + self.dt_bias
+                )
+                if attention_mask is not None:
+                    g = g.mul(attention_mask[:, -g.shape[-2]:, None])
+
+                # Interleave g with zeros for non-first transformations
+                g = interleave_multiple_sequences(
+                    [g] + [torch.zeros_like(g)] * (self.num_householder - 1)
+                )
+
+                o, recurrent_state = chunk_gated_delta_rule(
+                    q=q,
+                    k=k,
+                    v=v,
+                    g=g,
+                    beta=beta,
+                    initial_state=recurrent_state,
+                    output_final_state=use_cache,
+                    cu_seqlens=offsets,
+                    head_first=False,
+                    use_qk_l2norm_in_kernel=True
+                )
+            else:
+                o, recurrent_state = chunk_delta_rule(
+                    q=q,
+                    k=k,
+                    v=v,
+                    beta=beta,
+                    initial_state=recurrent_state,
+                    output_final_state=use_cache,
+                    cu_seqlens=offsets,
+                    head_first=False,
+                    use_qk_l2norm_in_kernel=True
+                )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        # Take every nth element for n householder transformations
+        o = o[:, self.num_householder - 1:: self.num_householder, :]
+
+        if past_key_values is not None:
+            past_key_values.update(
+                recurrent_state=recurrent_state,
+                conv_state=conv_states if self.use_short_conv else None,
+                layer_idx=self.layer_idx,
+                offset=q.shape[2],
+            )
+
+        if self.use_gate:
+            g = rearrange(
+                self.g_proj(hidden_states),
+                "... (h d) -> ... h d",
+                h=self.num_heads,
+            )
+            o = self.o_norm(o, g)
+        else:
+            o = self.o_norm(o)
+        o = rearrange(o, "b t h d -> b t (h d)")
+        o = self.o_proj(o)
+
+        return o, None, past_key_values

fla/ops/utils/cumsum.py

@@ -0,0 +1,400 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from fla.utils import check_shared_mem, input_guard
+
+BS_LIST = [32, 64] if check_shared_mem() else [16, 32]
+
+
+@triton.heuristics({
+    'USE_OFFSETS': lambda args: args['offsets'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [1, 2, 4, 8]
+    ],
+    key=['BT']
+)
+@triton.jit(do_not_specialize=['T'])
+def chunk_local_cumsum_scalar_kernel(
+    s,
+    o,
+    offsets,
+    indices,
+    T,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    HEAD_FIRST: tl.constexpr,
+    USE_OFFSETS: tl.constexpr,
+    REVERSE: tl.constexpr
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_h = i_bh // H, i_bh % H
+    if USE_OFFSETS:
+        i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32)
+        bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    if HEAD_FIRST:
+        p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+        p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    else:
+        p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
+        p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
+    # [BT]
+    b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32)
+    b_o = tl.cumsum(b_s, axis=0)
+    if REVERSE:
+        b_z = tl.sum(b_s, axis=0)
+        b_o = -b_o + b_z[None] + b_s
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,))
+
+
+@triton.heuristics({
+    'USE_OFFSETS': lambda args: args['offsets'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BS': BS}, num_warps=num_warps)
+        for BS in BS_LIST
+        for num_warps in [2, 4, 8]
+    ],
+    key=['S', 'BT'],
+)
+@triton.jit(do_not_specialize=['T'])
+def chunk_local_cumsum_vector_kernel(
+    s,
+    o,
+    offsets,
+    indices,
+    T,
+    H: tl.constexpr,
+    S: tl.constexpr,
+    BT: tl.constexpr,
+    BS: tl.constexpr,
+    HEAD_FIRST: tl.constexpr,
+    USE_OFFSETS: tl.constexpr,
+    REVERSE: tl.constexpr
+):
+    i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_b, i_h = i_bh // H, i_bh % H
+    if USE_OFFSETS:
+        i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32)
+        bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    o_i = tl.arange(0, BT)
+    if REVERSE:
+        m_s = tl.where(o_i[:, None] <= o_i[None, :], 1., 0.)
+    else:
+        m_s = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.)
+
+    if HEAD_FIRST:
+        p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+        p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+    else:
+        p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+        p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+    # [BT, BS]
+    b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)
+    b_o = tl.dot(m_s, b_s, allow_tf32=False)
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.heuristics({
+    'USE_OFFSETS': lambda args: args['offsets'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BT': 16}, num_warps=2),
+        triton.Config({'BT': 32}, num_warps=4),
+        triton.Config({'BT': 32}, num_warps=2),
+        triton.Config({'BT': 64}, num_warps=8),
+        triton.Config({'BT': 64}, num_warps=4),
+    ],
+    key=[]
+)
+@triton.jit(do_not_specialize=['T'])
+def chunk_global_cumsum_scalar_kernel(
+    s,
+    o,
+    offsets,
+    T,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    HEAD_FIRST: tl.constexpr,
+    USE_OFFSETS: tl.constexpr,
+    REVERSE: tl.constexpr
+):
+    i_bh = tl.program_id(0)
+    i_b, i_h = i_bh // H, i_bh % H
+    if USE_OFFSETS:
+        bos, eos = tl.load(offsets + i_b).to(tl.int32), tl.load(offsets + i_b + 1).to(tl.int32)
+    else:
+        bos, eos = i_b * T, i_b * T + T
+    T = eos - bos
+
+    b_z = tl.zeros([], dtype=tl.float32)
+    NT = tl.cdiv(T, BT)
+    for i_c in range(NT):
+        i_t = NT-1-i_c if REVERSE else i_c
+        if HEAD_FIRST:
+            p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+            p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+        else:
+            p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
+            p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
+        b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32)
+        b_o = tl.cumsum(b_s, axis=0)
+        b_ss = tl.sum(b_s, 0)
+        if REVERSE:
+            b_o = -b_o + b_ss + b_s
+        b_o += b_z
+        if i_c >= 0:
+            b_z += b_ss
+        tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,))
+
+
+@triton.heuristics({
+    'USE_OFFSETS': lambda args: args['offsets'] is not None,
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BT': BT}, num_warps=num_warps)
+        for BT in [16, 32, 64]
+        for num_warps in [2, 4, 8]
+    ],
+    key=['S']
+)
+@triton.jit(do_not_specialize=['T'])
+def chunk_global_cumsum_vector_kernel(
+    s,
+    z,
+    offsets,
+    T,
+    H: tl.constexpr,
+    S: tl.constexpr,
+    BT: tl.constexpr,
+    BS: tl.constexpr,
+    HEAD_FIRST: tl.constexpr,
+    USE_OFFSETS: tl.constexpr,
+    REVERSE: tl.constexpr
+):
+    i_s, i_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_h = i_bh // H, i_bh % H
+    if USE_OFFSETS:
+        bos, eos = tl.load(offsets + i_b).to(tl.int32), tl.load(offsets + i_b + 1).to(tl.int32)
+    else:
+        bos, eos = i_b * T, i_b * T + T
+    T = eos - bos
+
+    o_i = tl.arange(0, BT)
+    if REVERSE:
+        m_s = tl.where(o_i[:, None] <= o_i[None, :], 1., 0.)
+    else:
+        m_s = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.)
+
+    b_z = tl.zeros([BS], dtype=tl.float32)
+    NT = tl.cdiv(T, BT)
+    for i_c in range(NT):
+        i_t = NT-1-i_c if REVERSE else i_c
+        if HEAD_FIRST:
+            p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+            p_z = tl.make_block_ptr(z + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+        else:
+            p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+            p_z = tl.make_block_ptr(z + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+        # [BT, BS]
+        b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)
+        b_c = b_z[None, :] + tl.dot(m_s, b_s, allow_tf32=False)
+        tl.store(p_z, b_c.to(p_z.dtype.element_ty), boundary_check=(0, 1))
+        if i_c >= 0:
+            b_z += tl.sum(b_s, 0)
+
+
+def chunk_local_cumsum_scalar(
+    g: torch.Tensor,
+    chunk_size: int,
+    reverse: bool = False,
+    offsets: Optional[torch.Tensor] = None,
+    indices: Optional[torch.Tensor] = None,
+    head_first: bool = True,
+    output_dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    if head_first:
+        B, H, T = g.shape
+    else:
+        B, T, H = g.shape
+    if offsets is not None:
+        B = len(offsets) - 1
+    assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2"
+    BT = chunk_size
+    NT = triton.cdiv(T, BT) if offsets is None else len(indices)
+    g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
+    grid = (NT, B * H)
+    chunk_local_cumsum_scalar_kernel[grid](
+        g_org,
+        g,
+        offsets,
+        indices,
+        T=T,
+        H=H,
+        BT=BT,
+        HEAD_FIRST=head_first,
+        REVERSE=reverse
+    )
+    return g
+
+
+def chunk_local_cumsum_vector(
+    g: torch.Tensor,
+    chunk_size: int,
+    reverse: bool = False,
+    offsets: Optional[torch.Tensor] = None,
+    indices: Optional[torch.Tensor] = None,
+    head_first: bool = True,
+    output_dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    if head_first:
+        B, H, T, S = g.shape
+    else:
+        B, T, H, S = g.shape
+    BT = chunk_size
+    NT = triton.cdiv(T, BT) if offsets is None else len(indices)
+    assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2"
+
+    g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
+    def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H)
+    # keep cummulative normalizer in fp32
+    # this kernel is equivalent to
+    # g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1)
+    chunk_local_cumsum_vector_kernel[grid](
+        g_org,
+        g,
+        offsets,
+        indices,
+        T=T,
+        H=H,
+        S=S,
+        BT=BT,
+        HEAD_FIRST=head_first,
+        REVERSE=reverse
+    )
+    return g
+
+
+@input_guard
+def chunk_global_cumsum_scalar(
+    s: torch.Tensor,
+    dtype: Optional[torch.dtype] = None,
+    reverse: bool = False,
+    offsets: Optional[torch.Tensor] = None,
+    head_first: bool = True,
+    output_dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    dtype = dtype or s.dtype
+    if head_first:
+        B, H, T = s.shape
+    else:
+        B, T, H = s.shape
+    if offsets is not None:
+        B = len(offsets) - 1
+    grid = (B * H,)
+    z = torch.empty_like(s, dtype=output_dtype or dtype)
+    chunk_global_cumsum_scalar_kernel[grid](
+        s,
+        z,
+        offsets,
+        T=T,
+        H=H,
+        HEAD_FIRST=head_first,
+        REVERSE=reverse
+    )
+    return z
+
+
+@input_guard
+def chunk_global_cumsum_vector(
+    s: torch.Tensor,
+    dtype: Optional[torch.dtype] = None,
+    reverse: bool = False,
+    offsets: Optional[torch.Tensor] = None,
+    head_first: bool = True,
+    output_dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    dtype = dtype or s.dtype
+    if head_first:
+        B, H, T, S = s.shape
+    else:
+        B, T, H, S = s.shape
+    BS = min(32, triton.next_power_of_2(S))
+    if offsets is not None:
+        B = len(offsets) - 1
+    grid = (triton.cdiv(S, BS), B * H)
+    z = torch.empty_like(s, dtype=output_dtype or dtype)
+    chunk_global_cumsum_vector_kernel[grid](
+        s,
+        z,
+        offsets,
+        T=T,
+        H=H,
+        S=S,
+        BS=BS,
+        HEAD_FIRST=head_first,
+        REVERSE=reverse
+    )
+    return z
+
+
+@input_guard
+def chunk_global_cumsum(
+    s: torch.Tensor,
+    dtype: Optional[torch.dtype] = None,
+    reverse: bool = False,
+    offsets: Optional[torch.Tensor] = None,
+    head_first: bool = True,
+    output_dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    if offsets is not None:
+        assert s.shape[0] == 1, "Only batch size 1 is supported when offsets are provided"
+    if len(s.shape) == 3:
+        return chunk_global_cumsum_scalar(s, dtype, reverse, offsets, head_first, output_dtype)
+    elif len(s.shape) == 4:
+        return chunk_global_cumsum_vector(s, dtype, reverse, offsets, head_first, output_dtype)
+    else:
+        raise ValueError(f"Unsupported input shape {s.shape}. "
+                         f"which should be [B, H, T]/[B, H, T, D] if `head_first=True` "
+                         f"or [B, T, H]/[B, T, H, D] otherwise")
+
+
+@input_guard
+def chunk_local_cumsum(
+    g: torch.Tensor,
+    chunk_size: int,
+    reverse: bool = False,
+    offsets: Optional[torch.Tensor] = None,
+    indices: Optional[torch.Tensor] = None,
+    head_first: bool = True,
+    output_dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    if offsets is not None:
+        assert g.shape[0] == 1, "Only batch size 1 is supported when offsets are provided"
+    if len(g.shape) == 3:
+        return chunk_local_cumsum_scalar(g, chunk_size, reverse, offsets, indices, head_first, output_dtype)
+    elif len(g.shape) == 4:
+        return chunk_local_cumsum_vector(g, chunk_size, reverse, offsets, indices, head_first, output_dtype)
+    else:
+        raise ValueError(f"Unsupported input shape {g.shape}. "
+                         f"which should be (B, H, T, dim) if `head_first=True` "
+                         f"or (batch_size, num_heads, seq_len) otherwise")

flame/config_manager.py

@@ -0,0 +1,940 @@
+# 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.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 FSDP.
+                This feature only takes effect when data_parallel_shard_degree > 1
+            """,
+        )
+        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",
+        )
+        # 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(
+            "--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.interval",
+            type=int,
+            default=500,
+            help="Checkpointing interval in steps.",
+        )
+        self.parser.add_argument(
+            "--checkpoint.model_weights_only",
+            action="store_true",
+            help="""
+                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.
+            """,
+        )
+        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.
+            """,
+        )
+        self.parser.add_argument(
+            "--checkpoint.convert_to_hf_on_save",
+            action="store_true",
+            help="""
+                If true, automatically convert the saved DCP checkpoint to Hugging Face format
+                in a parallel directory (e.g., step-1000-hf) after each save.
+            """,
+        )
+        self.parser.add_argument(
+            "--checkpoint.hf_upload_enabled",
+            action="store_true",
+            help="Enable uploading converted Hugging Face checkpoints to the Hub.",
+        )
+        self.parser.add_argument(
+            "--checkpoint.hf_repo_base_name",
+            type=str,
+            default=None,
+            help="Hugging Face Hub repository ID to upload checkpoints to (e.g., 'username/repo').",
+        )
+        self.parser.add_argument(
+            "--checkpoint.hf_upload_format",
+            type=str,
+            default="dcp",
+            choices=["dcp", "hf"],
+            help="""
+                Format to upload to Hugging Face Hub. 'dcp' for DCP format, 'hf' for Hugging Face format.
+                Note: 'hf' is only supported for models with a single pipeline stage.
+            """,
+        )
+        # 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

flame/data.py

@@ -0,0 +1,570 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+import copy
+import pickle
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Iterable, List, Optional, Union
+
+import datasets
+import numpy as np
+import torch
+from datasets import Dataset, IterableDataset
+from datasets.iterable_dataset import ShufflingConfig
+from torch.distributed.checkpoint.stateful import Stateful
+from torchdata.stateful_dataloader import StatefulDataLoader
+from transformers import PreTrainedTokenizer
+
+from torchtitan.tools.logging import logger
+
+
+class BufferShuffledIterableDataset(IterableDataset):
+    def __init__(
+        self,
+        dataset: Dataset,
+        tokenizer: PreTrainedTokenizer,
+        seq_len: int = 2048,
+        rank: int = 0,
+        world_size: int = 1,
+        buffer_size: int = 1024,
+    ) -> BufferShuffledIterableDataset:
+        self.dataset = dataset
+        self.tokenizer = tokenizer
+
+        self.data = dataset.shard(world_size, rank)
+        self.seq_len = seq_len
+
+        self.rank = rank
+        self.world_size = world_size
+        self.buffer_size = buffer_size
+
+        if tokenizer.vocab_size < torch.iinfo(torch.int16).max:
+            self.dtype = torch.int16
+        elif tokenizer.vocab_size < torch.iinfo(torch.int32).max:
+            self.dtype = torch.int32
+        else:
+            self.dtype = torch.int64
+        self.states = None
+        self.buffer = torch.tensor([], dtype=self.dtype)
+        self.tokens = []
+        self.rand_id = 0
+        self.token_id = 0
+        self.rng_state = None
+        self._epoch = 0
+
+    def __iter__(self):
+        g = torch.Generator()
+        g.manual_seed(self._epoch + self.rank)
+        if self.rng_state is not None:
+            g.set_state(self.rng_state)
+
+        rand_it = self.randint(0, self.buffer_size, g=g)
+        if self.states is not None:
+            self.data.load_state_dict(self.states)
+
+        # max number of tokens allowed in the chunk buffer
+        n_tokens = self.buffer_size * self.seq_len
+
+        while True:
+            for sample in self.tokenize(self.data):
+                # keep appending the samples to the token buffer
+                self.tokens += sample
+                # if the token buffer is full, start sampling
+                # NOTE: we first convert the token ids to a tensor of shape [n_chunks, seq_len] for efficiency
+                if len(self.buffer) == 0 and len(self.tokens) >= n_tokens:
+                    self.buffer = torch.tensor(self.tokens[:n_tokens], dtype=self.dtype).view(self.buffer_size, -1)
+                    self.tokens = self.tokens[n_tokens:]
+                if len(self.buffer) == self.buffer_size:
+                    yield from self.sample(rand_it)
+
+            n_chunks = len(self.tokens) // self.seq_len
+            # handle the left tokens in the buffer
+            if n_chunks > 0:
+                n_tokens = n_chunks * self.seq_len
+                indices = torch.randperm(n_chunks, generator=g).tolist()
+                self.buffer = torch.tensor(self.tokens[:n_tokens], dtype=torch.long).view(n_chunks, -1)
+                self.tokens = self.tokens[n_tokens:]
+                for i in indices:
+                    yield {'input_ids': self.buffer[i]}
+
+    def tokenize(self, data, batch_size: int = 64):
+        texts, states = [], []
+        for sample in data:
+            texts.append(sample['text'])
+            states.append(self.data.state_dict())
+            if len(texts) == batch_size:
+                for s, tokenized in zip(states, self.tokenizer(texts, return_attention_mask=False)['input_ids']):
+                    self.states = s
+                    yield tokenized
+                texts, states = [], []
+        if len(texts) > 0:
+            for s, tokenized in zip(states, self.tokenizer(texts, return_attention_mask=False)['input_ids']):
+                self.states = s
+                yield tokenized
+
+    def sample(self, indices):
+        n_tokens = (len(self.tokens) // self.seq_len) * self.seq_len
+        while self.token_id < n_tokens:
+            i = next(indices)
+            start, end = self.token_id, self.token_id + self.seq_len
+            self.token_id += self.seq_len
+            yield {'input_ids': self.buffer[i].to(torch.long)}
+            self.buffer[i] = torch.tensor(self.tokens[start:end], dtype=self.dtype)
+        self.token_id = 0
+        self.tokens = self.tokens[n_tokens:]
+
+    def randint(self, low: int, high: int, buffer_size: int = 1024, g: torch.Generator = torch.Generator()) -> Iterable[int]:
+        indices = torch.empty(buffer_size, dtype=torch.long)
+        while True:
+            # record the generator states before sampling
+            self.rng_state = g.get_state()
+            indices = torch.randint(low, high, (buffer_size,), out=indices, generator=g)
+            for i in indices[self.rand_id:].tolist():
+                self.rand_id += 1
+                yield i
+            self.rand_id = 0
+
+    def set_epoch(self, epoch):
+        self._epoch = epoch
+        if hasattr(self.dataset, 'set_epoch'):
+            self.dataset.set_epoch(epoch)
+
+    def state_dict(self):
+        return {
+            'states': self.states,
+            'buffer': self.buffer.clone(),
+            'tokens': deepcopy(self.tokens),
+            'rand_id': self.rand_id,
+            'token_id': self.token_id,
+            'rng_state': self.rng_state,
+            'epoch': self._epoch,
+        }
+
+    def load_state_dict(self, state_dict):
+        self.states = state_dict['states']
+        self.buffer = state_dict['buffer'].clone()
+        self.tokens = deepcopy(state_dict['tokens'])
+        self.rand_id = state_dict['rand_id']
+        self.token_id = state_dict['token_id']
+        self.rng_state = state_dict['rng_state'].clone() if state_dict['rng_state'] is not None else None
+        self._epoch = state_dict['epoch']
+
+
+class OnlineTokenizedIterableDataset(IterableDataset):
+    def __init__(
+        self, dataset: Dataset, tokenizer: PreTrainedTokenizer, seq_len: int = 2048, rank: int = 0, world_size: int = 1
+    ) -> OnlineTokenizedIterableDataset:
+        self.dataset = dataset
+        self.tokenizer = tokenizer
+
+        self.data = dataset.shard(world_size, rank)
+        self.seq_len = seq_len
+        self.rank = rank
+        self.world_size = world_size
+
+        self.states = None
+        self.tokens = []
+
+    def __iter__(self):
+        if self.states is not None:
+            self.data.load_state_dict(self.states)
+
+        while True:
+            for sample in self.tokenize(self.data):
+                # keep appending the samples to the token buffer
+                self.tokens += sample
+
+                while len(self.tokens) >= self.seq_len:
+                    input_ids = torch.tensor(self.tokens[:self.seq_len], dtype=torch.long)
+                    self.tokens = self.tokens[self.seq_len:]
+                    yield {'input_ids': input_ids}
+
+    def tokenize(self, data, buffer_size: int = 64):
+        buffer, states = [], []
+        for sample in data:
+            if sample.get('text', None) is not None:
+                buffer.append(sample['text'])
+            elif sample.get('content', None) is not None:
+                buffer.append(sample['content'])
+            else:
+                raise ValueError(f"No 'text' or 'content' field found in sample:\n{sample}")
+            states.append(self.data.state_dict())
+            if len(buffer) == buffer_size:
+                for s, tokenized in zip(states, self.tokenizer(buffer, return_attention_mask=False)['input_ids']):
+                    self.states = s
+                    yield tokenized
+                buffer, states = [], []
+        if len(buffer) > 0:
+            for s, tokenized in zip(states, self.tokenizer(buffer, return_attention_mask=False)['input_ids']):
+                self.states = s
+                yield tokenized
+
+    def state_dict(self):
+        return {'states': self.states, 'tokens': deepcopy(self.tokens)}
+
+    def load_state_dict(self, state_dict):
+        self.states = state_dict['states']
+        self.tokens = deepcopy(state_dict['tokens'])
+
+
+class BufferShuffledExamplesIterable(datasets.iterable_dataset.BufferShuffledExamplesIterable):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _init_state_dict(self) -> dict:
+        self._state_dict = self.ex_iterable._init_state_dict()
+        self._state_dict['mem_buffer'] = ([],)
+        self._state_dict['bit_generator_state'] = self.generator.bit_generator.state
+        self._state_dict['bit_generator_index_offset'] = 0
+        self._state_dict['bit_generator_index_offset_shuffle'] = 0
+        return self._state_dict
+
+    def __iter__(self):
+        buffer_size = self.buffer_size
+        rng = deepcopy(self.generator)
+        # this is the shuffle buffer that we keep in memory
+        mem_buffer = self._state_dict['mem_buffer'][0]
+        # this is an infinite iterator that randomly samples the index of the source to pick examples from
+        index_offset = self._state_dict['bit_generator_index_offset'] if self._state_dict else 0
+        if self._state_dict:
+            rng.bit_generator.state = self._state_dict['bit_generator_state']
+        indices_iterator = self._iter_random_indices(rng, buffer_size, random_batch_size=buffer_size)
+        # skip already consumed ones
+        for _ in range(index_offset):
+            i = next(indices_iterator)
+
+        for x in self.ex_iterable:
+            if len(mem_buffer) < buffer_size:  # if the buffer is not full, keep filling the buffer
+                mem_buffer.append(x)
+            else:  # otherwise, pick an example from it
+                i = next(indices_iterator)
+                index_offset = (index_offset + 1) % buffer_size
+                if self._state_dict:
+                    self._state_dict['bit_generator_index_offset'] = index_offset
+                    if index_offset == 0:
+                        self._state_dict['bit_generator_state'] = rng.bit_generator.state
+                selected = mem_buffer[i]
+                mem_buffer[i] = x  # replace the picked example by a new one
+                yield selected
+
+        index_offset = self._state_dict['bit_generator_index_offset_shuffle'] if self._state_dict else 0
+        if self._state_dict:
+            rng.bit_generator.state = self._state_dict['bit_generator_state']
+
+        # when we run out of examples, we shuffle the remaining examples in the buffer and yield them
+        for i in rng.permutation(len(mem_buffer))[index_offset:].tolist():
+            index_offset = index_offset + 1
+            if self._state_dict:
+                self._state_dict['bit_generator_index_offset_shuffle'] = index_offset
+            yield mem_buffer[i]
+
+    def shuffle_data_sources(self, generator: np.random.Generator) -> BufferShuffledExamplesIterable:
+        """Shuffle the wrapped examples iterable as well as the shuffling buffer."""
+        return BufferShuffledExamplesIterable(
+            self.ex_iterable.shuffle_data_sources(generator), buffer_size=self.buffer_size, generator=generator
+        )
+
+    def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> BufferShuffledExamplesIterable:
+        """Keep only the requested shard."""
+        return BufferShuffledExamplesIterable(
+            self.ex_iterable.shard_data_sources(num_shards, index, contiguous=contiguous),
+            buffer_size=self.buffer_size,
+            generator=self.generator,
+        )
+
+    def load_state_dict(self, state_dict: dict) -> dict:
+        def _inner_load_state_dict(state, new_state):
+            if new_state is not None and isinstance(state, dict):
+                for key in new_state:
+                    state[key] = _inner_load_state_dict(state[key], new_state[key])
+                return state
+            elif new_state is not None and isinstance(state, list):
+                for i in range(len(state)):
+                    state[i] = _inner_load_state_dict(state[i], new_state[i])
+                return state
+            return new_state
+
+        return _inner_load_state_dict(self._state_dict, state_dict)
+
+
+def shuffle(
+    dataset: IterableDataset,
+    seed: int = 42,
+    generator: np.random.Generator = None,
+    buffer_size: int = 1024,
+):
+    generator = np.random.default_rng(seed) if generator is None else deepcopy(generator)
+    return IterableDataset(
+        ex_iterable=BufferShuffledExamplesIterable(dataset._ex_iterable, buffer_size=buffer_size, generator=generator),
+        info=dataset._info.copy(),
+        split=dataset._split,
+        formatting=dataset._formatting,
+        shuffling=ShufflingConfig(generator=generator, _original_seed=seed),
+        distributed=copy.deepcopy(dataset._distributed),
+        token_per_repo_id=dataset._token_per_repo_id,
+    )
+
+
+@dataclass
+class DataCollatorForLanguageModeling:
+    """
+    Data collator used for language modeling. Inputs are dynamically padded if `varlen=False`.
+    If `varlen=True`, sequences are expected to be concatenated, and labels match inputs.
+
+    Args:
+        tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):
+            The tokenizer used for encoding the data.
+        context_len (`int`, optional):
+            When `varlen=True`, sequences longer than this length within a document
+            (as determined by `cu_seqlens`) will be further chunked.
+        varlen (`bool`):
+            Whether to handle variable length concatenated sequences (`True`) or padded batches (`False`).
+
+    Returns:
+        A dictionary with the following keys:
+        - `input_ids`: Tensor of input IDs. Shape `[batch_size, seq_len]` if `varlen=False`, `[1, total_len]` if `varlen=True`.
+        - `labels`: Tensor of labels. Shape matches `input_ids`. Padding positions are masked with -100 if `varlen=False`.
+        - `attention_mask`: Tensor indicating non-padding tokens (only if `varlen=False`). Shape matches `input_ids`.
+        - `cu_seqlens`: Tensor of cumulative sequence lengths (only if `varlen=True`). Shape `[1, num_sequences + 1]`.
+
+    NOTE: When `varlen=True`, the `batch_size` must be 1.
+    """
+
+    tokenizer: PreTrainedTokenizer
+    context_len: Optional[int] = None
+    varlen: bool = False
+
+    def __call__(self, examples: List[Union[List[int], Dict[str, Any]]]) -> Dict[str, Any]:
+        if not isinstance(examples[0], Dict):
+            examples = [{'input_ids': example} for example in examples]
+
+        def tensorize(example: Dict[str, Any]) -> Dict[str, Any]:
+            tensorized = {}
+            for key in ['input_ids', 'cu_seqlens']:
+                if key not in example:
+                    continue
+                if isinstance(example[key], List):
+                    tensorized[key] = torch.tensor(example[key], dtype=torch.long)
+                elif isinstance(example[key], np.ndarray):
+                    tensorized[key] = torch.from_numpy(example[key])
+                else:
+                    tensorized[key] = example[key]
+            return tensorized
+
+        examples = list(map(tensorize, examples))
+
+        if not self.varlen:
+            # --- Handling for varlen=False (Batch Padding) ---
+            length_of_first = examples[0]['input_ids'].size(0)
+            needs_padding = not all(example['input_ids'].size(0) == length_of_first for example in examples)
+
+            if needs_padding:
+                # Check for pad token if padding is actually required
+                if self.tokenizer.pad_token_id is None:
+                    raise ValueError(
+                        f'You are attempting to pad samples but the tokenizer you are using '
+                        f'({self.tokenizer.__class__.__name__}) does not have a pad token.'
+                    )
+                # Pad using the tokenizer, ensuring attention_mask is returned
+                batch = self.tokenizer.pad(examples, return_tensors='pt', return_attention_mask=True)
+            else:
+                # No padding needed, stack directly and create a full attention mask
+                input_ids = torch.stack([example['input_ids'] for example in examples], dim=0)
+                batch = {
+                    'input_ids': input_ids,
+                    # Create attention mask of all ones
+                    'attention_mask': torch.ones_like(input_ids),
+                }
+
+            # Create labels by cloning input_ids
+            labels = batch['input_ids'].clone()
+            # Mask labels only where attention_mask is 0 (padding positions)
+            if 'attention_mask' in batch:
+                labels[batch['attention_mask'] == 0] = -100
+            batch['labels'] = labels
+
+        else:
+            # --- Handling for varlen=True (Concatenated Sequences) ---
+            if len(examples) > 1:
+                raise ValueError('The batch size must be 1 for inputs with variable lengths (varlen=True).')
+
+            batch = {'input_ids': torch.cat([example['input_ids'] for example in examples], dim=0).unsqueeze(0)}
+
+            # --- cu_seqlens calculation logic remains the same ---
+            if 'cu_seqlens' in examples[0]:
+                batch['cu_seqlens'] = (
+                    torch.cat([example['cu_seqlens'] for example in examples], dim=0).unsqueeze(0).to(dtype=torch.int32)
+                )  # Ensure int32
+            else:
+                # determine boundaries by bos/eos positions
+                # Check for bos_token_id first
+                if self.tokenizer.bos_token_id is not None:
+                    cu_seqlens = []
+                    # Handle case where the sequence doesn't start with BOS
+                    if batch['input_ids'][0, 0] != self.tokenizer.bos_token_id:
+                        cu_seqlens.append(torch.tensor([0], device=batch['input_ids'].device))  # Match device
+                    # Find all BOS token positions
+                    bos_positions = torch.where(batch['input_ids'].eq(self.tokenizer.bos_token_id))[1]
+                    # Ensure bos_positions is on the correct device if empty
+                    if bos_positions.numel() == 0 and len(cu_seqlens) > 0:
+                        cu_seqlens.append(bos_positions.to(cu_seqlens[0].device))
+                    elif bos_positions.numel() > 0:
+                        cu_seqlens.append(bos_positions)
+                    # Add the end of the entire batch
+                    cu_seqlens.append(
+                        torch.tensor([batch['input_ids'].size(1)], device=batch['input_ids'].device)
+                    )  # Match device and use size(1)
+                    # Filter out empty tensors before cat
+                    cu_seqlens = [t for t in cu_seqlens if t.numel() > 0]
+                    if not cu_seqlens:  # Handle case where input is empty or has no BOS
+                        batch['cu_seqlens'] = torch.tensor(
+                            [0, batch['input_ids'].size(1)], dtype=torch.int32, device=batch['input_ids'].device
+                        )
+                    else:
+                        batch['cu_seqlens'] = torch.cat(cu_seqlens, dim=0).to(dtype=torch.int32)
+
+                # Else, check for eos_token_id
+                elif self.tokenizer.eos_token_id is not None:
+                    cu_seqlens = [torch.tensor([0], device=batch['input_ids'].device)]  # Match device
+                    # Find positions *after* EOS tokens
+                    eos_positions = torch.where(batch['input_ids'].eq(self.tokenizer.eos_token_id))[1] + 1
+                    # Ensure eos_positions is on the correct device if empty
+                    if eos_positions.numel() > 0:
+                        cu_seqlens.append(eos_positions)
+                    # Handle case where the sequence doesn't end with EOS
+                    if batch['input_ids'][0, -1] != self.tokenizer.eos_token_id:
+                        # Only add the final length if the last found EOS wasn't already the end
+                        if eos_positions.numel() == 0 or eos_positions[-1] != batch['input_ids'].size(1):
+                            cu_seqlens.append(
+                                torch.tensor([batch['input_ids'].size(1)], device=batch['input_ids'].device)
+                            )  # Match device and use size(1)
+                    # Filter out empty tensors before cat
+                    cu_seqlens = [t for t in cu_seqlens if t.numel() > 0]
+                    if not cu_seqlens:  # Handle case where input is empty or has no EOS
+                        batch['cu_seqlens'] = torch.tensor(
+                            [0, batch['input_ids'].size(1)], dtype=torch.int32, device=batch['input_ids'].device
+                        )
+                    else:
+                        batch['cu_seqlens'] = torch.cat(cu_seqlens, dim=0).to(dtype=torch.int32)
+                # Else, neither BOS nor EOS is usable
+                else:
+                    raise ValueError(
+                        'For varlen=True without precomputed cu_seqlens, the tokenizer must have either a bos_token_id '
+                        'or an eos_token_id defined to act as sequence separators.'
+                    )
+
+                # --- cu_seqlens validation checks remain the same ---
+                if batch['cu_seqlens'].numel() < 2:
+                    raise ValueError(f'Calculated cu_seqlens must have at least start and end: {batch["cu_seqlens"]}')
+                if not torch.all(batch['cu_seqlens'][1:] >= batch['cu_seqlens'][:-1]):
+                    raise ValueError(f'Calculated cu_seqlens are not monotonically increasing: {batch["cu_seqlens"]}')
+                if batch['cu_seqlens'][0] != 0:
+                    raise ValueError(f'Calculated cu_seqlens do not start at 0: {batch["cu_seqlens"]}')
+                if batch['cu_seqlens'][-1] != batch['input_ids'].size(1):
+                    # Allow empty sequence case where cu_seqlens=[0, 0] and input_ids.size(1)=0
+                    if not (batch['cu_seqlens'].tolist() == [0, 0] and batch['input_ids'].size(1) == 0):
+                        raise ValueError(
+                            f'Calculated cu_seqlens do not end at total length {batch["input_ids"].size(1)}: '
+                            f'{batch["cu_seqlens"]}'
+                        )
+
+                # --- context_len splitting logic remains the same ---
+                if self.context_len is not None:
+                    # This logic splits sequences based on context_len *after* initial boundaries are found
+                    bos = batch['cu_seqlens'][:-1].tolist()
+                    eos = batch['cu_seqlens'][1:].tolist()
+                    # Handle empty sequences between boundaries
+                    split_boundaries = []
+                    for i, j in zip(bos, eos):
+                        if i < j:  # Only process non-empty sequences
+                            split_boundaries.append(torch.arange(i, j, self.context_len, device=batch['input_ids'].device))
+                    # Add the final end point if it wasn't included by arange
+                    final_end_point = torch.tensor([batch['input_ids'].size(1)], device=batch['input_ids'].device)
+                    # Concatenate all boundaries
+                    if not split_boundaries:  # Handle case of completely empty input
+                        batch['cu_seqlens'] = torch.tensor([0, 0], dtype=torch.int32, device=batch['input_ids'].device)
+                    else:
+                        batch['cu_seqlens'] = torch.cat(split_boundaries + [final_end_point]).to(dtype=torch.int32)
+                        # Ensure uniqueness and sort, as arange might duplicate the endpoint
+                        batch['cu_seqlens'] = torch.unique(batch['cu_seqlens'])
+
+            # Create labels directly from input_ids, NO padding mask needed for varlen
+            labels = batch['input_ids'].clone()
+            batch['labels'] = labels
+
+        return batch
+
+
+class ParallelAwareDataLoader(StatefulDataLoader, Stateful):
+    """
+    A wrapper around the StatefulDataLoader that ensures that the state is stored only once per DP rank.
+    """
+
+    def __init__(
+        self,
+        rank: int,
+        dataset: IterableDataset,
+        batch_size: int,
+        collate_fn: Callable,
+        num_workers: int = 0,
+        pin_memory: bool = False,
+        prefetch_factor: int = 2,
+        persistent_workers: bool = False,
+        snapshot_every_n_steps: Optional[int] = 1,
+    ):
+        super().__init__(
+            dataset=dataset,
+            batch_size=batch_size,
+            collate_fn=collate_fn,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            prefetch_factor=prefetch_factor,
+            persistent_workers=persistent_workers,
+            snapshot_every_n_steps=snapshot_every_n_steps,
+        )
+        self.rank = rank
+
+    def state_dict(self) -> Dict[str, Any]:
+        # Store state only for dp rank to avoid replicating the same state across other dimensions
+        return {f'rank_{self.rank}': pickle.dumps(super().state_dict())}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        # State being empty is valid
+        if not state_dict:
+            return
+
+        if f'rank_{self.rank}' not in state_dict:
+            logger.warning(f'DataLoader state is empty for dp rank {self.rank}, expected key rank_{self.rank}')
+            return
+        super().load_state_dict(pickle.loads(state_dict[f'rank_{self.rank}']))
+
+
+def build_dataloader(
+    dataset: IterableDataset,
+    tokenizer: PreTrainedTokenizer,
+    rank: int,
+    world_size: int,
+    batch_size: int,
+    seq_len: int,
+    context_len: Optional[int] = None,
+    varlen: bool = False,
+    num_workers: int = 0,
+    pin_memory: bool = False,
+    persistent_workers: bool = False,
+    snapshot_every_n_steps: Optional[int] = 1,
+):
+    dataset = OnlineTokenizedIterableDataset(
+        dataset=dataset, tokenizer=tokenizer, seq_len=seq_len, rank=rank, world_size=world_size
+    )
+    return ParallelAwareDataLoader(
+        rank=rank,
+        dataset=dataset,
+        batch_size=batch_size,
+        collate_fn=DataCollatorForLanguageModeling(tokenizer=tokenizer, context_len=context_len, varlen=varlen),
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        persistent_workers=persistent_workers,
+        snapshot_every_n_steps=snapshot_every_n_steps,
+    )

flame/utils/checkpoint.py

@@ -0,0 +1,50 @@
+import os
+import glob
+import re
+import shutil
+from torchtitan.tools.logging import logger
+
+
+def cleanup_local_checkpoints(checkpoint_dir: str, keep_latest_k: int):
+    """Removes older checkpoint directories locally, keeping only the latest k for both DCP and HF formats."""
+    if keep_latest_k <= 0:
+        return # Keep all checkpoints
+
+    logger.info(f"Cleaning up local checkpoints in {checkpoint_dir}, keeping latest {keep_latest_k}")
+
+    # Cleanup DCP checkpoints (step-*)
+    dcp_checkpoints = sorted(
+        glob.glob(os.path.join(checkpoint_dir, "step-*")),
+        key=lambda x: int(re.search(r"step-(\d+)", os.path.basename(x)).group(1)) if re.search(r"step-(\d+)", os.path.basename(x)) and not x.endswith("-hf") else -1,
+        reverse=True
+    )
+    # Filter out HF format directories
+    dcp_checkpoints = [d for d in dcp_checkpoints if not d.endswith("-hf")]
+
+    if len(dcp_checkpoints) > keep_latest_k:
+        checkpoints_to_delete = dcp_checkpoints[keep_latest_k:]
+        logger.info(f"Deleting {len(checkpoints_to_delete)} old DCP checkpoints: {[os.path.basename(c) for c in checkpoints_to_delete]}")
+        for ckpt_path in checkpoints_to_delete:
+            if os.path.isdir(ckpt_path): # Ensure it's a directory
+                 try:
+                     shutil.rmtree(ckpt_path)
+                 except OSError as e:
+                     logger.error(f"Error removing directory {ckpt_path}: {e}")
+
+
+    # Cleanup HF checkpoints (step-*-hf)
+    hf_checkpoints = sorted(
+        glob.glob(os.path.join(checkpoint_dir, "step-*-hf")),
+         key=lambda x: int(re.search(r"step-(\d+)-hf", os.path.basename(x)).group(1)) if re.search(r"step-(\d+)-hf", os.path.basename(x)) else -1,
+        reverse=True
+    )
+
+    if len(hf_checkpoints) > keep_latest_k:
+        checkpoints_to_delete = hf_checkpoints[keep_latest_k:]
+        logger.info(f"Deleting {len(checkpoints_to_delete)} old HF checkpoints: {[os.path.basename(c) for c in checkpoints_to_delete]}")
+        for ckpt_path in checkpoints_to_delete:
+             if os.path.isdir(ckpt_path): # Ensure it's a directory
+                 try:
+                     shutil.rmtree(ckpt_path)
+                 except OSError as e:
+                     logger.error(f"Error removing directory {ckpt_path}: {e}")

flame/utils/hf_utils.py

@@ -0,0 +1,77 @@
+import os
+import re
+from huggingface_hub import HfApi, HfFolder, logging as hf_logging, create_repo
+from torchtitan.tools.logging import logger 
+
+def upload_checkpoint_to_hf(
+    local_path: str,         
+    step: int,
+    hf_repo_id_for_run: str,
+    hf_keep_latest_k: int,
+    upload_format: str        
+):
+    """Uploads a checkpoint directory to HF Hub and manages retention."""
+    if not os.path.isdir(local_path):
+        logger.error(f"Local path for upload does not exist or is not a directory: {local_path}")
+        return
+
+    api = HfApi()
+    token = HfFolder.get_token() 
+    if not token:
+        logger.warning("Hugging Face Hub token not found. Skipping upload. Login via `huggingface-cli login` or set HF_TOKEN.")
+        return
+
+    # --- Ensure the specific repository for this run exists ---
+    try:
+        logger.info(f"Ensuring repository {hf_repo_id_for_run} exists...")
+        # Use create_repo which handles creation only if it doesn't exist
+        create_repo(repo_id=hf_repo_id_for_run, token=token, repo_type="model", exist_ok=True)
+        logger.info(f"Repository {hf_repo_id_for_run} ensured.")
+    except Exception as e:
+        logger.error(f"Failed to create or ensure repository {hf_repo_id_for_run}: {e}", exc_info=True)
+        return # Stop if repo interaction fails
+
+    commit_message = f"Upload {upload_format.upper()} checkpoint step {step}"
+    path_in_repo = f"step-{step}"
+
+    logger.info(f"Uploading {local_path} to {hf_repo_id_for_run}/{path_in_repo} on Hugging Face Hub...")
+    try:
+        api.upload_folder(
+            folder_path=local_path,
+            path_in_repo=path_in_repo,
+            repo_id=hf_repo_id_for_run,
+            repo_type="model", 
+            commit_message=commit_message,
+            token=token,
+        )
+        logger.info(f"Successfully uploaded step {step} to {hf_repo_id_for_run}.")
+    except Exception as e:
+        logger.error(f"Failed to upload checkpoint step {step} to {hf_repo_id_for_run}: {e}", exc_info=True)
+    if hf_keep_latest_k > 0:
+        logger.info(f"Cleaning up old checkpoints on {hf_repo_id_for_run}, keeping latest {hf_keep_latest_k}")
+        try:
+            repo_files = api.list_repo_tree(hf_repo_id_for_run, repo_type="model", token=token, recursive=False)
+            step_folders = [
+                item.path for item in repo_files
+                if item.path.startswith("step-") and item.path[5:].isdigit()
+            ]
+
+            step_folders.sort(key=lambda x: int(x.split('-')[1]), reverse=True)
+
+            if len(step_folders) > hf_keep_latest_k:
+                folders_to_delete = step_folders[hf_keep_latest_k:]
+                logger.info(f"Found {len(step_folders)} checkpoints on Hub. Deleting {len(folders_to_delete)} older ones: {folders_to_delete}")
+                for folder in folders_to_delete:
+                    # Deleting requires repo_id, path_in_repo, and token
+                    api.delete_folder(
+                        repo_id=hf_repo_id_for_run,
+                        path_in_repo=folder,
+                        repo_type="model",
+                        commit_message=f"Delete old checkpoint {folder}",
+                        token=token
+                    )
+                logger.info("Hub cleanup complete.")
+            else:
+                logger.info("No old checkpoints found on Hub to delete.")
+        except Exception as e:
+            logger.error(f"Error during Hub checkpoint cleanup for {hf_repo_id_for_run}: {e}", exc_info=True)

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "transformers_version": "4.51.3"
+}

measure_sink_rate.py

@@ -0,0 +1,137 @@
+import fla
+import torch
+import numpy as np
+import argparse
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from datasets import load_dataset
+
+def calculate_sink_rate(attention_maps, epsilon=0.3):
+    """
+    Calculate sink rate using the formula:
+    sink_rate = (1/(L*H))*sum_L_H(1_((1/T)*sum_T(a_l_h_1_t) > epsilon))
+    
+    Where:
+    - L is the number of layers
+    - H is the number of attention heads
+    - T is the sequence length
+    - 1_() is the indicator function
+    - a_ is the attention score at that index
+    - epsilon is the threshold (default: 0.3)
+    
+    Args:
+        attention_maps: Attention maps from the model as a list with length L of tensors with shape [batch, heads, seq_len, seq_len]
+        epsilon: Threshold for attention
+        
+    Returns:
+        sink_rate: The calculated sink rate
+    """
+    sink_rate = 0
+    for i, attention in enumerate(attention_maps):
+        # Extract attention on first token (BOS) across all heads
+        first_token_attention = attention[:, :, :, 0]  # [batch, heads, seq_len]
+        # print("first token attentions", first_token_attention)
+        
+        # Calculate mean attention on first token across sequence length
+        mean_first_token_attention = first_token_attention.mean(dim=-1)  # [batch, heads]
+        # print("mean first token attentions", mean_first_token_attention)
+        
+        # Apply indicator function - whether mean attention > epsilon
+        indicator = (mean_first_token_attention > epsilon).float()  # [batch, heads]
+        # print("indicator", indicator)
+        
+        # Average across heads
+        batch_sink_rates = indicator.mean(dim=(1))  # [batch]
+        
+        # Average across batch
+        sink_rate += batch_sink_rates.mean().item()
+
+    # Normalize by number of layers
+    num_layers = len(attention_maps)
+    sink_rate /= num_layers
+        
+    return sink_rate
+
+def main(args):
+    # Set device
+    device = torch.device("cuda" if torch.cuda.is_available() and not args.cpu else "cpu")
+    print(f"Using device: {device}")
+    
+    # Load model and tokenizer
+    print(f"Loading model: {args.model_name}")
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    model = AutoModelForCausalLM.from_pretrained(args.model_name, return_dict_in_generate=True, output_attentions=True).half()
+    model.to(device)
+    model.eval()
+    
+    # Add padding token if it doesn't exist
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    # Load dataset
+    print(f"Loading dataset: {args.dataset_name}")
+    dataset = load_dataset(args.dataset_name, split=args.split)
+    
+    # Take n_samples from the dataset
+    if args.n_samples < len(dataset):
+        dataset = dataset.select(range(args.n_samples))
+    else:
+        args.n_samples = len(dataset)
+    
+    print(f"Processing {args.n_samples} samples...")
+    
+    # Process samples
+    sink_rate = 0
+    
+    for i in tqdm(range(0, args.n_samples, args.batch_size)):
+        batch_samples = dataset[i:min(i + args.batch_size, args.n_samples)]
+        
+        # Tokenize
+        encodings = tokenizer(
+            batch_samples["text"], 
+            padding=True, 
+            truncation=True, 
+            max_length=args.max_length, 
+            return_tensors="pt"
+        ).to(device)
+        
+        # Forward pass
+        with torch.no_grad():
+            outputs = model(**encodings)
+        
+        # Calculate sink rate for this batch
+        batch_sink_rate = calculate_sink_rate(outputs.attentions, args.epsilon)
+        attention_maps = None  # Free memory
+        sink_rate += batch_sink_rate
+    
+    # Average sink rate
+    sink_rate /= args.n_samples
+    
+    print(f"Sink Rate (ε={args.epsilon}): {sink_rate:.4f}")
+    
+    # Optional: Save sink rate results
+    if args.output_file:
+        with open(args.output_file, 'w') as f:
+            f.write(f"Model: {args.model_name}\n")
+            f.write(f"Dataset: {args.dataset_name}\n")
+            f.write(f"Split: {args.split}\n")
+            f.write(f"Samples: {args.n_samples}\n")
+            f.write(f"Epsilon: {args.epsilon}\n")
+            f.write(f"Sink Rate: {sink_rate:.4f}\n")
+            
+        print(f"Results saved to {args.output_file}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Measure Sink Rate for Transformer Models")
+    parser.add_argument("--model_name", type=str, default="gpt2", help="Huggingface model name")
+    parser.add_argument("--dataset_name", type=str, default="wikitext", help="Huggingface dataset name")
+    parser.add_argument("--split", type=str, default="test", help="Dataset split to use")
+    parser.add_argument("--n_samples", type=int, default=100, help="Number of samples to process")
+    parser.add_argument("--batch_size", type=int, default=4, help="Batch size for processing")
+    parser.add_argument("--max_length", type=int, default=512, help="Maximum sequence length")
+    parser.add_argument("--epsilon", type=float, default=0.3, help="Threshold for sink rate calculation")
+    parser.add_argument("--output_file", type=str, default="", help="File to save results")
+    parser.add_argument("--cpu", action="store_true", help="Force CPU usage")
+    
+    args = parser.parse_args()
+    main(args)

passkey_retrieval.py

@@ -0,0 +1,158 @@
+import torch
+import argparse
+import random
+from numpy import random
+from tqdm import tqdm
+import fla
+import transformers
+from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
+
+def parse_config():
+    parser = argparse.ArgumentParser(description='arg parser')
+    parser.add_argument('--base_model', type=str, default="/data1/pretrained-models/llama-7b-hf")
+    parser.add_argument('--cache_dir', type=str, default="./cache")
+    parser.add_argument('--context_size', type=int, default=-1, help='context size during fine-tuning')
+    parser.add_argument('--num_tests', type=int, default=10, help='number of repeat testing for each length')
+    parser.add_argument('--test_k_tokens', type=int, default=2, help='test length')
+    
+    args = parser.parse_args()
+    return args
+
+
+def generate_prompt_landmark(tokenizer, n_garbage, seed):
+    """Generates a text file and inserts an passkey at a random position."""
+    rnd_state = random.get_state()
+    random.seed(seed)
+    n_garbage_prefix = random.randint(0, n_garbage)
+    n_garbage_suffix = n_garbage - n_garbage_prefix
+
+    task_description = "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there."
+    garbage = "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again."
+    garbage_inf = " ".join([garbage] * 100000)
+    assert len(garbage_inf) >= n_garbage
+    garbage_prefix = garbage_inf[:n_garbage_prefix]
+    garbage_suffix = garbage_inf[:n_garbage_suffix]
+    pass_key = random.randint(1, 50000)
+    
+    
+    information_line = f"The pass key is {pass_key}. Remember it. {pass_key} is the pass key."
+    final_question = "What is the pass key? The pass key is"
+    lines = [
+        task_description,
+        garbage_prefix,
+        information_line,
+        garbage_suffix, # generate: 1k token
+        final_question,
+    ]
+    
+    position = []
+    current_pos = 0
+    block = 0
+
+    for line in lines:
+        input_ids = tokenizer(line, return_tensors="pt").input_ids
+        line_length = input_ids.size(1)  # Get the length of the sequence
+        position.append((current_pos, current_pos + line_length - 1))  # Store the start and end positions
+        if line.startswith("The pass key is"):
+            block = current_pos // 256 
+        current_pos += line_length  # Update the current position
+    
+    random.set_state(rnd_state)
+    return "\n".join(lines), str(pass_key), position[2], block
+
+
+def passkey_retrieval_test(model, tokenizer, device, use_cache=False, n_garbage=60000, seed=666, sequence=None):
+    prompt, answer, position, block = generate_prompt_landmark(tokenizer, n_garbage, seed+n_garbage) # 修改 seed 为 seed+n，防止重复
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    input_ids = input_ids.to(device)
+    len_token = input_ids.shape[-1]
+
+    answer_ids = tokenizer(answer, return_tensors="pt").input_ids[:, 1:] # drop BOS
+
+    generation_output = model.generate(
+        input_ids=input_ids, max_new_tokens=answer_ids.shape[-1], num_beams=1, use_cache=True
+    )
+    
+    model_answer = generation_output[0, -answer_ids.shape[-1]:].cpu()
+    
+    is_correct = (model_answer == answer_ids[0]).all().item()
+    is_split = is_number_in_range(sequence, position)
+    
+        
+    print(f"The correct answer is {tokenizer.decode(answer_ids[0].cpu())}")
+    print(f"The model answer is {tokenizer.decode(model_answer.cpu())}, is_correct : {is_correct}")
+    return is_correct, is_split, len_token, position
+
+
+def generate_sequence():
+    sequence = [0]  
+    increment = 256  
+
+    for i in range(1, 201):  
+        next_number = sequence[-1] + increment
+        sequence.append(next_number)
+
+    return sequence
+
+def is_number_in_range(sequence, position):
+    for num in sequence:
+        if position[0] <= num and num <= position[1]:
+            return True
+    return False
+
+def main(args):
+    device = "cuda:0"
+    torch.cuda.set_device(device)
+
+    print("base model", args.base_model)
+
+    # Set RoPE scaling factor
+    # config = AutoConfig.from_pretrained(
+    #         args.base_model,
+    #         cache_dir=args.cache_dir,
+    #     )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        args.base_model,
+        # config=config,
+        # cache_dir=args.cache_dir,
+        # torch_dtype=torch.float16,
+    )
+    model = model.to('cuda:0').half()
+    
+    model.resize_token_embeddings(32001)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.base_model,
+        # cache_dir=args.cache_dir,
+        # model_max_length=args.context_size,
+        # padding_side="right",
+        # use_fast=False,
+    )
+
+
+    sequence = generate_sequence()
+
+    
+    # This is a rough ratio to control the number of texts and tokens
+    n_garbage = int(3.75 * args.test_k_tokens * 1024 // 1024 * 1024)
+    passed_tests = 0
+    total_tokens = 0
+    for j in tqdm(range(args.num_tests)):
+        is_correct, is_split, len_tokens, position = passkey_retrieval_test(model, tokenizer, device, use_cache=True, n_garbage=n_garbage, seed=j, sequence=sequence)
+        
+        passed_tests += is_correct
+        total_tokens += len_tokens
+        if is_correct:
+            print(f" Success: {position},\tis_split: {is_split}", end="", flush=True)
+        else:
+            print(f" [Fails]: {position},\tis_split: {is_split}", end="", flush=True)
+    avg_tokens = total_tokens//args.num_tests
+    accuracy = float(passed_tests)/args.num_tests
+    print("Accuracy on the token length %d is %f, max GPU allocate %f GB"%(avg_tokens, accuracy, torch.cuda.max_memory_allocated(0) / 1024 / 1024 / 1024), flush=True)
+
+
+
+if __name__ == "__main__":
+    args = parse_config()
+    main(args)

pyproject.toml

@@ -0,0 +1,42 @@
+[project]
+name = "flame"
+dynamic = ["version"]
+description = "A minimal training framework for scaling FLA models"
+readme = "README.md"
+authors = [
+    { name = "Songlin Yang", email = "yangsl66@mit.edu" },
+    { name = "Yu Zhang", email = "yzhang.cs@outlook.com" },
+]
+license = { file = "LICENSE" }
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+requires-python = ">=3.10"
+dependencies = [
+    'torch>=2.5',
+    'torchdata',
+    'transformers>=4.45.0',
+    'triton>=3.0',
+    'datasets>=3.3.0',
+    'einops',
+    'ninja',
+    'wandb',
+    'tiktoken',
+    'tensorboard',
+]
+
+[project.optional-dependencies]
+dev = ["pytest"]
+
+[project.urls]
+Homepage = "https://github.com/fla-org/flame"
+
+[build-system]
+requires = ["setuptools>=45", "wheel", "ninja", "torch"]
+
+[tool.isort]
+line_length = 127
+multi_line_output = 3

register_softpick.py

@@ -0,0 +1,72 @@
+from einops import rearrange
+from typing import Optional
+import torch
+from transformers import AttentionInterface
+from torch.nn import functional as F
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+def softpick(x, dim=-1, eps=1e-8):
+    # softpick function: relu(exp(x)-1) / sum(abs(exp(x)-1))
+    # numerically stable version
+    x_m = torch.max(x, dim=dim, keepdim=True).values
+    x_m_e_m = torch.exp(-x_m)
+    x_e_1 = torch.exp(x - x_m) - x_m_e_m
+    r_x_e_1 = F.relu(x_e_1)
+    a_x_e_1 = torch.where(x.isfinite(), torch.abs(x_e_1), 0)
+    return r_x_e_1 / (torch.sum(a_x_e_1, dim=dim, keepdim=True) + eps) # epsilon is only useful if all inputs are EXACTLY 0. we might not even need it
+
+def naive_softpick_attn(
+    module: torch.nn.Module,  # required arg
+    query: torch.Tensor,  # required arg
+    key: torch.Tensor,  # required arg
+    value: torch.Tensor,  # required arg
+    attention_mask: Optional[torch.Tensor],  # required arg
+    *args,
+    scale: Optional[float] = None,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    head_first: bool = False,
+    **kwargs
+) -> torch.Tensor:
+    head_dim = query.shape[-1]
+
+    # In transformers, the shape is (batch_size, num_heads, seq_len, head_dim)
+    num_query_heads = query.shape[1]
+    num_key_valye_heads = key.shape[1]
+
+
+    if num_query_heads != num_key_valye_heads:
+        # MQA or GQA
+        key = repeat_kv(key, num_query_heads // num_key_valye_heads)
+        value = repeat_kv(value, num_query_heads // num_key_valye_heads)
+
+    if scale is None:
+        scale = 1.0 / (head_dim ** 0.5)
+    if not head_first:
+        query, key, value = map(lambda x: rearrange(x, 'b t h d -> b h t d'), (query, key, value))
+    query_len = query.shape[-2]
+    key_len = key.shape[-2]
+    mask = torch.tril(torch.ones(key_len, key_len, device=query.device))
+    wei = torch.matmul(query, key.transpose(2, 3)) # shape: (batch_size, num_heads, query_len, key_len)
+    wei = wei * scale
+    wei = wei.masked_fill(mask[key_len-query_len:key_len, :key_len] == 0, float('-inf'))
+    wei = softpick(wei.float(), dim=-1).to(query.dtype)
+    o = torch.matmul(wei, value) # shape: (batch_size, num_heads, q_len, head_dim)
+    if not head_first:
+        o = rearrange(o, 'b h t d -> b t h d')
+    return o, wei
+
+def softpick_attention(*args, **kwargs):
+    # print("Using softpick attention") # NOTE: Add print statement here to check whether we actually use softpick or not
+    return naive_softpick_attn(*args, **kwargs)
+
+AttentionInterface.register("softpick", softpick_attention)

setup.py

@@ -0,0 +1,50 @@
+# -*- coding: utf-8 -*-
+
+import ast
+import os
+import re
+from pathlib import Path
+
+from setuptools import find_packages, setup
+
+with open('README.md') as f:
+    long_description = f.read()
+
+
+def get_package_version():
+    with open(Path(os.path.dirname(os.path.abspath(__file__))) / 'flame' / '__init__.py') as f:
+        version_match = re.search(r"^__version__\s*=\s*(.*)$", f.read(), re.MULTILINE)
+    return ast.literal_eval(version_match.group(1))
+
+
+setup(
+    name='flame',
+    version=get_package_version(),
+    description='A minimal training framework for scaling FLA models',
+    long_description=long_description,
+    long_description_content_type='text/markdown',
+    author='Songlin Yang, Yu Zhang',
+    author_email='yangsl66@mit.edu, yzhang.cs@outlook.com',
+    url='https://github.com/fla-org/flame',
+    packages=find_packages(),
+    license='MIT',
+    classifiers=[
+        'Programming Language :: Python :: 3',
+        'License :: OSI Approved :: MIT License',
+        'Operating System :: OS Independent',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence'
+    ],
+    python_requires='>=3.10',
+    install_requires=[
+        'torch>=2.5',
+        'torchdata',
+        'transformers>=4.45.0',
+        'triton>=3.0',
+        'datasets>=3.3.0',
+        'einops',
+        'ninja',
+        'wandb',
+        'tiktoken',
+        'tensorboard',
+    ],
+)