Delete gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/LICENSE

@@ -1,21 +0,0 @@
-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.

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/README.md

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/checkpoint/step-30720/.metadata

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/checkpoint/step-30720/__0_0.distcp

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/checkpoint/step-30720/__1_0.distcp

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/checkpoint/step-30720/__2_0.distcp

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/checkpoint/step-30720/__3_0.distcp

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/checkpoint/step-30720/__4_0.distcp

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/checkpoint/step-30720/__5_0.distcp

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/checkpoint/step-30720/__6_0.distcp

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/checkpoint/step-30720/__7_0.distcp

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/config.json

@@ -1,36 +0,0 @@
-{
-  "allow_neg_eigval": false,
-  "architectures": [
-    "GatedDeltaNetForCausalLM"
-  ],
-  "attn": null,
-  "attn_mode": "chunk",
-  "bos_token_id": 1,
-  "conv_size": 4,
-  "dtype": "float32",
-  "eos_token_id": 2,
-  "expand_v": 2,
-  "fuse_cross_entropy": true,
-  "fuse_linear_cross_entropy": false,
-  "fuse_norm": true,
-  "fuse_swiglu": true,
-  "head_dim": 256,
-  "hidden_act": "swish",
-  "hidden_ratio": 4,
-  "hidden_size": 2048,
-  "initializer_range": 0.02,
-  "intermediate_size": null,
-  "max_position_embeddings": 2048,
-  "model_type": "gated_deltanet",
-  "norm_eps": 1e-06,
-  "num_heads": 6,
-  "num_hidden_layers": 21,
-  "num_v_heads": null,
-  "tie_word_embeddings": false,
-  "transformers_version": "4.57.6",
-  "use_cache": true,
-  "use_gate": true,
-  "use_l2warp": false,
-  "use_short_conv": true,
-  "vocab_size": 32000
-}

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/blt_transformer_1000hash.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/blt_transformer_1_5B.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/blt_transformer_380M.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/delta_net_1B.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/delta_net_340M.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/gated_deltanet_1B.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/gated_deltanet_340M.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/gated_deltanet_h_340M.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/gla_340M.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/gla_7B.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/gsa_340M.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/mergenet_340M.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/mergenet_64M.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/qwen3_next_1B.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/qwen3_next_350M.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/transformer_1B.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/transformer_340M.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/configs/transformer_7B.json

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/flame/__init__.py

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/flame/c4_test.py

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/flame/components/checkpoint.py

@@ -1,59 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-from dataclasses import dataclass, field
-from datetime import timedelta
-from io import BytesIO
-from typing import Any, Dict, List
-
-import torch
-from torch.distributed.checkpoint.stateful import Stateful
-
-
-@dataclass
-class TrainState(Stateful):
-    step: int = 0
-    skipped_step: int = 0
-    token: int = 0
-    elapsed: timedelta = timedelta(0)
-    global_avg_losses: List[float] = field(default_factory=list)
-    global_max_losses: List[float] = field(default_factory=list)
-    log_steps: List[int] = field(default_factory=list)
-
-    def state_dict(self) -> Dict[str, Any]:
-        # Only checkpoint global_avg_losses and global_max_losses per log frequency
-        # to avoid sync overhead in every iteration.
-        global_avg_losses_bytes = BytesIO()
-        torch.save(self.global_avg_losses, global_avg_losses_bytes)
-        global_max_losses_bytes = BytesIO()
-        torch.save(self.global_max_losses, global_max_losses_bytes)
-        log_steps_bytes = BytesIO()
-        torch.save(self.log_steps, log_steps_bytes)
-        return {
-            "step": torch.tensor(self.step, dtype=torch.int32),
-            "skipped_step": torch.tensor(self.skipped_step, dtype=torch.int32),
-            "token": torch.tensor(self.token, dtype=torch.int64),
-            "elapsed": self.elapsed,
-            "global_avg_losses": global_avg_losses_bytes,
-            "global_max_losses": global_max_losses_bytes,
-            "log_steps": log_steps_bytes,
-        }
-
-    def load_state_dict(self, state_dict) -> None:
-        self.step = state_dict["step"].item()
-        self.skipped_step = state_dict.get("skipped_step", 0).item()
-        self.token = state_dict["token"].item()
-        self.elapsed = state_dict["elapsed"]
-        state_dict["global_avg_losses"].seek(0)
-        self.global_avg_losses = torch.load(
-            state_dict["global_avg_losses"], weights_only=False
-        )
-        state_dict["global_max_losses"].seek(0)
-        self.global_max_losses = torch.load(
-            state_dict["global_max_losses"], weights_only=False
-        )
-        state_dict["log_steps"].seek(0)
-        self.log_steps = torch.load(state_dict["log_steps"], weights_only=False)

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/flame/config_manager.py

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

