Add files using upload-large-folder tool

LICENSE

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

README.md

@@ -0,0 +1,471 @@
+<div align="center">
+
+# 🔥 Flame: Flash Linear Attention Made Easy
+
+</div>
+
+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,34 @@
+{
+  "architectures": [
+    "MTPTransformerForCausalLM"
+  ],
+  "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": "mtp_transformer",
+  "n_future_tokens": 4,
+  "norm_eps": 1e-06,
+  "num_heads": 32,
+  "num_hidden_layers": 32,
+  "num_kv_heads": null,
+  "pad_token_id": 2,
+  "qk_norm": false,
+  "qkv_bias": false,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.50.3",
+  "use_cache": true,
+  "use_custom_backward": false,
+  "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/mtp_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": "mtp_transformer",
+    "num_heads": 12,
+    "num_hidden_layers": 14,
+    "norm_eps": 1e-06,
+    "tie_word_embeddings": true,
+    "use_cache": true,
+    "vocab_size": 32000,
+    "n_future_tokens": 4
+}

configs/mtp_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": "mtp_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,
+    "n_future_tokens": 4
+}

configs/mtp_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": "mtp_transformer",
+    "num_heads": 16,
+    "num_hidden_layers": 24,
+    "norm_eps": 1e-06,
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "vocab_size": 32000,
+    "n_future_tokens": 4
+}

configs/mtp_transformer_7B.json

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

configs/top_transformer_120M.json

@@ -0,0 +1,20 @@
+{
+    "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": "top_transformer",
+    "num_heads": 12,
+    "num_hidden_layers": 14,
+    "norm_eps": 1e-06,
+    "tie_word_embeddings": true,
+    "use_cache": true,
+    "vocab_size": 32000,
+    "use_top_loss": true,
+    "top_window_size": 2048
+}

configs/top_transformer_1B.json

@@ -0,0 +1,24 @@
+{
+    "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": "top_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,
+    "use_top_loss": true,
+    "top_window_size": 4096
+}

configs/top_transformer_340M.json

@@ -0,0 +1,20 @@
+{
+    "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": "top_transformer",
+    "num_heads": 16,
+    "num_hidden_layers": 24,
+    "norm_eps": 1e-06,
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "vocab_size": 32000,
+    "use_top_loss": true,
+    "top_window_size": 4096
+}

configs/top_transformer_7B.json

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

configs/transformer_120M.json

@@ -0,0 +1,18 @@
+{
+    "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
+}

configs/transformer_1B.json

@@ -0,0 +1,22 @@
+{
+    "bos_token_id": 1,
+    "elementwise_affine": true,
+    "eos_token_id": 2,
+    "fuse_cross_entropy": true,
+    "fuse_norm": true,
+    "fuse_swiglu": true,
+    "hidden_act": "swish",
+    "hidden_ratio": 4,
+    "hidden_size": 2048,
+    "initializer_range": 0.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
+}

configs/transformer_340M.json

@@ -0,0 +1,18 @@
+{
+    "attention_bias": false,
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "fuse_cross_entropy": true,
+    "fuse_norm": true,
+    "hidden_act": "swish",
+    "hidden_size": 1024,
+    "initializer_range": 0.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
+}

configs/transformer_7B.json

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

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/__init__.py

@@ -0,0 +1,110 @@
+# -*- coding: utf-8 -*-
+
+from fla.layers import (
+    ABCAttention,
+    Attention,
+    BasedLinearAttention,
+    BitAttention,
+    DeltaNet,
+    GatedDeltaNet,
+    GatedDeltaProduct,
+    GatedLinearAttention,
+    GatedSlotAttention,
+    HGRN2Attention,
+    HGRNAttention,
+    LightNetAttention,
+    LinearAttention,
+    MultiScaleRetention,
+    NativeSparseAttention,
+    ReBasedLinearAttention,
+    RWKV6Attention,
+    RWKV7Attention
+)
+from fla.models import (
+    ABCForCausalLM,
+    ABCModel,
+    BitNetForCausalLM,
+    BitNetModel,
+    DeltaNetForCausalLM,
+    DeltaNetModel,
+    GatedDeltaNetForCausalLM,
+    GatedDeltaNetModel,
+    GatedDeltaProductForCausalLM,
+    GatedDeltaProductModel,
+    GLAForCausalLM,
+    GLAModel,
+    GSAForCausalLM,
+    GSAModel,
+    HGRN2ForCausalLM,
+    HGRN2Model,
+    HGRNForCausalLM,
+    LightNetForCausalLM,
+    LightNetModel,
+    LinearAttentionForCausalLM,
+    LinearAttentionModel,
+    NSAForCausalLM,
+    NSAModel,
+    RetNetForCausalLM,
+    RetNetModel,
+    RWKV6ForCausalLM,
+    RWKV6Model,
+    RWKV7ForCausalLM,
+    RWKV7Model,
+    TransformerForCausalLM,
+    TransformerModel
+)
+
+__all__ = [
+    'ABCAttention',
+    'Attention',
+    'BasedLinearAttention',
+    'BitAttention',
+    'DeltaNet',
+    'GatedDeltaNet',
+    'GatedDeltaProduct',
+    'GatedLinearAttention',
+    'GatedSlotAttention',
+    'HGRNAttention',
+    'HGRN2Attention',
+    'LightNetAttention',
+    'LinearAttention',
+    'MultiScaleRetention',
+    'NativeSparseAttention',
+    'ReBasedLinearAttention',
+    'RWKV6Attention',
+    'RWKV7Attention',
+    'ABCForCausalLM',
+    'ABCModel',
+    'BitNetForCausalLM',
+    'BitNetModel',
+    'DeltaNetForCausalLM',
+    'DeltaNetModel',
+    'GatedDeltaNetForCausalLM',
+    'GatedDeltaNetModel',
+    'GatedDeltaProductForCausalLM',
+    'GatedDeltaProductModel',
+    'GLAForCausalLM',
+    'GLAModel',
+    'GSAForCausalLM',
+    'GSAModel',
+    'HGRNForCausalLM',
+    'HGRNModel',
+    'HGRN2ForCausalLM',
+    'HGRN2Model',
+    'LightNetForCausalLM',
+    'LightNetModel',
+    'LinearAttentionForCausalLM',
+    'LinearAttentionModel',
+    'NSAForCausalLM',
+    'NSAModel',
+    'RetNetForCausalLM',
+    'RetNetModel',
+    'RWKV6ForCausalLM',
+    'RWKV6Model',
+    'RWKV7ForCausalLM',
+    'RWKV7Model',
+    'TransformerForCausalLM',
+    'TransformerModel',
+]
+
+__version__ = '0.1.2'

fla/utils.py

@@ -0,0 +1,223 @@
+# -*- coding: utf-8 -*-
+
+import contextlib
+import functools
+import os
+from enum import Enum
+from functools import lru_cache
+from typing import Any, Callable, Dict, Literal, Optional, Tuple
+
+import torch
+import triton
+from packaging import version
+
+
+def tensor_cache(
+    fn: Callable[..., torch.Tensor]
+) -> Callable[..., torch.Tensor]:
+    """
+    A decorator that caches the most recent result of a function with tensor inputs.
+
+    This decorator will store the output of the decorated function for the most recent set of input tensors.
+    If the function is called again with the same input tensors, it will return the cached result.
+
+
+    Args:
+        fn (Callable[..., torch.Tensor]):
+            The function to be decorated. It should take tensor inputs and return tensor outputs.
+
+    Returns:
+        Callable[..., torch.Tensor]:
+            A wrapped version of the input function with single-entry caching.
+    """
+    last_args: Optional[Tuple] = None
+    last_kwargs: Optional[Dict] = None
+    last_result: Any = None
+
+    @functools.wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        nonlocal last_args, last_kwargs, last_result
+
+        if last_args is not None and last_kwargs is not None:
+            if len(args) == len(last_args) and len(kwargs) == len(last_kwargs):
+                if all(a is b for a, b in zip(args, last_args)) and \
+                        all(k in last_kwargs and v is last_kwargs[k] for k, v in kwargs.items()):
+                    return last_result
+
+        result = fn(*args, **kwargs)
+        last_args, last_kwargs, last_result = args, kwargs, result
+        return result
+
+    return wrapper
+
+
+def input_guard(
+    fn: Callable[..., torch.Tensor]
+) -> Callable[..., torch.Tensor]:
+    """
+    A decorator to make sure all input tensors are contiguous and set the device based on input tensors.
+    """
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        contiguous_args = (i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args)
+        contiguous_kwargs = {k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) for k, v in kwargs.items()}
+
+        tensor = None
+        for arg in args:
+            if isinstance(arg, torch.Tensor):
+                tensor = arg
+                break
+        if tensor is None:
+            for value in kwargs.values():
+                if isinstance(value, torch.Tensor):
+                    tensor = value
+                    break
+
+        if tensor is not None:
+            ctx = custom_device_ctx(tensor.device.index)
+        else:
+            ctx = contextlib.nullcontext()
+
+        with ctx:
+            return fn(*contiguous_args, **contiguous_kwargs)
+
+    return wrapper
+
+
+contiguous = input_guard
+
+
+def require_version(version, hint):
+    """
+    Perform a runtime check of the dependency versions, using the exact same syntax used by pip.
+    """
+    def decorator(fn):
+        @functools.wraps(fn)
+        def wrapper(ctx, *args, **kwargs):
+            from transformers.utils.versions import require_version
+            require_version(version, hint)
+            return fn(ctx,
+                      *(i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args),
+                      **{k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) for k, v in kwargs.items()})
+        return wrapper
+    return decorator
+
+
+def checkpoint(fn):
+    def wrapper(*args, **kwargs):
+        return torch.utils.checkpoint.checkpoint(fn, *args, **kwargs)
+    return wrapper
+
+
+@lru_cache(maxsize=None)
+def check_pytorch_version(version_s: str = '2.4') -> bool:
+    return version.parse(torch.__version__) >= version.parse(version_s)
+
+
+def _cpu_device_warning():
+    import warnings
+    warnings.warn(('Triton is not supported on current platform, roll back to CPU.'), stacklevel=1)
+
+
+@lru_cache(maxsize=None)
+def get_multiprocessor_count(tensor_idx: int = 0) -> int:
+    try:
+        # Only works if Homogeneous hardware
+        # TEMPORARY FIX since old version introduce graph break
+        return torch.cuda.get_device_properties().multi_processor_count
+    except BaseException:
+        _cpu_device_warning()
+        return -1
+
+
+@lru_cache(maxsize=None)
+def get_available_device() -> str:
+    try:
+        return triton.runtime.driver.active.get_current_target().backend
+    except BaseException:
+        _cpu_device_warning()
+        return 'cpu'
+
+
+@lru_cache(maxsize=None)
+def _check_platform() -> Literal['nvidia', 'amd', 'intel', 'musa']:
+    device = get_available_device()
+    if device == 'cuda':
+        return 'nvidia'
+    elif device == 'hip':
+        return 'amd'
+    elif device == 'xpu':
+        return 'intel'
+    else:
+        return device
+
+
+# For AMD GPUs, the triton backend is 'hip', while for Nvidia GPUs, the triton backend is 'cuda'.
+# However, the torch backend is 'cuda' for both Nvidia and AMD GPUs.
+# Therefore, we need to check the triton backend to determine the actual GPU vendor.
+device = get_available_device() if get_available_device() != 'hip' else 'cuda'
+device_torch_lib = getattr(torch, device)
+device_platform = _check_platform()
+
+is_amd = (device_platform == 'amd')
+is_intel = (device_platform == 'intel')
+is_nvidia = (device_platform == 'nvidia')
+is_intel_alchemist = (is_intel and 'Intel(R) Arc(TM) A' in torch.xpu.get_device_name(0))
+is_nvidia_hopper = (is_nvidia and ('NVIDIA H' in torch.cuda.get_device_name(0) or torch.cuda.get_device_capability()[0] >= 9))
+use_cuda_graph = (is_nvidia and os.environ.get('FLA_USE_CUDA_GRAPH', '0') == '1')
+
+# Nvidia Ampere or newer, haven't check AMD and intel yet.
+is_tf32_supported = (is_nvidia and torch.cuda.get_device_capability(0)[0] >= 8)
+is_gather_supported = hasattr(triton.language, 'gather')
+
+
+def get_all_max_shared_mem():
+    try:
+        return [
+            triton.runtime.driver.active.utils.get_device_properties(i)['max_shared_mem']
+            for i in range(device_torch_lib.device_count())
+        ]
+    except BaseException:
+        _cpu_device_warning()
+        return [-1]
+
+
+class Backend(Enum):
+    ADA = 101376       # RTX 4090
+    AMPERE = 166912    # A100
+    HOPPER = 232448    # H100
+    DEFAULT = 102400   # Default
+
+    @classmethod
+    def get_shared_memory(cls, arch: str) -> int:
+        try:
+            return cls[arch.upper()].value
+        except KeyError:
+            return cls.DEFAULT.value
+
+
+@lru_cache(maxsize=None)
+def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool:
+    try:
+        device_shared_mem_list = get_all_max_shared_mem()
+        max_shared_memory = device_shared_mem_list[tensor_idx]
+        return max_shared_memory >= Backend.get_shared_memory(arch)
+    except Exception:
+        return False
+
+
+if check_pytorch_version('2.4'):
+    device = 'cuda' if device == 'cpu' else device
+    autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=device)
+    autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=device)
+
+    def custom_device_ctx(index: int):
+        return device_torch_lib.device(index)
+else:
+    assert device == 'cuda', 'Only cuda device is supported for PyTorch version < 2.4.0.'
+    autocast_custom_fwd = device_torch_lib.amp.custom_fwd
+    autocast_custom_bwd = device_torch_lib.amp.custom_bwd
+
+    def custom_device_ctx(index: int):
+        return torch.cuda.device(index)

flame/__init__.py

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

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/tools/utils.py

@@ -0,0 +1,41 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torch import nn
+from torchtitan.tools.logging import logger
+
+
+def get_nparams_and_flops(model: nn.Module, model_config, seq_len: int) -> tuple[int, int]:
+    nparams = sum(p.numel() for p in model.parameters())
+    nparams_embedding = sum(
+        sum(p.numel() for p in m.parameters())
+        for m in model.children()
+        if isinstance(m, nn.Embedding)
+    )
+    
+    if hasattr(model_config, "num_heads"):
+        num_heads = model_config.num_heads
+    elif hasattr(model_config, "num_attention_heads"):
+        num_heads = model_config.num_attention_heads
+    else:
+        num_heads = 1
+        logger.warning("num_heads not found in model_config, defaulting to 1. ")
+
+    l, h, q, t = (
+        model_config.num_hidden_layers,
+        num_heads,
+        model_config.hidden_size // num_heads,
+        seq_len,
+    )
+    # Reasoning behind the factor of 12 for the self-attention part of the formula:
+    # 1. each self-attention has 2 matmul in the forward and 4 in the backward (6)
+    # 2. the flash attention does 1 more matmul recomputation in the backward
+    #    but recomputation should not be counted in calculating MFU           (+0)
+    # 3. each matmul performs 1 multiplication and 1 addition                 (*2)
+    # 4. we follow the convention and do not account for sparsity in causal attention
+    num_flops_per_token = 6 * (nparams - nparams_embedding) + 12 * l * h * q * t
+
+    return nparams, num_flops_per_token

flame/train.py