gated_deltanet_1b_v3/gated_deltanet_1b_mars_shampoo_lr2e_2_b1_0_95_b2_0_99_eps_1e_15_scale2_0_rank512_20260506_140829/exp_data/flame/data.py

@@ -1,762 +0,0 @@
-# -*- 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, interleave_datasets, load_dataset
-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 import utils
-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.uint16).max:
-            self.dtype = torch.uint16
-        elif tokenizer.vocab_size < torch.iinfo(torch.uint32).max:
-            self.dtype = torch.uint32
-        else:
-            self.dtype = torch.uint64
-        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:
-            # ------ jinxin ------ #
-            # Sample: 
-            #        text: str
-            #        id: <xxx>
-            #        ...
-            #        language: en/cn
-            #        language score: float
-            #        token count: int
-            #        score / int_score : float/int
-            # print(sample)
-            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:   # 🌟Tips: self.tokenizer() 直接调用的 __call__
-                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_dataset(
-    dataset: str,
-    dataset_name: str = None,
-    dataset_split: str = 'train',
-    data_dir: str = None,
-    data_files: str = None,
-    data_probs: List[float] = None,
-    streaming: bool = False,
-    dp_degree: Optional[int] = None,
-    num_workers: int = 32,
-    seed: Optional[int] = None,
-) -> IterableDataset:
-    color = utils.Color
-    min_num_shards = dp_degree * num_workers if dp_degree else None
-    if len(dataset.split(',')) == 1:
-        dataset = load_dataset(
-            path=dataset,
-            name=dataset_name,
-            split=dataset_split,
-            data_dir=data_dir,
-            data_files=data_files,
-            streaming=streaming,
-            num_proc=num_workers if not streaming else None,
-        )
-        logger.info(f"Shuffling the dataset with seed {seed}")
-        if not streaming:
-            # the states of map-style dataset is recoverable after shuffling
-            if seed is not None:
-                dataset = dataset.shuffle(seed=seed)
-            if min_num_shards is not None:
-                dataset = dataset.to_iterable_dataset(num_shards=min_num_shards)
-        else:
-            if min_num_shards is not None and dataset.num_shards < min_num_shards:
-                logger.warning(
-                    f"{color.red}"
-                    f"Dataset {dataset} has insufficient shards ({dataset.num_shards}). "
-                    f"Need {min_num_shards} shards minimum for {dp_degree} data parallel workers × "
-                    f"{num_workers} dataloader workers. "
-                    f"Disabling the streaming mode and resharding dataset to {min_num_shards} shards."
-                    f"{color.reset}"
-                )
-                dataset = load_dataset(
-                    path=dataset,
-                    name=dataset_name,
-                    split=dataset_split,
-                    data_dir=data_dir,
-                    data_files=data_files,
-                    streaming=False,
-                    num_proc=num_workers,
-                )
-                if seed is not None:
-                    dataset = dataset.shuffle(seed=seed)
-                dataset = dataset.to_iterable_dataset(num_shards=min_num_shards)
-            else:
-                if seed is not None:
-                    dataset = shuffle(dataset, seed=seed)
-    else:
-        datasets = dataset.split(",")
-        if dataset_name is not None:
-            dataset_names = [
-                name or None for name in 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 dataset_split is not None:
-            dataset_splits = [split or "train"for split in 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 data_dir is not None:
-            data_dirs = [
-                data_dir or None for data_dir in 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 data_files is not None:
-            data_files = 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 data_probs is not None:
-            data_probs = [float(p) for p in 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],
-                split=dataset_splits[i],
-                data_dir=data_dirs[i],
-                data_files=data_files[i],
-                streaming=streaming,
-                num_proc=(
-                    num_workers
-                    if not 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 {seed}")
-            if not streaming:
-                # the states of map-style dataset is recoverable after shuffling
-                if seed is not None:
-                    subset = subset.shuffle(seed=seed)
-                if min_num_shards is not None:
-                    subset = subset.to_iterable_dataset(num_shards=min_num_shards)
-            else:
-                if min_num_shards is not None and 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 desired data parallel workers × "
-                        f"{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],
-                        split=dataset_splits[i],
-                        data_dir=data_dirs[i],
-                        data_files=data_files[i],
-                        streaming=False,
-                        num_proc=num_workers,
-                    )
-                    if seed is not None:
-                        subset = subset.shuffle(seed=seed)
-                    subset = subset.to_iterable_dataset(num_shards=min_num_shards)
-                else:
-                    # we set relatively small buffer size here as interleaving could provide some randomness
-                    if seed is not None:
-                        subset = shuffle(subset, seed=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=seed,
-        )
-    logger.info(f"{dataset}")
-    return dataset
-
-
-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,
-    )