@@ -0,0 +1,897 @@
+# 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 json
+import os
+import time
+from datetime import timedelta
+from collections import defaultdict
+import dataclasses
+
+import torch
+from datasets import interleave_datasets, load_dataset
+from torch.distributed.elastic.multiprocessing.errors import record
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+import fla  # noqa
+from fla.modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss
+from fla.ops.common.utils import prepare_position_ids
+from flame.components.checkpoint import TrainState
+from flame.config_manager import JobConfig
+from flame.data import build_dataloader, shuffle
+from flame.models.parallelize_fla import parallelize_fla
+from flame.models.pipeline_fla import pipeline_fla
+from flame.tools.utils import get_nparams_and_flops
+from flame.utils.checkpoint import cleanup_local_checkpoints
+from flame.utils.convert_dcp_to_hf import save_pretrained
+from flame.utils.hf_utils import upload_checkpoint_to_hf
+from datetime import datetime
+from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.ft import FTParallelDims, init_ft_manager
+from torchtitan.components.loss import build_cross_entropy_loss
+from torchtitan.components.lr_scheduler import build_lr_schedulers
+from torchtitan.components.metrics import build_device_memory_monitor, build_metrics_processor, ensure_pp_loss_visible
+from torchtitan.components.optimizer import build_optimizers
+from torchtitan.distributed import ParallelDims
+from torchtitan.distributed import utils as dist_utils
+from torchtitan.protocols.model_converter import build_model_converters
+from torchtitan.protocols.train_spec import TrainSpec, get_train_spec, register_train_spec
+from torchtitan.tools import utils
+from torchtitan.tools.logging import init_logger, logger
+from torchtitan.tools.profiling import maybe_enable_memory_snapshot, maybe_enable_profiling
+
+from dotenv import load_dotenv
+load_dotenv()
+
+import wandb
+wandb.login(key=os.environ["WANDB_API_KEY"])
+
+import huggingface_hub
+huggingface_hub.login(token=os.environ["HF_TOKEN"])
+
+
+def build_tokenizer(job_config: JobConfig) -> AutoTokenizer:
+    return AutoTokenizer.from_pretrained(job_config.model.tokenizer_path)
+
+
+register_train_spec(
+    TrainSpec(
+        name="fla",
+        cls=AutoModelForCausalLM,
+        config=AutoConfig,
+        parallelize_fn=parallelize_fla,
+        pipelining_fn=pipeline_fla,
+        build_optimizers_fn=build_optimizers,
+        build_lr_schedulers_fn=build_lr_schedulers,
+        build_dataloader_fn=build_dataloader,
+        build_tokenizer_fn=build_tokenizer,
+        build_loss_fn=build_cross_entropy_loss,
+    )
+)
+
+
+# Enable debug tracing on failure: https://pytorch.org/docs/stable/elastic/errors.html
+@record
+def main(job_config: JobConfig):
+    logger.info(f"Starting job: {job_config.job.description}")
+
+    if job_config.experimental.custom_model_path:
+        utils.import_module_from_path(job_config.experimental.custom_model_path)
+
+    # used for colorful printing
+    color = utils.NoColor if job_config.metrics.disable_color_printing else utils.Color
+
+    if job_config.job.print_args:
+        logger.info(
+            f"{color.green}{json.dumps(job_config.to_dict(), indent=2, sort_keys=True)}{color.reset}"
+        )
+
+    # take control of garbage collection to avoid stragglers
+    gc_handler = utils.GarbageCollection(gc_freq=job_config.training.gc_freq)
+
+    device_module, device_type = utils.device_module, utils.device_type
+    device = torch.device(f"{device_type}:{int(os.environ['LOCAL_RANK'])}")
+    # Device has to be set before creating TorchFT manager.
+    device_module.set_device(device)
+    ft_manager = init_ft_manager(job_config)
+
+    run_specific_repo_id = None
+    if getattr(job_config.checkpoint, "hf_upload_enabled", False):
+        hf_repo_base = getattr(job_config.checkpoint, "hf_repo_base_name", None)
+        if hf_repo_base:
+            # Generate timestamp (adjust format if desired)
+            timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
+            run_specific_repo_id = f"{hf_repo_base}-{timestamp}"
+            logger.info(f"Target Hugging Face repository for this run: {run_specific_repo_id}")
+        else:
+            logger.warning("HF Hub upload enabled, but 'checkpoint.hf_repo_base_name' is not set.")
+            # Disable upload if base name is missing
+            job_config.checkpoint.hf_upload_enabled = False
+
+    # init distributed
+    world_size = int(os.environ["WORLD_SIZE"])
+    if not ft_manager.enabled:
+        parallel_dims = ParallelDims(
+            dp_shard=job_config.training.data_parallel_shard_degree,
+            dp_replicate=job_config.training.data_parallel_replicate_degree,
+            cp=job_config.experimental.context_parallel_degree,
+            tp=job_config.training.tensor_parallel_degree,
+            pp=job_config.experimental.pipeline_parallel_degree,
+            world_size=world_size,
+            enable_loss_parallel=not job_config.training.disable_loss_parallel,
+        )
+    else:
+        parallel_dims = FTParallelDims(
+            dp_shard=job_config.training.data_parallel_shard_degree,
+            dp_replicate=job_config.training.data_parallel_replicate_degree,
+            cp=job_config.experimental.context_parallel_degree,
+            tp=job_config.training.tensor_parallel_degree,
+            pp=job_config.experimental.pipeline_parallel_degree,
+            world_size=world_size,
+            enable_loss_parallel=not job_config.training.disable_loss_parallel,
+            ft_manager=ft_manager,
+        )
+    dist_utils.init_distributed(job_config)
+    # initialize device memory monitor and get peak flops for MFU calculation
+    device_memory_monitor = build_device_memory_monitor()
+    gpu_peak_flops = utils.get_peak_flops(device_memory_monitor.device_name)
+    logger.info(f"Peak FLOPS used for computing MFU: {gpu_peak_flops:.3e}")
+
+    # build meshes
+    world_mesh = parallel_dims.build_mesh(device_type=device_type)
+    if parallel_dims.dp_enabled:
+        dp_mesh = world_mesh["dp"]
+        dp_degree, dp_rank = dp_mesh.size(), dp_mesh.get_local_rank()
+    else:
+        dp_degree, dp_rank = 1, 0
+
+    if parallel_dims.pp_enabled:
+        raise NotImplementedError(
+            "Pipeline parallelism is not supported in this version"
+        )
+        """
+        ! TODO[flame]: We need to fix the pipeline parallelism for flame
+        [x] Match the key of models' components with the actual naming
+        [ ] Fix the post-init and tie-embedding for pipeline parallelism, HF's transformer automatically
+            forces to tie if head is None, we need to handle this case
+        [ ]
+        """
+        pp_mesh = world_mesh["pp"]
+
+    # Set random seed, and maybe enable deterministic mode (mainly for debugging, expect perf loss)
+    dist_utils.set_determinism(
+        world_mesh, device, job_config.training.seed, job_config.training.deterministic
+    )
+    train_spec = get_train_spec(job_config.model.name)
+
+    logger.info("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        job_config.model.tokenizer_path,
+        trust_remote_code=True,
+        model_max_length=int(1e10),
+    )
+    logger.info(f"{tokenizer}")
+    logger.info(
+        f"Loading dataset {job_config.training.dataset}"
+        f":{job_config.training.dataset_name}"
+        if job_config.training.dataset_name is not None
+        else ""
+    )
+
+    min_num_shards = dp_degree * job_config.training.num_workers
+    if len(job_config.training.dataset.split(",")) == 1:
+        dataset = load_dataset(
+            path=job_config.training.dataset,
+            name=getattr(job_config.training, "dataset_name", None),
+            data_dir=getattr(job_config.training, "data_dir", None),
+            data_files=getattr(job_config.training, "data_files", None),
+            split=job_config.training.dataset_split or "train",
+            trust_remote_code=True,
+            streaming=job_config.training.streaming,
+            num_proc=(
+                job_config.training.num_workers
+                if not job_config.training.streaming
+                else None
+            ),
+        )
+        logger.info(f"{dataset}")
+
+        logger.info(f"Shuffling the dataset with seed {job_config.training.seed}")
+        if not job_config.training.streaming:
+            # the states of map-style dataset is recoverable after shuffling
+            dataset = dataset.shuffle(
+                seed=job_config.training.seed
+            ).to_iterable_dataset(num_shards=min_num_shards)
+        else:
+            if dataset.num_shards < min_num_shards:
+                logger.warning(
+                    f"{color.red}"
+                    f"Dataset {job_config.training.dataset} has insufficient shards ({dataset.num_shards}). "
+                    f"Need {min_num_shards} shards minimum for {dp_degree} data parallel workers × "
+                    f"{job_config.training.num_workers} dataloader workers. "
+                    f"Disabling the streaming mode and resharding dataset to {min_num_shards} shards."
+                    f"{color.reset}"
+                )
+                dataset = (
+                    load_dataset(
+                        path=job_config.training.dataset,
+                        name=getattr(job_config.training, "dataset_name", None),
+                        data_dir=getattr(job_config.training, "data_dir", None),
+                        data_files=getattr(job_config.training, "data_files", None),
+                        split=job_config.training.dataset_split or "train",
+                        trust_remote_code=True,
+                        streaming=False,
+                        num_proc=job_config.training.num_workers,
+                    )
+                    .shuffle(seed=job_config.training.seed)
+                    .to_iterable_dataset(num_shards=min_num_shards)
+                )
+            else:
+                dataset = shuffle(dataset, seed=job_config.training.seed)
+    else:
+        datasets = job_config.training.dataset.split(",")
+        if job_config.training.dataset_name is not None:
+            dataset_names = [
+                name or None for name in job_config.training.dataset_name.split(",")
+            ]
+            assert len(dataset_names) == len(datasets), (
+                "The number of dataset names must match the number of datasets"
+            )
+        else:
+            dataset_names = [None] * len(datasets)
+        if job_config.training.dataset_split is not None:
+            dataset_splits = [
+                split or "train"
+                for split in job_config.training.dataset_split.split(",")
+            ]
+            assert len(dataset_splits) == len(datasets), (
+                "The number of dataset splits must match the number of datasets"
+            )
+        else:
+            dataset_splits = ["train"] * len(datasets)
+        if job_config.training.data_dir is not None:
+            data_dirs = [
+                data_dir or None for data_dir in job_config.training.data_dir.split(",")
+            ]
+            assert len(data_dirs) == len(datasets), (
+                "The number of data dirs must match the number of datasets"
+            )
+        else:
+            data_dirs = [None] * len(datasets)
+        if job_config.training.data_files is not None:
+            data_files = job_config.training.data_files.split(",")
+            assert len(data_files) == len(datasets), (
+                "The number of data files must match the number of datasets"
+            )
+        else:
+            data_files = [None] * len(datasets)
+        if job_config.training.data_probs is not None:
+            data_probs = [float(p) for p in job_config.training.data_probs.split(",")]
+            assert len(data_probs) == len(datasets), (
+                "The number of data probabilities must match the number of datasets"
+            )
+        else:
+            raise ValueError(
+                "Data sampling probabilities are required if using multiple datasets"
+            )
+
+        subsets = []
+        for i, prob in enumerate(data_probs):
+            subset = load_dataset(
+                path=datasets[i],
+                name=dataset_names[i],
+                data_dir=data_dirs[i],
+                data_files=data_files[i],
+                split=dataset_splits[i],
+                trust_remote_code=True,
+                streaming=job_config.training.streaming,
+                num_proc=(
+                    job_config.training.num_workers
+                    if not job_config.training.streaming
+                    else None
+                ),
+            )
+            logger.info(
+                f"Subset {color.cyan}{datasets[i]}"
+                + (f":{dataset_names[i]} " if dataset_names[i] else " ")
+                + f"(p = {prob:.3f}){color.reset}:\n"
+                + f"{subset}"
+            )
+
+            logger.info(f"Shuffling the dataset with seed {job_config.training.seed}")
+            if not job_config.training.streaming:
+                # the states of map-style dataset is recoverable after shuffling
+                subset = subset.shuffle(
+                    seed=job_config.training.seed
+                ).to_iterable_dataset(num_shards=min_num_shards)
+            else:
+                if subset.num_shards < min_num_shards:
+                    logger.warning(
+                        f"{color.red}"
+                        f"Dataset {datasets[i]} has insufficient shards ({subset.num_shards}). "
+                        f"Need {min_num_shards} shards minimum for {dp_degree} data parallel workers × "
+                        f"{job_config.training.num_workers} dataloader workers. "
+                        f"Resharding dataset to {min_num_shards} shards and disabling streaming mode."
+                        f"{color.reset}"
+                    )
+                    # again, it's ok to directly shuffle the map-style dataset
+                    # we expect an error raised if the map-style dataset still has not enough data shards
+                    subset = (
+                        load_dataset(
+                            path=datasets[i],
+                            name=dataset_names[i],
+                            data_dir=data_dirs[i],
+                            data_files=data_files[i],
+                            split=dataset_splits[i],
+                            trust_remote_code=True,
+                            streaming=False,
+                            num_proc=job_config.training.num_workers,
+                        )
+                        .shuffle(seed=job_config.training.seed)
+                        .to_iterable_dataset(min_num_shards)
+                    )
+                else:
+                    # we set relatively small buffer size here as interleaving could provide some randomness
+                    subset = shuffle(
+                        subset,
+                        seed=job_config.training.seed,
+                        buffer_size=max(128, 1024 // len(datasets)),
+                    )
+
+            if "text" in subset.column_names:
+                subset = subset.select_columns("text")
+            elif "content" in subset.column_names:
+                subset = subset.select_columns("content")
+            else:
+                raise ValueError(
+                    f"Subset {datasets[i]} has no 'text' or 'content' column"
+                )
+            subsets.append(subset)
+
+        logger.info(
+            f"Interleaving {len(subsets)} datasets with probabilities {data_probs}"
+        )
+        dataset = interleave_datasets(
+            datasets=subsets,
+            probabilities=data_probs,
+            stopping_strategy="all_exhausted",
+            seed=job_config.training.seed,
+        )
+        logger.info(f"{dataset}")
+
+
+    logger.info(f"Loading model config from {job_config.model.config}")
+    model_config = AutoConfig.from_pretrained(job_config.model.config)
+
+    logger.info("Building dataloader...")
+    dataloader = build_dataloader(
+        dataset=dataset,
+        tokenizer=tokenizer,
+        rank=dp_rank,
+        world_size=dp_degree,
+        batch_size=job_config.training.batch_size,
+        # TODO: Make this more modular
+        # seq_len=job_config.training.seq_len if not model_config.use_top_loss else job_config.training.seq_len*2,
+        seq_len=job_config.training.seq_len * 2,
+        context_len=job_config.training.context_len,
+        varlen=job_config.training.varlen,
+        num_workers=job_config.training.num_workers,
+        pin_memory=job_config.training.pin_memory,
+        persistent_workers=job_config.training.persistent_workers,
+        snapshot_every_n_steps=job_config.checkpoint.interval,
+    )
+
+    # set the model configs from training inputs:
+    # 1. norm type to decide which norm layer to use
+    # 2. disable fused norm if TP is enabled
+    # 3. vocab size from tokenizer
+    # 4. context_len base on inputs
+    if parallel_dims.tp_enabled:
+        if model_config.fuse_norm:
+            logger.warning(
+                f"{color.red}"
+                f"Fused norm is not compatible with tensor parallelism. "
+                f"Disabling it for now."
+                f"{color.reset}"
+            )
+            model_config.fuse_norm = False
+    if parallel_dims.loss_parallel_enabled:
+        if model_config.fuse_cross_entropy:
+            logger.warning(
+                f"{color.red}"
+                f"Loss parallel enabled. Disabling fused cross entropy for now."
+                f"{color.reset}"
+            )
+            model_config.fuse_cross_entropy = False
+    model_config.vocab_size = max(tokenizer.vocab_size, model_config.vocab_size)
+
+    logger.info(
+        f"Building model from the config\n{color.green}{model_config}{color.reset}"
+    )
+    with torch.device("meta"):
+        model = AutoModelForCausalLM.from_config(model_config)
+        if (
+            getattr(model_config, "fuse_cross_entropy", False)
+            and FusedLinearCrossEntropyLoss is not None
+        ):
+            model.criterion = FusedLinearCrossEntropyLoss(
+                num_chunks=8 // parallel_dims.tp
+            )
+        # defer weight initialization until after parallelisms are applied
+        model.apply(lambda m: setattr(m, "_is_hf_initialized", False))
+    logger.info(f"{color.blue}\n{model}{color.reset}\n")
+
+    # Build the collection of model converters. No-op if `model.converters` empty
+    model_converters = build_model_converters(job_config, parallel_dims)
+    model_converters.convert(model)
+
+    # calculate model size and flops per token
+    model_param_count, num_flops_per_token = get_nparams_and_flops(
+        model, model_config, job_config.training.context_len
+    )
+
+    # move sharded model to CPU/GPU and initialize weights via DTensor
+    if job_config.checkpoint.create_seed_checkpoint:
+        init_device = "cpu"
+    elif job_config.training.enable_cpu_offload:
+        init_device = "cpu"
+    else:
+        init_device = device_type
+
+    # apply parallelisms and initialization
+    if parallel_dims.pp_enabled:
+        # apply PT-D Pipeline Parallel
+        (
+            pp_schedule,
+            model_parts,
+            has_first_stage,
+            has_last_stage,
+        ) = train_spec.pipelining_fn(
+            model,
+            pp_mesh,
+            parallel_dims,
+            job_config,
+            device,
+            model_config,
+            train_spec.loss_fn,
+        )
+        # when PP is enabled, `model` obj is no longer used after this point, model_parts is used instead
+        del model
+
+        # For PP with looped schedules, each item in model_parts is one stage-model-chunk.
+        # We need to iterate through model_parts to apply SPMD parallelisms, compilation,
+        # optimizer, and checkpointing
+        for m in model_parts:
+            # apply SPMD-style PT-D techniques
+            train_spec.parallelize_fn(m, world_mesh, parallel_dims, job_config)
+            m.to_empty(device=init_device)
+            with torch.no_grad():
+                m.post_init()
+            m.train()
+
+        # confirm that user will be able to view loss metrics on the console
+        ensure_pp_loss_visible(parallel_dims, job_config, color)
+    else:
+        # apply PT-D Tensor Parallel, activation checkpointing, torch.compile, Data Parallel
+        train_spec.parallelize_fn(model, world_mesh, parallel_dims, job_config)
+        model.to_empty(device=init_device)
+        with torch.no_grad():
+            model.post_init()
+        model.train()
+
+        model_parts = [model]
+
+    device_mem_stats = device_memory_monitor.get_peak_stats()
+    logger.info(
+        f"{device_type.upper()} memory usage for model: "
+        f"{device_mem_stats.max_reserved_gib:.2f}GiB"
+        f"({device_mem_stats.max_reserved_pct:.2f}%)"
+    )
+
+    # build optimizer after applying parallelisms to the model
+    optimizers = train_spec.build_optimizers_fn(model_parts, job_config, ft_manager)
+    lr_schedulers = train_spec.build_lr_schedulers_fn(optimizers, job_config)
+    # Post optimizer step model converters hook.
+    # e.g. calculate float8 dynamic amax/scale for all-parameter for FSDP2
+    # where it issues a single all-reduce for all parameters at once for better performance
+    optimizers.register_step_post_hook(
+        lambda *args, **kwargs: model_converters.post_optimizer_hook(model_parts)
+    )
+
+    train_state = TrainState()
+
+    # load initial checkpoint
+    checkpoint = CheckpointManager(
+        dataloader=dataloader,
+        model_parts=model_parts,
+        optimizers=optimizers,
+        lr_schedulers=lr_schedulers,
+        states={"train_state": train_state},
+        job_config=job_config,
+        ft_manager=ft_manager,
+    )
+
+    if job_config.checkpoint.create_seed_checkpoint:
+        assert world_size == 1, (
+            "Must create seed checkpoint using a single device, to disable sharding"
+        )
+        assert job_config.checkpoint.enable_checkpoint, (
+            "Must enable checkpointing when creating a seed checkpoint"
+        )
+        checkpoint.save(curr_step=0, force=True)
+        logger.info("Created seed checkpoint")
+        return
+
+    checkpoint.load(step=job_config.checkpoint.load_step)
+    metric_logger = build_metrics_processor(job_config, parallel_dims)
+    # Set dependent attributes for metric_logger
+    metric_logger.num_flops_per_token = num_flops_per_token
+    metric_logger.optimizers = optimizers  # Pass optimizers if needed by logger logic
+    metric_logger.lr_schedulers = (
+        lr_schedulers  # Pass schedulers if needed by logger logic
+    )
+
+    # plot losses loaded from checkpoint (if any) to TensorBoard
+    # NOTE: Loss info after the last log step before checkpoint saving will not be ploted.
+    #       This can be avoided by setting checkpoint.interval to be a multiple of metrics.log_freq
+    if train_state.step > 0 and len(metric_logger.data_loading_times) > 0:
+        for idx, step in enumerate(train_state.log_steps):
+            metric_logger.log(
+                step,
+                global_avg_loss=train_state.global_avg_losses[idx],
+                global_max_loss=train_state.global_max_losses[idx],
+            )
+
+    data_iterator = iter(dataloader)
+
+    train_context = dist_utils.get_train_context(
+        parallel_dims.loss_parallel_enabled,
+        job_config.experimental.enable_compiled_autograd,
+    )
+
+    # variables used to keep info for metrics logging
+    device_memory_monitor.reset_peak_stats()
+
+    global_batch_size = (
+        job_config.training.batch_size
+        * dp_degree
+        * job_config.training.gradient_accumulation_steps
+    )
+    num_tokens_per_step = global_batch_size * job_config.training.seq_len
+    # train loop
+    logger.info(f"{color.red}***** Running training *****{color.reset}")
+    logger.info(f"{color.green}  Training starts at step {train_state.step + 1}")
+    logger.info(
+        f"{color.green}  Number of tokens per sequence = {job_config.training.seq_len:,}"
+    )
+    logger.info(
+        f"{color.green}  Gradient Accumulation steps = {job_config.training.gradient_accumulation_steps}"
+    )
+    logger.info(
+        f"{color.green}  Instantaneous batch size (per device) = {job_config.training.batch_size:,}"
+    )
+    logger.info(
+        f"{color.green}  Global batch size (w. parallel, distributed & accumulation) = {global_batch_size:,}"
+        f" ({num_tokens_per_step:,} tokens)"
+    )
+    logger.info(
+        f"{color.green}  Total optimization steps = {job_config.training.steps:,} "
+        f"({job_config.training.steps * num_tokens_per_step:,} tokens)"
+    )
+    logger.info(
+        f"{color.green}  Warmup steps = {job_config.lr_scheduler.warmup_steps:,}"
+        f" ({job_config.lr_scheduler.warmup_steps * num_tokens_per_step:,} tokens)"
+    )
+    logger.info(
+        f"{color.green}  Number of parameters = {model_param_count:,} {color.reset}"
+    )
+
+    with (
+        maybe_enable_profiling(
+            job_config, global_step=train_state.step
+        ) as torch_profiler,
+        maybe_enable_memory_snapshot(
+            job_config, global_step=train_state.step
+        ) as memory_profiler,
+    ):
+        while train_state.step < job_config.training.steps:
+            train_state.step += 1
+            gc_handler.run(train_state.step)
+
+            optimizers.zero_grad()
+
+            losses = defaultdict(list)
+            actual_loss = []
+            # do gradient accumulation if enabled
+            for _ in range(job_config.training.gradient_accumulation_steps):
+                # get batch
+                data_load_start = time.perf_counter()
+                batch = next(data_iterator)
+                # Recall that this is, for top and MTP, it will be 
+                # input_ids : (B, seq_len)
+                # labels : (B, seq_len * 2)
+                input_ids, labels = batch["input_ids"][:, :job_config.training.seq_len], batch["labels"]
+
+                # Update metrics processor state before forward/backward
+                metric_logger.ntokens_since_last_log += input_ids.numel()
+                metric_logger.data_loading_times.append(
+                    time.perf_counter() - data_load_start
+                )
+
+                input_ids = input_ids.to(device_type)
+
+                """
+                TODO[flame]: We need to carefully handle the position_ids for TP/CP
+                Depending on the Models'PE, the position_ids might be different.
+
+                e.g. for TP
+                    For RoPE, all ranks have the same position_ids. [FOR HF model]
+                    For sinusoidal, each rank has the coresponding chunked  position_ids. [FOR HF model]
+
+                e.g. for CP, [optional_context_parallel_ctx shoudl automatically distbute the position_ids]
+                    Each rank has the coresponding chunked position_ids. [FOR All model]
+
+                """
+                labels = labels.to(device_type)
+                cu_seqlens = (
+                    batch["cu_seqlens"].to(device_type)
+                    if "cu_seqlens" in batch
+                    else None
+                )
+                if cu_seqlens is not None:
+                    position_ids = prepare_position_ids(cu_seqlens).to(torch.int32)
+                else:
+                    position_ids = (
+                        torch.arange(0, input_ids.shape[1], device=device_type)
+                        .repeat(input_ids.shape[0], 1)
+                        .to(torch.int32)
+                    )
+                # apply context parallelism if cp is enabled
+                # ensure CP handles the separate freqs_cis buffer for each pp stage
+                optional_context_parallel_ctx = (
+                    dist_utils.create_context_parallel_ctx(
+                        cp_mesh=world_mesh["cp"],
+                        cp_buffers=[input_ids, labels, position_ids],
+                        cp_seq_dims=[1, 1, 1],
+                        cp_no_restore_buffers={input_ids, labels, position_ids},
+                        cp_rotate_method=job_config.experimental.context_parallel_rotate_method,
+                    )
+                    if parallel_dims.cp_enabled
+                    else None
+                )
+
+                # #! TODO[flame], we should distribute the position_ids as well with CP
+                if parallel_dims.pp_enabled:
+                    raise NotImplementedError(
+                        "Pipeline parallelism is not supported in this version"
+                    )
+                    # Pipeline Parallel forward / backward inside step() call
+                    with train_context(optional_context_parallel_ctx):
+                        targets, losses = (
+                            (labels, []) if has_last_stage else (None, None)
+                        )
+
+                        if has_first_stage:
+                            pp_schedule.step(input_ids, target=targets, losses=losses)
+                        else:
+                            pp_schedule.step(target=targets, losses=losses)
+
+                    # accumulate losses across pipeline microbatches
+                    # TODO: PP+FSDP unexpectedly puts the loss back to the CPU
+                    loss = (
+                        torch.mean(torch.stack(losses)).to(device)
+                        if has_last_stage
+                        else torch.tensor([-1.0], device=device)
+                    )
+                else:
+                    # Non-PP forward / backward
+                    with train_context(optional_context_parallel_ctx):
+                        output = model(
+                            input_ids=input_ids,
+                            labels=labels,
+                            position_ids=position_ids,
+                            cu_seqlens=cu_seqlens,
+                        )
+                        output_attributes = [field.name for field in dataclasses.fields(output)]
+                        losses_atributes = [x for x in output_attributes if "loss" in x and x != "loss"]
+                        loss = (
+                            output.loss
+                            / job_config.training.gradient_accumulation_steps
+                        )
+                        loss.backward()
+
+                    actual_loss.append(loss)
+                    for loss_attr in losses_atributes:
+                        custom_loss = getattr(output, loss_attr, None)
+                        if custom_loss is not None:
+                            custom_loss = custom_loss / job_config.training.gradient_accumulation_steps
+                            custom_loss = custom_loss
+                            losses[loss_attr].append(custom_loss)
+
+            loss = sum(actual_loss)
+            for loss_attr, loss_values in losses.items():
+                losses[loss_attr] = sum(loss_values)
+
+            # clip gradients
+            grad_norm = dist_utils.clip_grad_norm_(
+                [p for m in model_parts for p in m.parameters()],
+                job_config.training.max_norm,
+                foreach=True,
+                pp_mesh=pp_mesh if parallel_dims.pp_enabled else None,
+            )
+
+            # optimizer step
+            checkpoint.maybe_wait_for_staging()
+            if job_config.training.skip_nan_inf and (
+                grad_norm.isnan() or grad_norm.isinf()
+            ):
+                logger.warning(
+                    f"Skipping optimizer step - detected invalid gradient norm: {grad_norm:.4f}"
+                )
+                optimizers.zero_grad()
+                train_state.skipped_step += 1
+            else:
+                optimizers.step()
+            lr_schedulers.step()
+
+            # log metrics - Use MetricsProcessor
+            global_avg_custom_loss = {}
+            global_max_custom_loss = {}
+            if metric_logger.should_log(train_state.step):
+                if (
+                    parallel_dims.dp_replicate_enabled
+                    or parallel_dims.dp_shard_enabled
+                    or parallel_dims.cp_enabled
+                ):
+                    loss = loss.detach()
+                    # Use dist_mean/max on the accumulated loss for the step
+                    global_avg_loss, global_max_loss = (
+                        dist_utils.dist_mean(
+                            loss,
+                            world_mesh["dp_cp"],
+                        ),
+                        dist_utils.dist_max(
+                            loss,
+                            world_mesh["dp_cp"],
+                        ),
+                    )
+                    for loss_attr, loss_value in losses.items():
+                        global_avg_custom_loss[loss_attr] = dist_utils.dist_mean(
+                            loss_value, world_mesh["dp_cp"]
+                        )
+                        global_max_custom_loss[loss_attr] = dist_utils.dist_max(
+                            loss_value, world_mesh["dp_cp"]
+                        )
+                else:
+                    # Scale back the loss before logging
+                    global_avg_loss = global_max_loss = loss.item()
+                    for loss_attr, loss_value in losses.items():
+                        global_avg_custom_loss[loss_attr] = global_max_custom_loss[
+                            loss_attr
+                        ] = loss_value.item()
+
+                # Update train state tokens and elapsed time
+                time_now = time.perf_counter()
+                time_delta = (
+                    time_now - metric_logger.time_last_log
+                )  # Use metric_logger's time
+                train_state.token += (
+                    metric_logger.ntokens_since_last_log  # Use tokens tracked by metric_logger
+                    * parallel_dims.world_size
+                    / parallel_dims.non_data_parallel_size
+                )
+                train_state.elapsed += timedelta(seconds=time_delta)
+                train_state.log_steps.append(train_state.step)
+                train_state.global_avg_losses.append(global_avg_loss)
+                train_state.global_max_losses.append(global_max_loss)
+
+                # Log using the metric processor
+                last_lr = lr_schedulers.schedulers[0].get_last_lr()[0]
+                eta = (
+                    train_state.elapsed
+                    * (job_config.training.steps - train_state.step)
+                    / train_state.step
+                )
+                extra_metrics = {
+                    "optimizer/lr": last_lr,
+                    "optimizer/grad_norm": grad_norm.item(),
+                    "optimizer/skipped_step": train_state.skipped_step,
+                }
+                for loss_attr, loss_value in global_avg_custom_loss.items():
+                    extra_metrics[f"loss_metrics/global_avg_{loss_attr}"] = loss_value.item() if isinstance(loss_value, torch.Tensor) else loss_value
+                metric_logger.log(
+                    train_state.step,
+                    global_avg_loss,
+                    global_max_loss,
+                    extra_metrics=extra_metrics,
+                )
+
+                logger.info(
+                    f"{color.blue}lr: {last_lr:.4e} gnorm: {grad_norm:5.2f} "
+                    f"{color.magenta}[{str(train_state.elapsed).split('.')[0]:>8}<{str(eta).split('.')[0]:>8}]{color.reset}"
+                )
+
+            checkpoint.save(
+                train_state.step, force=(train_state.step == job_config.training.steps)
+            )
+            
+            if torch.distributed.get_rank() == 0:
+                if job_config.checkpoint.enable_checkpoint:
+                    hf_target_path = None
+                    dcp_save_path = os.path.join(job_config.job.dump_folder, job_config.checkpoint.folder, f"step-{train_state.step}") 
+
+                    # TODO: Haven't tested this one yet
+                    if getattr(job_config.checkpoint, "convert_to_hf_on_save", False):
+                        try:
+                            # Get the path where DCP was just saved
+                            # Check CheckpointManager API for the best way, assuming get_save_path exists
+                            hf_target_path = f"{dcp_save_path}" # e.g., .../checkpoint/step-1000-hf
+
+                            logger.info(f"Converting step {train_state.step} DCP checkpoint to HF format at: {hf_target_path}")
+                            save_pretrained( # Call the imported function
+                                path=hf_target_path, # Pass target HF path as 'path'
+                                step=train_state.step,
+                                config=job_config.model.config, # Pass model config path/id
+                                tokenizer=job_config.model.tokenizer_path # Pass tokenizer path/id
+                            )
+                            logger.info(f"Successfully converted step {train_state.step} to HF format.")
+
+                        except Exception as e:
+                            logger.error(f"Failed to convert checkpoint step {train_state.step} to HF format: {e}", exc_info=True)
+
+                    base_checkpoint_dir = os.path.join(job_config.job.dump_folder, job_config.checkpoint.folder)
+                    if getattr(job_config.checkpoint, "hf_upload_enabled", True):
+                        upload_format = getattr(job_config.checkpoint, "hf_upload_format", "hf")
+                        keep_k_hub = getattr(job_config.checkpoint, "hf_keep_latest_k", 5)
+
+                        local_path_to_upload = None
+                        if upload_format == "hf":
+                            if hf_target_path and os.path.isdir(hf_target_path):
+                                local_path_to_upload = hf_target_path
+                        elif upload_format == "dcp":
+                            if dcp_save_path and os.path.isdir(dcp_save_path):
+                                local_path_to_upload = dcp_save_path
+
+                        if local_path_to_upload:
+                            try:
+                                upload_checkpoint_to_hf(
+                                    local_path=local_path_to_upload,
+                                    step=train_state.step,
+                                    hf_repo_id_for_run=run_specific_repo_id,
+                                    upload_format=upload_format,
+                                    hf_keep_latest_k=job_config.checkpoint.keep_latest_k,
+                                )                               
+                            except Exception as e:
+                                logger.error(f"Failed during HF Hub upload for step {train_state.step}: {e}", exc_info=True)
+
+            # signal the profiler that the next profiling step has started
+            if torch_profiler:
+                torch_profiler.step()
+            if memory_profiler:
+                memory_profiler.step()
+
+            # reduce timeout after first train step for faster signal
+            # (assuming lazy init and compilation are finished)
+            if train_state.step == 1:
+                dist_utils.set_pg_timeouts(
+                    timeout=timedelta(seconds=job_config.comm.train_timeout_seconds),
+                    world_mesh=world_mesh,
+                )
+
+    if torch.distributed.get_rank() == 0:
+        logger.info("Sleeping 2 seconds for other ranks to complete")
+        time.sleep(2)
+
+    metric_logger.close()
+    logger.info("Training completed")
+
+
+if __name__ == "__main__":
+    init_logger()
+    config = JobConfig()
+    config.parse_args()
+    main(config)
+    torch.distributed.destroy_process_group()

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 2,
+  "transformers_version": "4.50.3"
+}

model.safetensors.index.json

@@ -0,0 +1,298 @@
+{
+  "metadata": {
+    "total_size": 7101489152
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00002-of-00002.safetensors",
+    "model.embeddings.weight": "model-00001-of-00002.safetensors",
+    "model.extra_heads.0.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.0.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.0.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.0.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.0.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.0.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.0.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.0.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.0.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.1.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.1.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.1.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.1.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.1.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.1.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.1.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.1.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.1.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.2.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.2.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.2.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.2.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.2.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.2.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.2.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.2.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.2.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.3.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.3.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.3.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.3.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.3.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.3.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.3.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.3.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.extra_heads.3.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.0.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.23.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.23.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.23.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.23.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.23.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.24.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.25.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.26.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.26.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.26.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.26.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.26.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.26.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.27.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.27.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.27.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.27.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.27.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.27.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.3.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "model.norm.weight": "model-00002-of-00002.safetensors"
+  }
+}

pyproject.toml

@@ -0,0 +1,43 @@
+[project]
+name = "flame"
+dynamic = ["version"]
+description = "A minimal training framework for scaling FLA models"
+readme = "README.md"
+authors = [
+    { name = "Songlin Yang", email = "[email protected]" },
+    { name = "Yu Zhang", email = "[email protected]" },
+]
+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.6',
+    'torchdata',
+    'transformers==4.51.3',
+    'triton>=3.0',
+    'datasets>=3.3.0',
+    'einops',
+    'ninja',
+    'wandb',
+    'tiktoken',
+    'tensorboard',
+    'python-dotenv'
+]
+
+[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

setup.py

@@ -0,0 +1,51 @@
+# -*- 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='[email protected], [email protected]',
+    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.6',
+        'torchdata',
+        'transformers==4.51.3',
+        'triton>=3.0',
+        'datasets>=3.3.0',
+        'einops',
+        'ninja',
+        'wandb',
+        'tiktoken',
+        'tensorboard',
+        'python-dotenv'
+    ],
+)

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tb/20250716-2210/wandb/debug-internal.log

@@ -0,0 +1,90 @@
+{"time":"2025-07-16T22:10:00.785425491Z","level":"INFO","msg":"stream: starting","core version":"0.21.0"}
+{"time":"2025-07-16T22:10:01.508654924Z","level":"INFO","msg":"stream: created new stream","id":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201"}
+{"time":"2025-07-16T22:10:01.508690211Z","level":"INFO","msg":"stream: started","id":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201"}
+{"time":"2025-07-16T22:10:01.508739999Z","level":"INFO","msg":"writer: Do: started","stream_id":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201"}
+{"time":"2025-07-16T22:10:01.508759314Z","level":"INFO","msg":"handler: started","stream_id":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201"}
+{"time":"2025-07-16T22:10:01.508803829Z","level":"INFO","msg":"sender: started","stream_id":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201"}
+{"time":"2025-07-16T23:09:45.740737848Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-16T23:18:29.56428269Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": context deadline exceeded"}
+{"time":"2025-07-16T23:19:01.917480335Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-16T23:19:36.868918826Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": context deadline exceeded"}
+{"time":"2025-07-16T23:20:16.297827588Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": context deadline exceeded"}
+{"time":"2025-07-16T23:20:18.619477493Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp: lookup api.wandb.ai on 127.0.0.53:53: read udp 127.0.0.1:46470->127.0.0.53:53: i/o timeout"}
+{"time":"2025-07-16T23:20:30.740650327Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp: lookup api.wandb.ai on 127.0.0.53:53: read udp 127.0.0.1:47482->127.0.0.53:53: i/o timeout"}
+{"time":"2025-07-16T23:21:04.536690541Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-16T23:21:49.291673175Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-16T23:22:07.542159208Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-16T23:23:23.103733736Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-16T23:23:37.543151076Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-16T23:25:07.544031298Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-16T23:26:37.545971769Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-16T23:27:42.194377246Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": net/http: request canceled while waiting for connection (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-16T23:27:59.564813743Z","level":"WARN","msg":"sender: taking a long time","seconds":600.000912631,"work":"WorkRecord(*service_go_proto.Request_StopStatus); Control(local:true  mailbox_slot:\"ft8cf3fgtodg\"  connection_id:\"1(@)\")"}
+{"time":"2025-07-16T23:28:07.547697617Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": context deadline exceeded"}
+{"time":"2025-07-16T23:29:37.549836886Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": context deadline exceeded (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-16T23:31:01.930916994Z","level":"WARN","msg":"runwork: taking a long time","seconds":600.000672411,"work":"WorkRecord(*service_go_proto.Record_Stats); Control(always_send:true)"}
+{"time":"2025-07-16T23:31:02.101966833Z","level":"WARN","msg":"runwork: taking a long time","seconds":600.000995925,"work":"WorkRecord(*service_go_proto.Record_Stats); Control(always_send:true)"}
+{"time":"2025-07-16T23:31:07.103368571Z","level":"WARN","msg":"runwork: taking a long time","seconds":600.000796336,"work":"WorkRecord(*service_go_proto.Request_PartialHistory); Control(local:true  connection_id:\"1(@)\")"}
+{"time":"2025-07-16T23:31:07.551682713Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-16T23:32:37.553473869Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-16T23:33:58.248779065Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": unexpected EOF"}
+{"time":"2025-07-16T23:34:58.351555112Z","level":"INFO","msg":"sender: succeeded after taking longer than expected","seconds":1018.787711083,"work":"WorkRecord(*service_go_proto.Request_StopStatus); Control(local:true  mailbox_slot:\"ft8cf3fgtodg\"  connection_id:\"1(@)\")"}
+{"time":"2025-07-16T23:34:58.351650283Z","level":"INFO","msg":"runwork: succeeded after taking longer than expected","seconds":836.421498346,"work":"WorkRecord(*service_go_proto.Record_Stats); Control(always_send:true)"}
+{"time":"2025-07-16T23:34:58.351778293Z","level":"INFO","msg":"runwork: succeeded after taking longer than expected","seconds":831.249242004,"work":"WorkRecord(*service_go_proto.Request_PartialHistory); Control(local:true  connection_id:\"1(@)\")"}
+{"time":"2025-07-16T23:34:58.351785775Z","level":"INFO","msg":"runwork: succeeded after taking longer than expected","seconds":836.250829923,"work":"WorkRecord(*service_go_proto.Record_Stats); Control(always_send:true)"}
+{"time":"2025-07-17T01:31:13.353253854Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": net/http: request canceled while waiting for connection (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-17T08:06:16.748740406Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-17T09:50:19.526737851Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": read tcp 10.0.2.15:54882->35.186.228.49:443: read: connection reset by peer"}
+{"time":"2025-07-17T09:52:30.348552703Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-17T09:53:02.422139335Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-17T09:53:36.600890938Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": context deadline exceeded"}
+{"time":"2025-07-17T09:54:16.203516351Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": context deadline exceeded"}
+{"time":"2025-07-17T09:55:05.357439477Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": context deadline exceeded"}
+{"time":"2025-07-17T09:56:15.05960959Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-17T09:57:45.061688428Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-17T09:59:15.063226591Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": context deadline exceeded"}
+{"time":"2025-07-17T10:00:45.065259852Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": context deadline exceeded"}
+{"time":"2025-07-17T10:01:04.518171545Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": net/http: request canceled while waiting for connection (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-17T10:02:00.347889757Z","level":"WARN","msg":"sender: taking a long time","seconds":600.000372919,"work":"WorkRecord(*service_go_proto.Request_StopStatus); Control(local:true  mailbox_slot:\"it0uq1ptdf5l\"  connection_id:\"1(@)\")"}
+{"time":"2025-07-17T10:02:15.066174619Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-17T10:03:45.067145051Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-17T10:05:02.098970791Z","level":"WARN","msg":"runwork: taking a long time","seconds":600.000073665,"work":"WorkRecord(*service_go_proto.Record_Stats); Control(always_send:true)"}
+{"time":"2025-07-17T10:05:07.474477054Z","level":"WARN","msg":"runwork: taking a long time","seconds":600.000841939,"work":"WorkRecord(*service_go_proto.Request_PartialHistory); Control(local:true  connection_id:\"1(@)\")"}
+{"time":"2025-07-17T10:05:15.068468165Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": context deadline exceeded"}
+{"time":"2025-07-17T10:05:16.930808745Z","level":"WARN","msg":"runwork: taking a long time","seconds":600.000229861,"work":"WorkRecord(*service_go_proto.Record_Stats); Control(always_send:true)"}
+{"time":"2025-07-17T10:06:07.008582668Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": net/http: request canceled while waiting for connection (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-17T10:06:45.070340311Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-17T10:07:57.799911415Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": unexpected EOF"}
+{"time":"2025-07-17T10:08:57.969386735Z","level":"INFO","msg":"sender: succeeded after taking longer than expected","seconds":1017.621908973,"work":"WorkRecord(*service_go_proto.Request_StopStatus); Control(local:true  mailbox_slot:\"it0uq1ptdf5l\"  connection_id:\"1(@)\")"}
+{"time":"2025-07-17T10:08:57.969579361Z","level":"INFO","msg":"runwork: succeeded after taking longer than expected","seconds":835.870728331,"work":"WorkRecord(*service_go_proto.Record_Stats); Control(always_send:true)"}
+{"time":"2025-07-17T10:08:57.969680501Z","level":"INFO","msg":"runwork: succeeded after taking longer than expected","seconds":821.039158554,"work":"WorkRecord(*service_go_proto.Record_Stats); Control(always_send:true)"}
+{"time":"2025-07-17T10:08:57.969682134Z","level":"INFO","msg":"runwork: succeeded after taking longer than expected","seconds":830.496074059,"work":"WorkRecord(*service_go_proto.Request_PartialHistory); Control(local:true  connection_id:\"1(@)\")"}
+{"time":"2025-07-17T12:53:12.780364188Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-17T16:43:31.998287109Z","level":"INFO","msg":"api: retrying HTTP error","status":502,"url":"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream","body":"\n<html><head>\n<meta http-equiv=\"content-type\" content=\"text/html;charset=utf-8\">\n<title>502 Server Error</title>\n</head>\n<body text=#000000 bgcolor=#ffffff>\n<h1>Error: Server Error</h1>\n<h2>The server encountered a temporary error and could not complete your request.<p>Please try again in 30 seconds.</h2>\n<h2></h2>\n</body></html>\n"}
+{"time":"2025-07-18T00:01:06.015630566Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-18T06:56:24.118529653Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-18T14:32:12.830145916Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-18T19:51:31.703829065Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-07-19T03:35:03.743864446Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-19T21:22:32.639517404Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": read tcp 10.0.2.15:51870->35.186.228.49:443: read: connection reset by peer"}
+{"time":"2025-07-19T21:31:32.643369264Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/graphql\": read tcp 10.0.2.15:38482->35.186.228.49:443: read: connection reset by peer"}
+{"time":"2025-07-20T00:27:42.221361901Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-20T09:40:16.319872482Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-20T09:45:18.218885403Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-20T19:19:37.674808147Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-20T20:26:46.102126738Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-20T21:40:42.245223721Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-20T21:42:31.526229193Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-20T22:42:07.859288654Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-21T03:41:28.397742169Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-21T04:49:16.742257697Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-21T05:48:28.62347913Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-21T06:22:31.529351974Z","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream\": dial tcp 35.186.228.49:443: connect: connection refused"}
+{"time":"2025-07-21T14:47:44.545628902Z","level":"INFO","msg":"api: retrying HTTP error","status":502,"url":"https://api.wandb.ai/files/zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/file_stream","body":"\n<html><head>\n<meta http-equiv=\"content-type\" content=\"text/html;charset=utf-8\">\n<title>502 Server Error</title>\n</head>\n<body text=#000000 bgcolor=#ffffff>\n<h1>Error: Server Error</h1>\n<h2>The server encountered a temporary error and could not complete your request.<p>Please try again in 30 seconds.</h2>\n<h2></h2>\n</body></html>\n"}
+{"time":"2025-07-21T21:19:44.840025606Z","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"}
+{"time":"2025-07-21T21:19:44.94975041Z","level":"INFO","msg":"handler: operation stats","stats":{}}
+{"time":"2025-07-21T21:19:44.958211652Z","level":"INFO","msg":"stream: closing","id":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201"}
+{"time":"2025-07-21T21:19:44.958407771Z","level":"INFO","msg":"writer: Close: closed","stream_id":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201"}
+{"time":"2025-07-21T21:19:44.958426934Z","level":"INFO","msg":"handler: closed","stream_id":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201"}
+{"time":"2025-07-21T21:19:44.958428316Z","level":"INFO","msg":"sender: closed","stream_id":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201"}
+{"time":"2025-07-21T21:19:44.958480192Z","level":"INFO","msg":"stream: closed","id":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201"}

tb/20250716-2210/wandb/debug.log

@@ -0,0 +1,28 @@
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_setup.py:_flush():80] Current SDK version is 0.21.0
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_setup.py:_flush():80] Configure stats pid to 1336753
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_setup.py:_flush():80] Loading settings from /home/cvm/.config/wandb/settings
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_setup.py:_flush():80] Loading settings from /home/cvm/flame/wandb/settings
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_setup.py:_flush():80] Loading settings from environment variables
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_init.py:setup_run_log_directory():703] Logging user logs to exp/mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine/tb/20250716-2210/wandb/run-20250716_221000-mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/logs/debug.log
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_init.py:setup_run_log_directory():704] Logging internal logs to exp/mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine/tb/20250716-2210/wandb/run-20250716_221000-mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/logs/debug-internal.log
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_init.py:init():830] calling init triggers
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_init.py:init():835] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_init.py:init():871] starting backend
+2025-07-16 22:10:00,777 INFO    MainThread:1336753 [wandb_init.py:init():874] sending inform_init request
+2025-07-16 22:10:00,781 INFO    MainThread:1336753 [wandb_init.py:init():882] backend started and connected
+2025-07-16 22:10:00,782 INFO    MainThread:1336753 [wandb_init.py:init():953] updated telemetry
+2025-07-16 22:10:00,786 INFO    MainThread:1336753 [wandb_init.py:init():977] communicating run to backend with 90.0 second timeout
+2025-07-16 22:10:01,926 INFO    MainThread:1336753 [wandb_init.py:init():1029] starting run threads in backend
+2025-07-16 22:10:02,009 INFO    MainThread:1336753 [wandb_run.py:_console_start():2458] atexit reg
+2025-07-16 22:10:02,009 INFO    MainThread:1336753 [wandb_run.py:_redirect():2306] redirect: wrap_raw
+2025-07-16 22:10:02,009 INFO    MainThread:1336753 [wandb_run.py:_redirect():2375] Wrapping output streams.
+2025-07-16 22:10:02,009 INFO    MainThread:1336753 [wandb_run.py:_redirect():2398] Redirects installed.
+2025-07-16 22:10:02,011 INFO    MainThread:1336753 [wandb_init.py:init():1075] run started, returning control to user process
+2025-07-21 21:19:44,102 INFO    MainThread:1336753 [wandb_run.py:_finish():2224] finishing run zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201
+2025-07-21 21:19:44,103 INFO    MainThread:1336753 [wandb_run.py:_atexit_cleanup():2423] got exitcode: 0
+2025-07-21 21:19:44,104 INFO    MainThread:1336753 [wandb_run.py:_restore():2405] restore
+2025-07-21 21:19:44,104 INFO    MainThread:1336753 [wandb_run.py:_restore():2411] restore done
+2025-07-21 21:19:44,955 INFO    MainThread:1336753 [wandb_run.py:_footer_history_summary_info():3903] rendering history
+2025-07-21 21:19:44,956 INFO    MainThread:1336753 [wandb_run.py:_footer_history_summary_info():3935] rendering summary
+2025-07-21 21:19:44,957 INFO    MainThread:1336753 [wandb_run.py:_footer_sync_info():3864] logging synced files

tb/20250716-2210/wandb/run-20250716_221000-mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/files/config.yaml

@@ -0,0 +1,150 @@
+_wandb:
+    value:
+        cli_version: 0.21.0
+        e:
+            ynnjkeia1kakdpk58ub5v7vb16scnioi:
+                args:
+                    - --job.config_file
+                    - flame/models/fla.toml
+                    - --job.dump_folder
+                    - exp/mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine
+                    - --model.config
+                    - configs/mtp_transformer_1B.json
+                    - --model.tokenizer_path
+                    - fla-hub/transformer-1.3B-100B
+                    - --optimizer.name
+                    - AdamW
+                    - --optimizer.eps
+                    - "1e-15"
+                    - --optimizer.lr
+                    - "2e-4"
+                    - --lr_scheduler.warmup_steps
+                    - "2000"
+                    - --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
+                    - "200000"
+                    - --training.max_norm
+                    - "1.0"
+                    - --training.skip_nan_inf
+                    - --training.dataset
+                    - /home/cvm/.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/mtp-1B-4096-batch16x1-steps200000
+                    - --comm.init_timeout_seconds
+                    - "1800"
+                    - --comm.train_timeout_seconds
+                    - "1800"
+                cpu_count: 64
+                cpu_count_logical: 128
+                cudaVersion: "12.8"
+                disk:
+                    /:
+                        total: "3242363822080"
+                        used: "1518440218624"
+                email: [email protected]
+                executable: /home/cvm/miniconda3/envs/flame-env/bin/python3.12
+                git:
+                    commit: aa4d5932e54fad8a568e10aa6895e69e0664fcf1
+                    remote: https://github.com/zaydzuhri/flame.git
+                gpu: NVIDIA H200
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "150754820096"
+                      name: NVIDIA H200
+                      uuid: GPU-eddf9f4c-ffde-5f10-3c76-12ebce1f042b
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "150754820096"
+                      name: NVIDIA H200
+                      uuid: GPU-b532c850-7343-8f67-7eb1-a69024695a99
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "150754820096"
+                      name: NVIDIA H200
+                      uuid: GPU-751a6bdf-72f3-4f5a-fefd-d2b98c338579
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "150754820096"
+                      name: NVIDIA H200
+                      uuid: GPU-0cd9d3c7-1d2e-1925-91eb-8ec99a4ed277
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "150754820096"
+                      name: NVIDIA H200
+                      uuid: GPU-fba7e7ab-8340-13b0-b893-c3686cfec728
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "150754820096"
+                      name: NVIDIA H200
+                      uuid: GPU-12ca11c0-9080-3877-2bd5-3775573a4134
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "150754820096"
+                      name: NVIDIA H200
+                      uuid: GPU-32b3ec8b-9dc8-c6f6-5c19-74fa2ce10ffd
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "150754820096"
+                      name: NVIDIA H200
+                      uuid: GPU-d0021141-e4f4-14ab-c2ab-0ef3e30d6dd5
+                host: mbzuai-2
+                memory:
+                    total: "1913833029632"
+                os: Linux-6.8.0-63-generic-x86_64-with-glibc2.39
+                program: -m flame.train
+                python: CPython 3.12.11
+                root: exp/mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine/tb/20250716-2210
+                startedAt: "2025-07-16T22:10:00.535907Z"
+                writerId: ynnjkeia1kakdpk58ub5v7vb16scnioi
+        m: []
+        python_version: 3.12.11
+        t:
+            "1":
+                - 1
+                - 11
+                - 49
+                - 51
+            "2":
+                - 1
+                - 11
+                - 49
+                - 51
+            "3":
+                - 2
+                - 13
+                - 14
+                - 61
+            "4": 3.12.11
+            "5": 0.21.0
+            "6": 4.50.3
+            "12": 0.21.0
+            "13": linux-x86_64

tb/20250716-2210/wandb/run-20250716_221000-mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/files/requirements.txt

@@ -0,0 +1,101 @@
+flame==0.1.0
+triton==3.2.0
+sympy==1.13.1
+nvidia-cusolver-cu12==11.6.1.9
+idna==3.10
+regex==2024.11.6
+wandb==0.21.0
+nvidia-cuda-cupti-cu12==12.4.127
+protobuf==6.31.1
+Jinja2==3.1.6
+packaging==25.0
+Markdown==3.8.2
+hf-xet==1.1.5
+sentry-sdk==2.33.0
+networkx==3.5
+certifi==2025.7.14
+ninja==1.11.1.4
+PyYAML==6.0.2
+smmap==5.0.2
+numpy==2.3.1
+tiktoken==0.9.0
+nvidia-cuda-nvrtc-cu12==12.4.127
+frozenlist==1.7.0
+tzdata==2025.2
+tokenizers==0.21.2
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-cusparse-cu12==12.3.1.170
+pandas==2.3.1
+attrs==25.3.0
+tensorboard-data-server==0.7.2
+aiohappyeyeballs==2.6.1
+aiosignal==1.4.0
+platformdirs==4.3.8
+python-dotenv==1.1.1
+charset-normalizer==3.4.2
+requests==2.32.4
+MarkupSafe==3.0.2
+GitPython==3.1.44
+nvidia-cufft-cu12==11.2.1.3
+click==8.2.1
+wheel==0.45.1
+nvidia-nccl-cu12==2.21.5
+nvidia-cuda-runtime-cu12==12.4.127
+typing-inspection==0.4.1
+gitdb==4.0.12
+datasets==4.0.0
+multidict==6.6.3
+Werkzeug==3.1.3
+grpcio==1.73.1
+tqdm==4.67.1
+absl-py==2.3.1
+multiprocess==0.70.16
+fsspec==2025.3.0
+dill==0.3.8
+propcache==0.3.2
+yarl==1.20.1
+transformers==4.50.3
+mpmath==1.3.0
+pydantic_core==2.33.2
+flame==0.1.0
+pip==25.1
+torch==2.6.0
+pytz==2025.2
+python-dateutil==2.9.0.post0
+safetensors==0.5.3
+nvidia-curand-cu12==10.3.5.147
+pyarrow==20.0.0
+nvidia-cusparselt-cu12==0.6.2
+einops==0.8.1
+torchdata==0.11.0
+six==1.17.0
+aiohttp==3.12.14
+urllib3==2.5.0
+nvidia-cublas-cu12==12.4.5.8
+filelock==3.18.0
+flash-attn==2.7.3
+nvidia-nvtx-cu12==12.4.127
+xxhash==3.5.0
+tensorboard==2.19.0
+pydantic==2.11.7
+nvidia-cudnn-cu12==9.1.0.70
+typing_extensions==4.14.1
+setuptools==78.1.1
+huggingface-hub==0.33.4
+annotated-types==0.7.0
+jaraco.context==5.3.0
+autocommand==2.2.2
+more-itertools==10.3.0
+tomli==2.0.1
+jaraco.functools==4.0.1
+zipp==3.19.2
+backports.tarfile==1.2.0
+wheel==0.45.1
+platformdirs==4.2.2
+inflect==7.3.1
+typing_extensions==4.12.2
+jaraco.text==3.12.1
+typeguard==4.3.0
+importlib_metadata==8.0.0
+packaging==24.2
+jaraco.collections==5.1.0

tb/20250716-2210/wandb/run-20250716_221000-mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/files/wandb-metadata.json

@@ -0,0 +1,146 @@
+{
+  "os":  "Linux-6.8.0-63-generic-x86_64-with-glibc2.39",
+  "python":  "CPython 3.12.11",
+  "startedAt":  "2025-07-16T22:10:00.535907Z",
+  "args":  [
+    "--job.config_file",
+    "flame/models/fla.toml",
+    "--job.dump_folder",
+    "exp/mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine",
+    "--model.config",
+    "configs/mtp_transformer_1B.json",
+    "--model.tokenizer_path",
+    "fla-hub/transformer-1.3B-100B",
+    "--optimizer.name",
+    "AdamW",
+    "--optimizer.eps",
+    "1e-15",
+    "--optimizer.lr",
+    "2e-4",
+    "--lr_scheduler.warmup_steps",
+    "2000",
+    "--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",
+    "200000",
+    "--training.max_norm",
+    "1.0",
+    "--training.skip_nan_inf",
+    "--training.dataset",
+    "/home/cvm/.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/mtp-1B-4096-batch16x1-steps200000",
+    "--comm.init_timeout_seconds",
+    "1800",
+    "--comm.train_timeout_seconds",
+    "1800"
+  ],
+  "program":  "-m flame.train",
+  "git":  {
+    "remote":  "https://github.com/zaydzuhri/flame.git",
+    "commit":  "aa4d5932e54fad8a568e10aa6895e69e0664fcf1"
+  },
+  "email":  "[email protected]",
+  "root":  "exp/mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine/tb/20250716-2210",
+  "host":  "mbzuai-2",
+  "executable":  "/home/cvm/miniconda3/envs/flame-env/bin/python3.12",
+  "cpu_count":  64,
+  "cpu_count_logical":  128,
+  "gpu":  "NVIDIA H200",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "3242363822080",
+      "used":  "1518440218624"
+    }
+  },
+  "memory":  {
+    "total":  "1913833029632"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H200",
+      "memoryTotal":  "150754820096",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-eddf9f4c-ffde-5f10-3c76-12ebce1f042b"
+    },
+    {
+      "name":  "NVIDIA H200",
+      "memoryTotal":  "150754820096",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-b532c850-7343-8f67-7eb1-a69024695a99"
+    },
+    {
+      "name":  "NVIDIA H200",
+      "memoryTotal":  "150754820096",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-751a6bdf-72f3-4f5a-fefd-d2b98c338579"
+    },
+    {
+      "name":  "NVIDIA H200",
+      "memoryTotal":  "150754820096",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-0cd9d3c7-1d2e-1925-91eb-8ec99a4ed277"
+    },
+    {
+      "name":  "NVIDIA H200",
+      "memoryTotal":  "150754820096",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-fba7e7ab-8340-13b0-b893-c3686cfec728"
+    },
+    {
+      "name":  "NVIDIA H200",
+      "memoryTotal":  "150754820096",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-12ca11c0-9080-3877-2bd5-3775573a4134"
+    },
+    {
+      "name":  "NVIDIA H200",
+      "memoryTotal":  "150754820096",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-32b3ec8b-9dc8-c6f6-5c19-74fa2ce10ffd"
+    },
+    {
+      "name":  "NVIDIA H200",
+      "memoryTotal":  "150754820096",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-d0021141-e4f4-14ab-c2ab-0ef3e30d6dd5"
+    }
+  ],
+  "cudaVersion":  "12.8",
+  "writerId":  "ynnjkeia1kakdpk58ub5v7vb16scnioi"
+}

tb/20250716-2210/wandb/run-20250716_221000-mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/files/wandb-summary.json

@@ -0,0 +1 @@
+{"memory/num_ooms":0,"loss_metrics/global_avg_loss":14.232244491577148,"tflops":431.0192442993196,"_step":200000,"throughput(tps)":31067.90008791172,"loss_metrics/global_avg_ntp_loss":2.0935492515563965,"loss_metrics/global_max_loss":14.996297836303711,"optimizer/grad_norm":1.1316726207733154,"memory/max_reserved(%)":85.28248535606636,"memory/max_active(%)":83.05995444766508,"memory/num_alloc_retries":0,"time_metrics/end_to_end(s)":2.109444146999158,"optimizer/skipped_step":0,"time_metrics/data_loading(%)":0.22571920630939155,"loss_metrics/global_avg_mtp_loss":12.13869571685791,"mfu(%)":43.58131893825275,"memory/max_reserved(GiB)":118.845703125,"_runtime":428982.182300563,"_wandb":{"runtime":428982},"time_metrics/data_loading(s)":0.004761420586146414,"optimizer/lr":2e-05,"_timestamp":1.7531325156953037e+09,"memory/max_active(GiB)":115.74848747253418}

tb/20250716-2210/wandb/run-20250716_221000-mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/logs/debug-core.log

@@ -0,0 +1,16 @@
+{"time":"2025-07-16T22:10:00.603271728Z","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmp4i2mlclc/port-1336753.txt","pid":1336753,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false}
+{"time":"2025-07-16T22:10:00.604042762Z","level":"INFO","msg":"server: will exit if parent process dies","ppid":1336753}
+{"time":"2025-07-16T22:10:00.604005987Z","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-1336753-1346688-770015975/socket","Net":"unix"}}
+{"time":"2025-07-16T22:10:00.768139222Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"}
+{"time":"2025-07-16T22:10:00.785307739Z","level":"INFO","msg":"handleInformInit: received","streamId":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201","id":"1(@)"}
+{"time":"2025-07-16T22:10:01.508694644Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201","id":"1(@)"}
+{"time":"2025-07-21T21:19:44.957947072Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201","id":"1(@)"}
+{"time":"2025-07-21T21:19:44.959492358Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201","id":"1(@)"}
+{"time":"2025-07-21T21:20:04.73362632Z","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"}
+{"time":"2025-07-21T21:20:04.73396049Z","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"}
+{"time":"2025-07-21T21:20:04.733969545Z","level":"INFO","msg":"server is shutting down"}
+{"time":"2025-07-21T21:20:04.734079592Z","level":"INFO","msg":"connection: closing","id":"1(@)"}
+{"time":"2025-07-21T21:20:04.734239368Z","level":"INFO","msg":"connection: closed successfully","id":"1(@)"}
+{"time":"2025-07-21T21:20:04.734245487Z","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"}
+{"time":"2025-07-21T21:20:04.734574344Z","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-1336753-1346688-770015975/socket","Net":"unix"}}
+{"time":"2025-07-21T21:20:04.734618146Z","level":"INFO","msg":"server is closed"}

tb/20250716-2210/wandb/run-20250716_221000-mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/logs/debug.log

@@ -0,0 +1,28 @@
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_setup.py:_flush():80] Current SDK version is 0.21.0
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_setup.py:_flush():80] Configure stats pid to 1336753
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_setup.py:_flush():80] Loading settings from /home/cvm/.config/wandb/settings
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_setup.py:_flush():80] Loading settings from /home/cvm/flame/wandb/settings
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_setup.py:_flush():80] Loading settings from environment variables
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_init.py:setup_run_log_directory():703] Logging user logs to exp/mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine/tb/20250716-2210/wandb/run-20250716_221000-mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/logs/debug.log
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_init.py:setup_run_log_directory():704] Logging internal logs to exp/mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine/tb/20250716-2210/wandb/run-20250716_221000-mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201/logs/debug-internal.log
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_init.py:init():830] calling init triggers
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_init.py:init():835] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2025-07-16 22:10:00,536 INFO    MainThread:1336753 [wandb_init.py:init():871] starting backend
+2025-07-16 22:10:00,777 INFO    MainThread:1336753 [wandb_init.py:init():874] sending inform_init request
+2025-07-16 22:10:00,781 INFO    MainThread:1336753 [wandb_init.py:init():882] backend started and connected
+2025-07-16 22:10:00,782 INFO    MainThread:1336753 [wandb_init.py:init():953] updated telemetry
+2025-07-16 22:10:00,786 INFO    MainThread:1336753 [wandb_init.py:init():977] communicating run to backend with 90.0 second timeout
+2025-07-16 22:10:01,926 INFO    MainThread:1336753 [wandb_init.py:init():1029] starting run threads in backend
+2025-07-16 22:10:02,009 INFO    MainThread:1336753 [wandb_run.py:_console_start():2458] atexit reg
+2025-07-16 22:10:02,009 INFO    MainThread:1336753 [wandb_run.py:_redirect():2306] redirect: wrap_raw
+2025-07-16 22:10:02,009 INFO    MainThread:1336753 [wandb_run.py:_redirect():2375] Wrapping output streams.
+2025-07-16 22:10:02,009 INFO    MainThread:1336753 [wandb_run.py:_redirect():2398] Redirects installed.
+2025-07-16 22:10:02,011 INFO    MainThread:1336753 [wandb_init.py:init():1075] run started, returning control to user process
+2025-07-21 21:19:44,102 INFO    MainThread:1336753 [wandb_run.py:_finish():2224] finishing run zaydzuhri/fla/mtp_transformer-mtp.1B.batch16.seqlen4096.context4096.warmup2000.update1.steps200000.lr2e-4.cosine-202507162201
+2025-07-21 21:19:44,103 INFO    MainThread:1336753 [wandb_run.py:_atexit_cleanup():2423] got exitcode: 0
+2025-07-21 21:19:44,104 INFO    MainThread:1336753 [wandb_run.py:_restore():2405] restore
+2025-07-21 21:19:44,104 INFO    MainThread:1336753 [wandb_run.py:_restore():2411] restore done
+2025-07-21 21:19:44,955 INFO    MainThread:1336753 [wandb_run.py:_footer_history_summary_info():3903] rendering history
+2025-07-21 21:19:44,956 INFO    MainThread:1336753 [wandb_run.py:_footer_history_summary_info():3935] rendering summary
+2025-07-21 21:19:44,957 INFO    MainThread:1336753 [wandb_run.py:_footer_sync_info():3864] logging synced files

tokenizer_config.json

@@ -0,0 +1,44 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [],
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "legacy": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": null,
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

torchtitan/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+#
+# Copyright (c) Meta Platforms, Inc. All Rights Reserved.
+
+# Import to register Float8Converter.
+import torchtitan.components.float8  # noqa: F401
+
+# Import the built-in models here so that the corresponding register_model_spec()
+# will be called.
+import torchtitan.experiments  # noqa: F401
+import torchtitan.models  # noqa: F401

torchtitan/config_manager.py

@@ -0,0 +1,947 @@
+# 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 importlib
+import inspect
+import os
+import sys
+from collections import defaultdict
+from typing import Tuple, Union
+
+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",
+        )
+
+        # 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",
+        )
+
+        # 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",
+        )
+
+        # model configs
+        self.parser.add_argument(
+            "--model.name",
+            type=str,
+            default="llama3",
+            help="Which model to train",
+        )
+        self.parser.add_argument(
+            "--model.flavor",
+            type=str,
+            default="debugmodel",
+            help="Which model config to train",
+        )
+        self.parser.add_argument(
+            "--model.norm_type",
+            type=str,
+            default="rmsnorm",
+            choices=["layernorm", "np_layernorm", "rmsnorm"],
+            help="Type of layer normalization to use [layernorm, np_layernorm, rmsnorm]",
+        )
+        self.parser.add_argument(
+            "--model.use_flex_attn",
+            action="store_true",
+            help="""
+                Whether to use Flex Attention.
+                Mixed usage of SDPA and FlexAttention is not upported yet.
+            """,
+        )
+        self.parser.add_argument(
+            "--model.attn_mask_type",
+            type=str,
+            default="causal",
+            choices=["causal", "block_causal"],
+            help="""
+                Specifies the type of bias/mask used for attention. If SDPA is used,
+                only the causal mask is supported by default. If FlexAttention is used,
+                both causal and block_causal masks are supported.
+            """,
+        )
+        self.parser.add_argument(
+            "--model.tokenizer_path",
+            type=str,
+            default="./assets/tokenizer/original/tokenizer.model",
+            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.
+            """,
+        )
+
+        # optimizer configs
+        self.parser.add_argument(
+            "--optimizer.name", type=str, default="AdamW", help="Optimizer to use"
+        )
+        self.parser.add_argument(
+            "--optimizer.lr", type=float, default=8e-4, help="Learning rate to use"
+        )
+        self.parser.add_argument(
+            "--optimizer.eps", type=float, default=1e-8, help="Epsilon value 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.dataset", type=str, default="c4_test", help="Dataset to use"
+        )
+        self.parser.add_argument(
+            "--training.dataset_path",
+            type=str,
+            help="""
+                Path to the dataset in the file system. If provided, data will be
+                loaded from this path instead of downloaded.""",
+        )
+        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.max_norm",
+            type=Union[float, int],
+            default=1.0,
+            help="Max norm for gradient clipping",
+        )
+        self.parser.add_argument(
+            "--training.steps",
+            type=int,
+            default=10000,
+            help="How many train steps to run",
+        )
+        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.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=None,
+            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",
+        )
+
+        # parallelism configs
+        self.parser.add_argument(
+            "--parallelism.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(
+            "--parallelism.enable_compiled_autograd",
+            action="store_true",
+            help="Enable CompiledAutograd to compile the backward.",
+        )
+        self.parser.add_argument(
+            "--parallelism.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(
+            "--parallelism.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(
+            "--parallelism.tensor_parallel_degree",
+            type=int,
+            default=1,
+            help="Tensor Parallelism degree. 1 means disabled.",
+        )
+        self.parser.add_argument(
+            "--parallelism.disable_loss_parallel",
+            action="store_true",
+            help="Whether to apply loss parallel when sequence parallel is enabled",
+        )
+        self.parser.add_argument(
+            "--parallelism.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(
+            "--parallelism.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(
+            "--parallelism.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(
+            "--parallelism.pipeline_parallel_layers_per_stage",
+            type=int,
+            default=None,
+            help="""
+                The number of layers per stage. If specified, the split points will be calculated from
+                the number of layers and pipeline_parallel_degree. If not specified, the layers per stage will
+                be inferred from the model, schedule, and pipeline_parallel_degree.""",
+        )
+        self.parser.add_argument(
+            "--parallelism.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(
+            "--parallelism.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(
+            "--parallelism.pipeline_parallel_microbatch_size",
+            type=int,
+            default=1,
+            help="""
+                The size of each pipeline parallel microbatch (default 1).
+
+                This value is used to compute the total number of microbatches by dividing batch_size with
+                pipeline_parallel_microbatch_size.
+
+                The global training batch size must be evenly divisible by pipeline_parallel_microbatch_size.
+            """,
+        )
+        self.parser.add_argument(
+            "--parallelism.context_parallel_degree",
+            type=int,
+            default=1,
+            help="Context parallelism degree. 1 means disabled.",
+        )
+        self.parser.add_argument(
+            "--parallelism.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'.
+            """,
+        )
+
+        # 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=10,
+            help="""
+                Keeps only the latest k checkpoints, and purging older ones. If 0, keep all checkpoints.
+                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. The default value is 10 to avoid
+                filling up the disk.
+            """,
+        )
+        self.parser.add_argument(
+            "--checkpoint.load_step",
+            type=int,
+            default=-1,
+            help="Load the checkpoint at the specified step. If -1, load the latest checkpoint.",
+        )
+        self.parser.add_argument(
+            "--checkpoint.exclude_from_loading",
+            type=string_list,
+            nargs="*",
+            default=[],
+            help="""
+                Exclude specific keys from being loaded from the checkpoint.
+                Provide a comma-separated list of keys to exclude, e.g. 'optimizer,lr_scheduler,dataloader'.
+                This will load the model only, excluding the specified keys.
+            """,
+        )
+
+        # activation checkpointing configs
+        self.parser.add_argument(
+            "--activation_checkpoint.mode",
+            type=str,
+            default="selective",
+            help="Type of activation checkpointing to use ['none', 'full', 'selective']",
+        )
+        self.parser.add_argument(
+            "--activation_checkpoint.selective_ac_option",
+            type=str,
+            default="2",  # 2 = checkpoint every other layer
+            help="""
+                Selective activation checkpointing options ['int', 'op'].
+                'int' (e.g., 2) for every nth layer, or 'op' for op level ac.
+            """,
+        )
+
+        # 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`.
+            """,
+        )
+        self.parser.add_argument(
+            "--float8.filter_fqns",
+            type=string_list,
+            default=[],
+            nargs="+",
+            help="""
+            Comma-separated list of fully qualified names of modules to skip applying float8 training to.
+            nn.Linear modules with any dim size not divisible by 16 are always skipped due to hardware requirements.
+            Example: --float8.module_filter_fqns "attention.wq,attention.wk,attention.wv,output"
+            """,
+        )
+
+        # 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 configs
+        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.
+            """,
+        )
+
+        # torchft configs
+        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.",
+        )
+
+        self.parser.add_argument(
+            "--experimental.custom_import",
+            type=str,
+            default="",
+            help="""
+            This option enables the importation of external modules.
+            Currently, it only supports dotted import modules (e.g., some_package.model_x).
+            It is the user's responsibility to ensure that the specified path can be
+            successfully imported. One method to achieve this, you can place your module
+            inside the ``torchtitan/torchtitan`` folder and execute ``pip install -e .`` to
+            make it available for import.
+            """,
+        )
+
+        self.parser.add_argument(
+            "--experimental.custom_args_module",
+            type=str,
+            default="",
+            help="""
+                This option allows users to extend TorchTitan's existing JobConfig by importing
+                a customized module. Similar to ``--experimental.custom_model_path``, the user
+                needs to ensure that the path can be imported. The module should contain exactly
+                one public function and the function has the signature
+                ``def func(parser: argparse.ArgumentParser) -> None:``. The user can use the
+                given parser to add new argument by calling``parser.add_argument``, as wish.
+            """,
+        )
+
+        self._is_parsed = False
+        self._allow_unkown_args = False
+
+    def maybe_add_custom_args(self) -> None:
+        """Add custom arguments to the parser if --experimental.custom_args_module is set.
+
+        Note: This function should be called before the parser is used to parse arguments.
+        """
+        if self._is_parsed:
+            raise RuntimeError(
+                "JobConfig has already been parsed. We could not add new arguments."
+            )
+
+        self._allow_unkown_args = True
+        self.parse_args(sys.argv[1:])
+        self._allow_unkown_args = False
+
+        if self.experimental.custom_args_module:
+            module = importlib.import_module(self.experimental.custom_args_module)
+            public_functions = [
+                name
+                for name, func in inspect.getmembers(module)
+                if inspect.isfunction(func) and not name.startswith("_")
+            ]
+            func = getattr(module, public_functions[0])
+            func(self.parser)
+
+    def to_dict(self):
+        return self.args_dict
+
+    def parse_args(self, args_list: list = sys.argv[1:]):
+        self._is_parsed = True
+        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: temporary mitigation of BC breaking change in
+        #       tokenizer default path, need to remove later
+        if not os.path.exists(self.model.tokenizer_path):
+            logger.warning(
+                f"Tokenizer path {self.model.tokenizer_path} does not exist!"
+            )
+            old_tokenizer_path = (
+                "torchtitan/datasets/tokenizer/original/tokenizer.model"
+            )
+            if os.path.exists(old_tokenizer_path):
+                self.model.tokenizer_path = old_tokenizer_path
+                logger.warning(
+                    f"Temporarily switching to previous default tokenizer path {old_tokenizer_path}. "
+                    "Please update your config."
+                )
+
+    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
+        """
+        if self._allow_unkown_args:
+            args, _ = self.parser.parse_known_args(args_list)
+        else:
+            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

torchtitan/train.py

@@ -0,0 +1,482 @@
+# 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 importlib
+import os
+import time
+from datetime import timedelta
+from typing import Any, Generator, Iterable, Optional
+
+import torch
+from torch.distributed.elastic.multiprocessing.errors import record
+
+import torchtitan.components.ft as ft
+import torchtitan.protocols.train_spec as train_spec_module
+
+from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.metrics import (
+    build_metrics_processor,
+    ensure_pp_loss_visible,
+)
+from torchtitan.config_manager import JobConfig
+from torchtitan.distributed import ParallelDims, utils as dist_utils
+from torchtitan.protocols.model_converter import build_model_converters
+from torchtitan.tools import utils
+from torchtitan.tools.logging import init_logger, logger
+from torchtitan.tools.profiling import (
+    maybe_enable_memory_snapshot,
+    maybe_enable_profiling,
+)
+
+
+class Trainer(torch.distributed.checkpoint.stateful.Stateful):
+    job_config: JobConfig
+    gc_handler: utils.GarbageCollection
+
+    parallel_dims: ParallelDims
+    train_spec: train_spec_module.TrainSpec
+    world_mesh: torch.distributed.DeviceMesh
+
+    dataloader: train_spec_module.BaseDataLoader
+    metrics_processor: train_spec_module.MetricsProcessor
+    checkpointer: CheckpointManager
+    train_context: Generator[None, None, None]
+
+    model_parts: list[torch.nn.Module]
+    loss_fn: train_spec_module.LossFunction
+    optimizers: train_spec_module.OptimizersContainer
+    lr_schedulers: train_spec_module.LRSchedulersContainer
+
+    pp_has_first_stage: bool
+    pp_has_last_stage: bool
+
+    device: torch.device
+
+    # states
+    step: int
+
+    # Enable debug tracing on failure: https://pytorch.org/docs/stable/elastic/errors.html
+    @record
+    def __init__(self, job_config: JobConfig):
+        self.job_config = job_config
+
+        logger.info(f"Starting job: {job_config.job.description}")
+
+        if job_config.experimental.custom_import:
+            importlib.import_module(job_config.experimental.custom_import)
+
+        if job_config.job.print_args:
+            logger.info(f"Running with args: {job_config.to_dict()}")
+
+        # take control of garbage collection to avoid stragglers
+        self.gc_handler = utils.GarbageCollection(gc_freq=job_config.training.gc_freq)
+
+        device_module, device_type = utils.device_module, utils.device_type
+        self.device = torch.device(f"{device_type}:{int(os.environ['LOCAL_RANK'])}")
+        # Device has to be set before creating TorchFT manager.
+        device_module.set_device(self.device)
+        ft_manager = ft.init_ft_manager(job_config)
+
+        # init distributed
+        world_size = int(os.environ["WORLD_SIZE"])
+        parallelism_config = job_config.parallelism
+        if not ft_manager.enabled:
+            self.parallel_dims = parallel_dims = ParallelDims(
+                dp_shard=parallelism_config.data_parallel_shard_degree,
+                dp_replicate=parallelism_config.data_parallel_replicate_degree,
+                cp=parallelism_config.context_parallel_degree,
+                tp=parallelism_config.tensor_parallel_degree,
+                pp=parallelism_config.pipeline_parallel_degree,
+                world_size=world_size,
+                enable_loss_parallel=not parallelism_config.disable_loss_parallel,
+            )
+        else:
+            self.parallel_dims = parallel_dims = ft.FTParallelDims(
+                dp_shard=parallelism_config.data_parallel_shard_degree,
+                dp_replicate=parallelism_config.data_parallel_replicate_degree,
+                cp=parallelism_config.context_parallel_degree,
+                tp=parallelism_config.tensor_parallel_degree,
+                pp=parallelism_config.pipeline_parallel_degree,
+                world_size=world_size,
+                enable_loss_parallel=not parallelism_config.disable_loss_parallel,
+                ft_manager=ft_manager,
+            )
+        dist_utils.init_distributed(job_config)
+
+        # build meshes
+        self.world_mesh = world_mesh = parallel_dims.build_mesh(device_type=device_type)
+        if parallel_dims.dp_enabled:
+            dp_mesh = world_mesh["dp"]
+            dp_degree, dp_rank = dp_mesh.size(), dp_mesh.get_local_rank()
+        else:
+            dp_degree, dp_rank = 1, 0
+
+        # Set random seed, and maybe enable deterministic mode
+        # (mainly for debugging, expect perf loss).
+        dist_utils.set_determinism(
+            world_mesh,
+            self.device,
+            job_config.training.seed,
+            job_config.training.deterministic,
+        )
+        self.train_spec = train_spec_module.get_train_spec(job_config.model.name)
+
+        # build dataloader
+        tokenizer = (
+            self.train_spec.build_tokenizer_fn(job_config)
+            if self.train_spec.build_tokenizer_fn is not None
+            else None
+        )
+
+        # If TorchFT is enabled, the dp_rank and dp_degree, which are used for
+        # dataloader must be changed.
+        if ft_manager.enabled:
+            dp_degree, dp_rank = ft_manager.get_dp_info(dp_degree, dp_rank)
+
+        self.dataloader = self.train_spec.build_dataloader_fn(
+            dp_world_size=dp_degree,
+            dp_rank=dp_rank,
+            tokenizer=tokenizer,
+            job_config=job_config,
+        )
+
+        # build model (using meta init)
+        model_cls = self.train_spec.cls
+        model_args = self.train_spec.config[job_config.model.flavor]
+        # set the model args from training job configs
+        model_args.update_from_config(job_config, tokenizer)
+
+        logger.info(
+            f"Building {self.train_spec.name} {job_config.model.flavor} with {model_args}"
+        )
+        with torch.device("meta"):
+            model = model_cls.from_model_args(model_args)
+
+        # Build the collection of model converters. No-op if `model.converters` empty
+        model_converters = build_model_converters(job_config, parallel_dims)
+        model_converters.convert(model)
+
+        # metrics logging
+        build_metrics_processor_fn = (
+            build_metrics_processor
+            if self.train_spec.build_metrics_processor_fn is None
+            else self.train_spec.build_metrics_processor_fn
+        )
+        self.metrics_processor = build_metrics_processor_fn(job_config, parallel_dims)
+        color = self.metrics_processor.color
+
+        # calculate model size and flops per token
+        (
+            model_param_count,
+            self.metrics_processor.num_flops_per_token,
+        ) = model_args.get_nparams_and_flops(model, job_config.training.seq_len)
+
+        logger.info(
+            f"{color.blue}Model {self.train_spec.name} {job_config.model.flavor} "
+            f"{color.red}size: {model_param_count:,} total parameters{color.reset}"
+        )
+
+        # move sharded model to CPU/GPU and initialize weights via DTensor
+        if job_config.checkpoint.create_seed_checkpoint:
+            init_device = "cpu"
+            buffer_device = None
+        elif job_config.training.enable_cpu_offload:
+            init_device = "cpu"
+            buffer_device = device_type
+        else:
+            init_device = device_type
+            buffer_device = None
+
+        self.loss_fn = self.train_spec.build_loss_fn(job_config)
+
+        # apply parallelisms and initialization
+        if parallel_dims.pp_enabled:
+            if not self.train_spec.pipelining_fn:
+                raise RuntimeError(
+                    f"Pipeline Parallel is enabled but {self.train_spec.name} "
+                    f"does not support pipelining"
+                )
+
+            # apply both PT-D Pipeline Parallel and SPMD-style PT-D techniques
+            (
+                self.pp_schedule,
+                self.model_parts,
+                self.pp_has_first_stage,
+                self.pp_has_last_stage,
+            ) = self.train_spec.pipelining_fn(
+                model,
+                world_mesh,
+                parallel_dims,
+                job_config,
+                self.device,
+                model_args,
+                self.train_spec.parallelize_fn,
+                self.loss_fn,
+            )
+            # when PP is enabled, `model` obj is no longer used after this point,
+            # model_parts is used instead
+            del model
+
+            for m in self.model_parts:
+                m.to_empty(device=init_device)
+                with torch.no_grad():
+                    m.init_weights(buffer_device=buffer_device)
+                m.train()
+
+            # confirm that user will be able to view loss metrics on the console
+            ensure_pp_loss_visible(parallel_dims, job_config, color)
+        else:
+            # apply PT-D Tensor Parallel, activation checkpointing, torch.compile, Data Parallel
+            model = self.train_spec.parallelize_fn(
+                model, world_mesh, parallel_dims, job_config
+            )
+
+            model.to_empty(device=init_device)
+            with torch.no_grad():
+                model.init_weights(buffer_device=buffer_device)
+            model.train()
+
+            self.model_parts = [model]
+
+        # initialize device memory monitor and get peak flops for MFU calculation
+        device_memory_monitor = self.metrics_processor.device_memory_monitor
+        gpu_peak_flops = utils.get_peak_flops(device_memory_monitor.device_name)
+        logger.info(f"Peak FLOPS used for computing MFU: {gpu_peak_flops:.3e}")
+        device_mem_stats = device_memory_monitor.get_peak_stats()
+        logger.info(
+            f"{device_type.upper()} memory usage for model: "
+            f"{device_mem_stats.max_reserved_gib:.2f}GiB"
+            f"({device_mem_stats.max_reserved_pct:.2f}%)"
+        )
+
+        # build optimizer after applying parallelisms to the model
+        self.optimizers = self.train_spec.build_optimizers_fn(
+            self.model_parts, job_config, ft_manager
+        )
+        self.lr_schedulers = self.train_spec.build_lr_schedulers_fn(
+            self.optimizers, job_config
+        )
+        # Post optimizer step model converters hook.
+        # e.g. calculate float8 dynamic amax/scale for all-parameter for FSDP2
+        # where it issues a single all-reduce for all parameters at once for better performance
+        self.optimizers.register_step_post_hook(
+            lambda *args, **kwargs: model_converters.post_optimizer_hook(
+                self.model_parts
+            )
+        )
+        self.metrics_processor.optimizers = self.optimizers
+
+        # Initialize trainer states that will be saved in checkpoint.
+        # These attributes must be initialized before checkpoint loading.
+        self.step = 0
+
+        self.checkpointer = CheckpointManager(
+            dataloader=self.dataloader,
+            model_parts=self.model_parts,
+            optimizers=self.optimizers,
+            lr_schedulers=self.lr_schedulers,
+            states={"train_state": self},
+            job_config=job_config,
+            ft_manager=ft_manager,
+        )
+
+        self.train_context = dist_utils.get_train_context(
+            parallel_dims.loss_parallel_enabled,
+            parallelism_config.enable_compiled_autograd,
+        )
+
+        logger.info(
+            "Trainer is initialized with "
+            f"local batch size {job_config.training.batch_size}, "
+            f"global batch size {job_config.training.batch_size * dp_degree}, "
+            f"sequence length {job_config.training.seq_len}, "
+            f"total steps {job_config.training.steps} "
+            f"(warmup {job_config.lr_scheduler.warmup_steps})."
+        )
+
+    def next_batch(
+        self, data_iterator: Iterable
+    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+        data_load_start = time.perf_counter()
+        batch = next(data_iterator)
+        input_dict, labels = batch
+        self.metrics_processor.ntokens_since_last_log += labels.numel()
+        self.metrics_processor.data_loading_times.append(
+            time.perf_counter() - data_load_start
+        )
+
+        device_type = utils.device_type
+        for k, _ in input_dict.items():
+            input_dict[k] = input_dict[k].to(device_type)
+        labels = labels.to(device_type)
+        return input_dict, labels
+
+    def train_step(self, input_dict: dict[str, torch.Tensor], labels: torch.Tensor):
+        self.optimizers.zero_grad()
+
+        # Keep these variables local to shorten the code as these are
+        # the major variables that are used in the training loop.
+        model_parts = self.model_parts
+        world_mesh = self.world_mesh
+        parallel_dims = self.parallel_dims
+
+        # apply context parallelism if cp is enabled
+        # ensure CP handles the separate freqs_cis buffer for each pp stage
+        inputs = input_dict["input"]
+        optional_context_parallel_ctx = (
+            dist_utils.create_context_parallel_ctx(
+                cp_mesh=world_mesh["cp"],
+                cp_buffers=[inputs, labels] + [m.freqs_cis for m in model_parts],
+                cp_seq_dims=[1, 1] + [0 for _ in model_parts],
+                cp_no_restore_buffers={inputs, labels},
+                cp_rotate_method=self.job_config.parallelism.context_parallel_rotate_method,
+            )
+            if parallel_dims.cp_enabled
+            else None
+        )
+
+        if parallel_dims.pp_enabled:
+            # Pipeline Parallel forward / backward inside step() call
+            with self.train_context(optional_context_parallel_ctx):
+                targets, losses = (
+                    (labels, []) if self.pp_has_last_stage else (None, None)
+                )
+                if self.pp_has_first_stage:
+                    self.pp_schedule.step(inputs, target=targets, losses=losses)
+                else:
+                    self.pp_schedule.step(target=targets, losses=losses)
+
+            # accumulate losses across pipeline microbatches
+            # TODO: PP+FSDP unexpectedly puts the loss back to the CPU
+            loss = (
+                torch.mean(torch.stack(losses)).to(self.device)
+                if self.pp_has_last_stage
+                else torch.tensor([-1.0], device=self.device)
+            )
+        else:
+            # Non-PP forward / backward
+            with self.train_context(optional_context_parallel_ctx):
+                assert len(model_parts) == 1
+                pred = model_parts[0](inputs)
+                loss = self.loss_fn(pred, labels)
+                # pred.shape=(bs, seq_len, vocab_size)
+                # need to free to before bwd to avoid peaking memory
+                del pred
+                loss.backward()
+
+        dist_utils.clip_grad_norm_(
+            [p for m in model_parts for p in m.parameters()],
+            self.job_config.training.max_norm,
+            foreach=True,
+            pp_mesh=self.world_mesh["pp"] if parallel_dims.pp_enabled else None,
+        )
+        self.checkpointer.maybe_wait_for_staging()
+        self.optimizers.step()
+        self.lr_schedulers.step()
+
+        # log metrics
+        if not self.metrics_processor.should_log(self.step):
+            return
+
+        if (
+            parallel_dims.dp_replicate_enabled
+            or parallel_dims.dp_shard_enabled
+            or parallel_dims.cp_enabled
+        ):
+            loss = loss.detach()
+            global_avg_loss, global_max_loss = (
+                dist_utils.dist_mean(loss, world_mesh["dp_cp"]),
+                dist_utils.dist_max(loss, world_mesh["dp_cp"]),
+            )
+        else:
+            global_avg_loss = global_max_loss = loss.item()
+
+        self.metrics_processor.log(self.step, global_avg_loss, global_max_loss)
+
+    @record
+    def train(self):
+        job_config = self.job_config
+
+        self.checkpointer.load(step=job_config.checkpoint.load_step)
+        logger.info(f"Training starts at step {self.step + 1}.")
+
+        with maybe_enable_profiling(
+            job_config, global_step=self.step
+        ) as torch_profiler, maybe_enable_memory_snapshot(
+            job_config, global_step=self.step
+        ) as memory_profiler:
+            data_iterator = iter(self.dataloader)
+            while self.step < job_config.training.steps:
+                self.step += 1
+                self.gc_handler.run(self.step)
+                inputs, labels = self.next_batch(data_iterator)
+                self.train_step(inputs, labels)
+                self.checkpointer.save(
+                    self.step, force=(self.step == job_config.training.steps)
+                )
+
+                # signal the profiler that the next profiling step has started
+                if torch_profiler:
+                    torch_profiler.step()
+                if memory_profiler:
+                    memory_profiler.step()
+
+                # reduce timeout after first train step for faster signal
+                # (assuming lazy init and compilation are finished)
+                if self.step == 1:
+                    dist_utils.set_pg_timeouts(
+                        timeout=timedelta(
+                            seconds=job_config.comm.train_timeout_seconds
+                        ),
+                        world_mesh=self.world_mesh,
+                    )
+
+        if torch.distributed.get_rank() == 0:
+            logger.info("Sleeping 2 seconds for other ranks to complete")
+            time.sleep(2)
+
+        self.metrics_processor.close()
+        logger.info("Training completed")
+
+    def state_dict(self) -> dict[str, Any]:
+        return {"step": self.step}
+
+    def load_state_dict(self, state_dict: dict[str, Any]):
+        self.step = state_dict["step"]
+
+    def close(self) -> None:
+        if self.checkpointer:
+            self.checkpointer.close()
+
+
+if __name__ == "__main__":
+    init_logger()
+    config = JobConfig()
+    config.maybe_add_custom_args()
+    config.parse_args()
+    trainer: Optional[Trainer] = None
+
+    try:
+        trainer = Trainer(config)
+
+        if config.checkpoint.create_seed_checkpoint:
+            assert int(
+                os.environ["WORLD_SIZE"]
+            ), "Must create seed checkpoint using a single device, to disable sharding."
+            assert (
+                config.checkpoint.enable_checkpoint
+            ), "Must enable checkpointing when creating a seed checkpoint."
+            trainer.checkpointer.save(curr_step=0, force=True)
+            logger.info("Created seed checkpoint")
+        else:
+            trainer.train()
+    finally:
+        if trainer:
+            trainer.close()
+
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()
+            logger.info("Process group destroyed.")

train.sh

@@ -0,0 +1,121 @@
+#!/usr/bin/bash
+
+params=""
+if [ $# -ne 0 ]; then
+    params="$*"
+fi
+
+# use envs as local params for convenience
+# e.g.
+# NNODE=1 NGPU=8 LOG_RANK=0 ./train.sh
+NNODE=${NNODE:-"1"}
+NGPU=${NGPU:-"8"}
+LOG_RANK=${LOG_RANK:-0}
+
+if [[ -z "${MASTER_ADDR}" ]]; then
+  export MASTER_ADDR="localhost"
+fi
+if [[ -z "${MASTER_PORT}" ]]; then
+  export MASTER_PORT="0"
+fi
+
+: '
+Usage:
+
+bash train.sh -h
+
+Training a 340M model:
+
+NNODE=1 NGPU=8 LOG_RANK=0 bash train.sh \
+  --job.config_file flame/models/fla.toml \
+  --job.dump_folder exp/transformer-340M-10B/batch32.seqlen2048.warmup1024.update1.steps20480.lr3e-4 \
+  --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 32 \
+  --training.seq_len 2048 \
+  --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 default \
+  --training.dataset_split train \
+  --training.streaming \
+  --training.num_workers 32 \
+  --training.prefetch_factor 2 \
+  --training.seed 42 \
+  --training.compile \
+  --training.tensor_parallel_degree 1 \
+  --training.disable_loss_parallel \
+  --checkpoint.interval 2048 \
+  --checkpoint.load_step -1 \
+  --metrics.log_freq 1
+'
+
+echo "Launching training..."
+
+set -x
+path=$(grep -oP '(?<=--job.dump_folder )[^ ]+' <<< "$params")
+steps=$(grep -oP '(?<=--training.steps )[^ ]+' <<< "$params")
+config=$(grep -oP '(?<=--model.config )[^ ]+' <<< "$params")
+tokenizer=$(grep -oP '(?<=--model.tokenizer_path )[^ ]+' <<< "$params")
+model=$(
+  python -c "import fla, sys; from transformers import AutoConfig; print(AutoConfig.from_pretrained(sys.argv[1]).to_json_string())" "$config" | jq -r '.model_type'
+)
+
+mkdir -p $path
+cp * $path
+cp -r configs $path
+cp -r flame   $path
+cp -r 3rdparty/flash-linear-attention/fla $path
+cp -r 3rdparty/torchtitan/torchtitan $path
+
+# for offline systems
+# export TRANSFORMERS_OFFLINE=1
+# export HF_DATASETS_OFFLINE=1
+# export HF_HUB_OFFLINE=1
+if [ "$date" == "" ]; then
+  date=$(date +%Y%m%d%H%M)
+fi
+RUN_NAME="$model-$(basename $path)"
+RUN_ID="$RUN_NAME-$date"
+
+export WANDB_RESUME=allow
+if [[ -z "${WANDB_PROJECT}" ]]; then
+  export WANDB_PROJECT="fla"
+fi
+if [[ -z "${WANDB_NAME}" ]]; then
+  export WANDB_NAME="$RUN_NAME"
+fi
+if [[ -z "${WANDB_RUN_ID}" ]]; then
+  export WANDB_RUN_ID="$RUN_ID"
+fi
+
+PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True" \
+torchrun --nnodes=${NNODE} \
+  --nproc_per_node=${NGPU} \
+  --rdzv_backend c10d \
+  --rdzv_endpoint "${MASTER_ADDR}:${MASTER_PORT}" \
+  --local-ranks-filter ${LOG_RANK} \
+  --role rank \
+  --tee 3 \
+  --log-dir $path/logs \
+  -m flame.train \
+  $params
+
+echo "TRAINING DONE!"
+echo "Converting the DCP checkpoints to HF format..."
+
+python -m flame.utils.convert_dcp_to_hf \
+  --path $path \
+  --step $steps \
+  --config $config \
+  --tokenizer $tokenizer
+
+echo "RUNNING DONE!"