Upload 50 files

flame/__init__.py

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

flame/components/checkpoint.py

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

flame/config_manager.py

@@ -0,0 +1,960 @@
+# 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",
+        )
+        # metrics configs
+        self.parser.add_argument(
+            "--metrics.log_freq",
+            type=int,
+            default=10,
+            help="How often to log metrics to TensorBoard, in iterations",
+        )
+        self.parser.add_argument(
+            "--metrics.enable_tensorboard",
+            action="store_true",
+            help="Whether to log metrics to TensorBoard",
+        )
+        self.parser.add_argument(
+            "--metrics.disable_color_printing",
+            action="store_true",
+            help="Whether to disable color printing in logs",
+        )
+        self.parser.add_argument(
+            "--metrics.save_tb_folder",
+            type=str,
+            default="tb",
+            help="Folder to dump TensorBoard states",
+        )
+        self.parser.add_argument(
+            "--metrics.save_for_all_ranks",
+            action="store_true",
+            default=False,
+            help="""
+                Whether to save TensorBoard/Wandb metrics only for rank 0 or for all ranks.
+                When this option is False and pipeline_parallel_degree is > 1, the metrics
+                component uses the 0th rank of the last stage pipeline group, which is the
+                only stage that computes loss metrics.
+            """,
+        )
+        self.parser.add_argument(
+            "--metrics.enable_wandb",
+            action="store_true",
+            help="Whether to log metrics to Weights & Biases",
+        )
+
+        self.parser.add_argument(
+            "--experimental.enable_async_tensor_parallel",
+            action="store_true",
+            help="Whether to apply async tensor parallel (currently only effective when compile is enabled)",
+        )
+        self.parser.add_argument(
+            "--experimental.pipeline_parallel_degree",
+            type=int,
+            default=1,
+            help="""
+                Pipeline Parallelism degree, or number of ranks. 1 means disabled.
+                If using looped schedules, this still specifies the number of physical ranks, not the number
+                of stages.  Stages per rank are inferred from split points degree, and schedule.""",
+        )
+        self.parser.add_argument(
+            "--experimental.pipeline_parallel_split_points",
+            type=string_list,
+            nargs="+",
+            default=[],
+            help="""
+                Specify comma-separated names of modules to use as the beginning of a split point.
+
+                e.g. "layers.0,layers.2" will cause the model to be split into 3 stages,
+                the first containing all the layers up to layers.0,
+                the second containing layers.0 and up to layers.2,
+                the third containing layers.2 and all the remaining layers.
+
+                Note: fully-automated splitting may be enabled in the future,
+                but currently the split points must be specified manually.""",
+        )
+        self.parser.add_argument(
+            "--experimental.pipeline_parallel_schedule",
+            type=str,
+            default="1F1B",
+            help="""
+                Specify the Pipeline Parallel schedule to use. The supported schedules are:
+                https://github.com/pytorch/pytorch/blob/de4c2a3b4e89d96334dc678d1c3f2ae51a6630a0/torch/distributed/pipelining/schedules.py#L2161.
+                The schedule must be compatible with the split points and stages_per_rank.
+
+                Looped schedules (e.g. Interleaved1F1B) require specifying pipeline_parallel_degree = number of ranks,
+                and split_points = number of stages - 1
+                """,
+        )
+        self.parser.add_argument(
+            "--experimental.pipeline_parallel_schedule_csv",
+            type=str,
+            default="",
+            help="""
+                Specify the path to the pipeline parallel schedule csv file to use.
+                The pipeline_parallel_schedule argument must be either
+                PipelineScheduleSingle, PipelineScheduleMulti, or _PipelineScheduleRuntime.
+            """,
+        )
+
+        self.parser.add_argument(
+            "--experimental.pipeline_parallel_microbatches",
+            type=int,
+            default=None,
+            help="""
+                How many microbatches to split the global training batch into when using pipeline parallelism.
+
+                The global training batch size must be evenly divisible by the number of microbatches.
+
+                The default value will be the number of pipeline stages, if unspecified.
+            """,
+        )
+        self.parser.add_argument(
+            "--experimental.enable_compiled_autograd",
+            action="store_true",
+            help="Enable CompiledAutograd to compile the backward.",
+        )
+        self.parser.add_argument(
+            "--experimental.context_parallel_degree",
+            type=int,
+            default=1,
+            help="Context parallelism degree. 1 means disabled.",
+        )
+        self.parser.add_argument(
+            "--experimental.context_parallel_rotate_method",
+            type=str,
+            default="allgather",
+            help="""
+                The collective to use in context parallel SDPA for kv shards exchange.
+
+                'allgather' means to all-gather all kv shards on ranks after the first sub-SDPA computation,
+
+                'alltoall' means to all-to-all shuffle the kv shards.
+
+                The default value is 'allgather'.
+            """,
+        )
+        # I'm not particularly fond of this. Users can choose to write their own wrapper
+        # module and import TorchTitan training loop and execute it, which look cleaner.
+        # One reason to provide this option is to allow users to use the existing run script.
+        # While the script is pretty trivial now, we may add more logic when integrating
+        # with TorchFT.
+        # This option is subject to change and may be deleted in the future.
+        self.parser.add_argument(
+            "--experimental.custom_model_path",
+            type=str,
+            default="",
+            help="""
+                The --custom_model_path option allows to specify a custom path to a model module
+                that is not natively implemented within TorchTitan.
+                Acceptable values are the file system path to the module (e.g., my_models/model_x)
+                dotted import module  (e.g., some_package.model_x).
+            """,
+        )
+        # checkpointing configs
+        self.parser.add_argument(
+            "--checkpoint.enable_checkpoint",
+            action="store_true",
+            help="Whether to enable checkpoint",
+        )
+        self.parser.add_argument(
+            "--checkpoint.folder",
+            type=str,
+            default="checkpoint",
+            help="""
+                The folder to store the checkpoints.
+                When enable_checkpoint is set to true, checkpoints will be in {--job.dump_folder}/{--checkpoint.folder}.
+            """,
+        )
+        self.parser.add_argument(
+            "--checkpoint.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

flame/data.py

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

flame/models/activation_offloading.py

@@ -0,0 +1,447 @@
+# Adapted from https://github.com/pytorch/torchtune/blob/main/torchtune/training/_activation_offloading.py
+# 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 contextlib
+from typing import Union
+from warnings import warn
+
+import psutil
+import torch
+from torch import nn
+from torch.autograd.graph import saved_tensors_hooks
+
+from torchtitan.tools.logging import logger
+
+try:
+    import torchao
+    from torchao.dtypes.nf4tensor import NF4Tensor
+except ImportError:
+    torchao = None
+    NF4Tensor = None
+    logger.warning("torchao not found. ")
+
+# from torchtune.modules import TiedLinear
+
+
+class OffloadActivations(saved_tensors_hooks):
+    """Context manager under which activation tensors created in the forward pass will be offloaded.
+
+    Enable the memory efficiency technique of activation offloading, where activations bigger than
+    min_offload_size bytes will be offloaded to CPU in the forward and brought back in the backward.
+    This is in contrast to maintaining the activation on GPU VRAM throughout the program.
+
+    This manager contains the option of using one additional CUDA stream to handle the communication
+    between CUDA and CPU, which is intended to overlap with the default computation stream to improve
+    runtime. We designed synchronization with a few heuristics for optimizing the tradeoff between
+    runtime vs memory usage.
+
+    Args:
+        use_pin_memory (bool): Whether or not the offloaded Tensor will be placed in pinned
+            memory on the CPU. Pinned memory allows the Tensor to be moved back onto GPU more quickly
+            but is a limited resource. Default: True.
+
+        use_streams (bool): Whether or not to use streams for performance optimization where
+            the communications get overlapped with the computation. Requires a torch build
+            after torch-2.5.0.]. Default: True.
+
+        max_fwd_stash_size (int): The maximum size of the forward stash, or the maximum number of
+            consecutive activations to keep alive during the forward pass. This number must be at
+            least 1. Keeping alive more activations will potentially allow more overlap between the
+            communication and compute streams at the cost of increasing memory usage. Keeping alive
+            fewer activations will conserve memory, but may cause poor overlap between the streams,
+            increasing runtime. Default: 5.
+
+        min_offload_size (int): The minimum number of bytes a Tensor must be in order to qualify
+            for offloading. If the tensor is too small, we do not want to waste bandwidth and resources
+            moving it to CPU and back. Default: 1024 bytes.
+
+    Raises:
+        ValueError: if max_fwd_stash_size is not at least 1.
+
+    Example:
+        >>> with OffloadActivations():
+        >>>     logits = model(inputs)
+        >>> loss = ...
+        >>> loss.backward()
+    """
+
+    def __init__(
+        self,
+        use_pin_memory: bool = True,
+        use_streams: bool = True,
+        max_fwd_stash_size: int = 5,
+        min_offload_size: int = 1024,
+    ) -> None:
+
+        self.use_streams: bool = use_streams
+
+        self.min_tensor_size_bytes = (
+            min_offload_size  # we don't want to bother with small tensors
+        )
+        self.tracker = (
+            {}
+        )  # tensor_id => (new_tensor, if_modified)  ---> track what saved/offloaded tensors are where
+        self.tensor_id: int = 0
+        self.is_first_forward_call = True
+        self.is_first_backward_call = True
+        self.is_first_forward_pass = True
+
+        # managing cpu memory
+        self.use_pin_memory: bool = use_pin_memory
+        self.virtual_memory_safe_pct = (
+            60  # we should not exceed this percentage of memory
+        )
+
+        self.s0 = torch.cuda.default_stream()  # comp stream
+
+        # for streaming
+        if self.use_streams:
+            self.s1 = torch.cuda.Stream()  # comms stream
+            self.fwd_stash = {}  # tensor_id => (activation, ev1)
+            if max_fwd_stash_size < 1:
+                raise ValueError(
+                    f"max_fwd_stash_size should be at least 1 but is {max_fwd_stash_size}"
+                )
+            self.max_fwd_stash_size = max_fwd_stash_size
+            self.bwd_tensor_stash = {}  # tensor_id => activation
+            self.bwd_ev_stash = {}  # tensor_id => ev0
+            self.curr_graph_id = None
+            self.curr_autograd_node = None
+
+        # -------- platform util functions -------- #
+        def verify_sufficient_virtual_memory():
+            curr_pct = get_cpu_ram_pct()
+            if curr_pct > self.virtual_memory_safe_pct:
+                warn(
+                    f"***** WARNING: {curr_pct=}% > {self.virtual_memory_safe_pct=}% of virtual memory used"
+                )
+
+        def get_cpu_ram_pct() -> float:
+            # get the percentage of memory used by the system
+            return psutil.virtual_memory().percent
+
+        def get_tensor_id() -> int:
+            # create a unique id for each tensor we are managing
+            self.tensor_id += 1
+            return self.tensor_id
+
+        def get_num_bytes_tensor(x: torch.Tensor) -> int:
+            # get the number of bytes in a tensor, for memory management purposes
+            return (
+                x.element_size() * x.nelement()
+            )  # x.element_size() * x._base_storage().nbytes()
+
+        # -------- core pack / unpack work -------- #
+        def pack_tensor(activation: torch.Tensor) -> int:
+            # activations are passed in during forward pass - from here we take over and return a unique id
+            if self.is_first_forward_call:
+                assert (
+                    len(self.tracker) == 0
+                ), "backward pass should have cleared tracker of all tensors"
+
+                # set training phase trackers
+                self.is_first_forward_call = False
+                self.is_first_backward_call = True
+
+            # query for basic tensor info
+            num_bytes = get_num_bytes_tensor(activation)
+            tensor_id = get_tensor_id()
+
+            # only offload hefty bois if they're activations on CUDA (our heuristic
+            # for that is to check if they're not params or buffers)!
+            if (
+                activation.is_cuda
+                and num_bytes >= self.min_tensor_size_bytes
+                and (
+                    not isinstance(activation, torch.nn.Parameter)
+                    and not isinstance(activation, torch.nn.Buffer)
+                )
+            ):
+                if self.use_streams:
+                    # First, sync back and dereference previously offloaded tensors
+                    # as the offloading should be done sufficiently long ago.
+                    for id in [k for k in self.fwd_stash.keys()]:
+                        if id <= tensor_id - self.max_fwd_stash_size:
+                            _, ev = self.fwd_stash[id]
+                            self.s0.wait_event(ev)
+                            del self.fwd_stash[id]
+                        else:
+                            break
+
+                    # Sync in, offload, and add an event to sync back later
+                    self.s1.wait_stream(self.s0)
+
+                stream = self.s1 if self.use_streams else self.s0
+                with torch.cuda.stream(stream):
+                    try:
+                        cpu_tensor = torch.empty_like(
+                            activation, pin_memory=self.use_pin_memory, device="cpu"
+                        )
+                    except NotImplementedError as e:
+                        if (
+                            isinstance(activation, NF4Tensor)
+                            and torchao.__version__ < "0.6.0.dev20240917"
+                        ):
+                            raise RuntimeError(
+                                "Offloading NF4Tensors requires torchao-0.6.0.dev20240917 or later"
+                            ) from e
+                        raise e
+                    cpu_tensor.copy_(activation, non_blocking=True)
+                    self.tracker[tensor_id] = (
+                        cpu_tensor,
+                        True,
+                    )  # True = (in future) modified
+
+                if self.use_streams:
+                    event = self.s1.record_event()
+
+                    # Stash to keep activation alive til s1 is done
+                    self.fwd_stash[tensor_id] = (activation, event)
+            else:
+                self.tracker[tensor_id] = (
+                    activation,
+                    False,
+                )  # False = not modified, tensor is as is
+
+            return tensor_id
+
+        def unpack_tensor_single_stream(unpack_tensor_id: int) -> torch.Tensor:
+            # backward pass - we are called with the tensor_id, which
+            # we will use to retrieve the saved/offloaded tensor
+            if self.is_first_backward_call:
+                if self.is_first_forward_pass:
+                    self.is_first_forward_pass = False
+                    if self.use_pin_memory:
+                        verify_sufficient_virtual_memory()
+
+                self.is_first_backward_call = False
+                self.is_first_forward_call = True
+
+            assert (
+                unpack_tensor_id in self.tracker
+            ), f"untracked tensor with id {unpack_tensor_id}"
+
+            maybe_gpu_tensor, modified = self.tracker[unpack_tensor_id]
+            if modified:
+                gpu_tensor = maybe_gpu_tensor.to("cuda", non_blocking=True)
+                maybe_gpu_tensor = gpu_tensor
+
+            # clear tensor from tracking
+            del self.tracker[unpack_tensor_id]
+            return maybe_gpu_tensor
+
+        def unpack_tensor_with_streams(unpack_tensor_id: int) -> torch.Tensor:
+            # backward pass - we are called with the tensor_id, which
+            # we will use to retrieve the saved/offloaded tensor
+            if self.is_first_backward_call:
+                self.curr_graph_id = torch._C._current_graph_task_id()
+
+                def wait_and_del_remaining_references() -> None:
+                    for id in [k for k in self.bwd_tensor_stash.keys()]:
+                        event = self.bwd_ev_stash[id]
+                        self.s1.wait_event(event)
+                        del self.bwd_tensor_stash[id]
+
+                # Register a callback to the end of autograd to clean everything up
+                torch.autograd.variable.Variable._execution_engine.queue_callback(
+                    wait_and_del_remaining_references
+                )
+
+                if self.is_first_forward_pass:
+                    self.is_first_forward_pass = False
+                    if self.use_pin_memory:
+                        verify_sufficient_virtual_memory()
+
+                self.is_first_backward_call = False
+                self.is_first_forward_call = True
+
+            assert (
+                unpack_tensor_id in self.tracker
+            ), f"untracked tensor with id {unpack_tensor_id}"
+
+            maybe_gpu_tensor, modified = self.tracker[unpack_tensor_id]
+            if modified:
+                # Get data on the current autograd node
+                graph_id = torch._C._current_graph_task_id()
+                node = torch._C._current_autograd_node()
+                prev_node_ids = []
+
+                # If we're on a new node, mark prev node's tensors to be freed later
+                if graph_id == self.curr_graph_id and self.curr_autograd_node != node:
+                    self.curr_autograd_node = node
+                    prev_node_ids = [id for id in self.bwd_tensor_stash.keys()]
+
+                brought_back_from_cpu = True
+                if unpack_tensor_id in self.fwd_stash:
+                    maybe_gpu_tensor = self.fwd_stash[unpack_tensor_id][0]
+                    brought_back_from_cpu = False
+                else:
+                    # Kick off the process to bring tensors back
+                    with torch.cuda.stream(self.s1):
+                        gpu_tensor = maybe_gpu_tensor.to("cuda", non_blocking=True)
+                        maybe_gpu_tensor = gpu_tensor
+
+                    # Tell comp stream to wait for the info to be loaded before executing
+                    self.s0.wait_stream(self.s1)
+
+                    # Stash the tensor to keep memory alive until compute stream is complete
+                    self.bwd_tensor_stash[unpack_tensor_id] = maybe_gpu_tensor
+
+                    # Note: [Track views of the unpacked]
+                    # Why do we get the use count of the unpacked tensor here? We want an
+                    # initial count to compare to later, during the post-hook of the
+                    # backward node, when we need to decide whether we're allowed to free
+                    # the tensor yet. In what obscure cases must we delay freeing the
+                    # tensor (and thus call record_stream)?
+                    # 1. Any of the outputs of the backward node is a view of the unpacked
+                    #    tensor.
+                    # 2. In the case that this unpacked tensor will be used in a
+                    #    checkpointed region, if one of the recomputed saved tensors ends
+                    #    up as a view of the unpacked tensor.
+                    # 3. The user abuses the system somehow and manually relies on the
+                    #    unpacked tensor to exist after the backward node has executed.
+                    storage_refcount = torch._C._storage_Use_Count(
+                        maybe_gpu_tensor.untyped_storage()._cdata
+                    )
+
+                def hook(outputs, inputs):
+                    # create events for the current node inputs/outputs if they were streamed in
+                    if brought_back_from_cpu:
+                        # See Note: [Track views of the unpacked]
+                        # IF any of the outputs is a view of the tensor, OR if a view of
+                        # the tensor has been saved as a part of checkpoint's recompute
+                        # process, OR the user has abusedly incurred a reference on the
+                        # unpacked tensor, THEN the tensor might be used later and we
+                        # cannot presume to delete it after only the current node is
+                        # done! So we use our frenemy, record_stream, to ensure the
+                        # Tensor stays unmessed with until it's done getting used in the
+                        # compute stream (s0 here). Note that the con here is we introduce
+                        # non-deterministic (thus higher) memory usage, but this case
+                        # should not happen often.
+                        unpacked_tensor = self.bwd_tensor_stash[unpack_tensor_id]
+                        if (
+                            torch._C._storage_Use_Count(
+                                unpacked_tensor.untyped_storage()._cdata
+                            )
+                            > storage_refcount
+                        ):
+                            unpacked_tensor.record_stream(self.s0)
+                            del self.bwd_tensor_stash[unpack_tensor_id]
+                        else:
+                            event = self.s0.record_event()
+                            self.bwd_ev_stash[unpack_tensor_id] = event
+
+                    # if there are still things in the fwd_stash, get rid of them as we're in bwd now
+                    for id in [k for k in self.fwd_stash.keys()]:
+                        _, ev = self.fwd_stash[id]
+                        self.s0.wait_event(ev)
+                        del self.fwd_stash[id]
+
+                    # wait on prev node's events and del those
+                    for id in prev_node_ids:
+                        event = self.bwd_ev_stash[id]
+                        self.s1.wait_event(event)
+                        del self.bwd_tensor_stash[id]
+
+                    return outputs
+
+                node.register_hook(hook)
+
+            # clear tensor from tracking
+            del self.tracker[unpack_tensor_id]
+            return maybe_gpu_tensor
+
+        unpack_tensor = (
+            unpack_tensor_with_streams
+            if self.use_streams
+            else unpack_tensor_single_stream
+        )
+        super().__init__(pack_tensor, unpack_tensor)
+
+
+class NoOpManager(saved_tensors_hooks):
+    """
+    A saved_tensors_hook manager used to disable any other saved_tensors_hook manager
+    applied before. This relies on the behavior that only the most recently registered
+    saved_tensors_hook will run.
+
+    One example usage is to opt a local region of code out of activations offloading,
+    which is usually applied globally to best track state.
+    """
+
+    def __init__(self) -> None:
+        def noop(tensor):
+            return tensor
+
+        super().__init__(noop, noop)
+
+
+def get_act_offloading_ctx_manager(
+    model: nn.Module, enable_activation_offloading: bool
+) -> Union[OffloadActivations, contextlib.nullcontext]:
+    """Returns the activation offloading context manager for the model, which will be
+    a null context if enable_activation_offloading is False.
+
+    If activation offloading is enabled, we return the OffloadActivations context manager.
+    If activation offloading is disabled, we return a NoOpManager context manager.
+
+    Args:
+        model (nn.Module): the model to wrap with the activation offloading context manager.
+        enable_activation_offloading (bool): whether or not to enable activation offloading
+            for the model.
+
+    Returns:
+        contextlib.ContextDecorator: the activation offloading context manager for the model.
+
+    Raises:
+        NotImplementedError: If the model is a multimodal model and activation offloading is enabled.
+    """
+    if enable_activation_offloading:
+        activations_handling_ctx = OffloadActivations()
+
+        # Below is our hack to disable offloading the last output Linear in every
+        # step, as the cost for offloading the activation and then soon after bringing
+        # it back is expensive. Moreover, due to heuristics in our streaming API,
+        # we actually use more memory if we offload it as it interferes with chunkedCE.
+        output_head_detected = False
+        noop_ctx = NoOpManager()
+
+        if hasattr(model, "output"):
+            if isinstance(model.output, nn.Module):
+                model.output.register_forward_pre_hook(
+                    lambda *args: noop_ctx.__enter__()
+                )
+                model.output.register_forward_hook(
+                    lambda *args: noop_ctx.__exit__(), always_call=True
+                )
+                print("registering hooks for model.output ============  ")
+                output_head_detected = True
+            # ================================
+            # ! TODO[flame] check if we need to detal with TiedLinear
+            # The following code appears in `torchtune`
+            # elif isinstance(model.output, TiedLinear):
+            #     model.output.linear.register_forward_pre_hook(
+            #         lambda *args: noop_ctx.__enter__()
+            #     )
+            #     model.output.linear.register_forward_hook(
+            #         lambda *args: noop_ctx.__exit__(), always_call=True
+            #     )
+            #     output_head_detected = True
+
+        if not output_head_detected:
+            logger.warning(
+                "During activation offloading, no output head was detected. "
+                "If your model has an output head, it will be offloaded. "
+                "This usually greatly slows training, given the large vocabulary size. "
+                "To change this behavior, set your output head as model.output and make it "
+                "an nn.Module."
+            )
+
+    else:
+        activations_handling_ctx = contextlib.nullcontext()
+
+    return activations_handling_ctx

flame/models/fla.toml

@@ -0,0 +1,67 @@
+[model]
+config = "fla-hub/hamilton-350M-15B"
+tokenizer_path = "mistralai/Mistral-7B-v0.1"
+
+[job]
+dump_folder = "exp"
+print_args = true
+
+[training]
+batch_size = 2
+seq_len = 2048
+context_len = 2048
+gradient_accumulation_steps = 1
+steps = 20480
+max_norm = 1.0
+skip_nan_inf = true
+data_parallel_replicate_degree = 1
+data_parallel_shard_degree = -1
+tensor_parallel_degree = 1
+compile = false
+dataset = "SlimPajama-627B"
+dataset_name = "default"
+num_workers = 32
+pin_memory = false
+persistent_workers = false
+prefetch_factor = 2
+seed = 42
+varlen = false
+
+[optimizer]
+name = "AdamW"
+eps = 1e-15
+lr = 3e-4
+
+[lr_scheduler]
+warmup_steps = 1024
+decay_type = "cosine"
+lr_min = 0.1
+
+[checkpoint]
+enable_checkpoint = true
+folder = "checkpoint"
+interval_type = "steps"
+interval = 2048
+model_weights_only = false
+export_dtype = "float32"
+async_mode = "disabled"    # ["disabled", "async", "async_with_pinned_mem"]
+
+[profiling]
+enable_profiling = false
+save_traces_folder = "profile_trace"
+profile_freq = 256
+
+[metrics]
+log_freq = 32
+enable_wandb = true
+
+[experimental]
+context_parallel_degree = 1
+pipeline_parallel_degree = 1
+
+[float8]
+enable_fsdp_float8_all_gather = false
+precompute_float8_dynamic_scale_for_fsdp = false
+
+[activation_checkpoint]
+mode = "none"

flame/models/parallelize_fla.py

@@ -0,0 +1,550 @@
+# 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.
+
+# This file applies the PT-D parallelisms (except pipeline parallelism) and various
+# training techniques (e.g. activation checkpointing and compile) to the Llama model.
+
+from collections import defaultdict
+
+import torch
+import torch.nn as nn
+from torch.distributed import DeviceMesh
+from torch.distributed._composable.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard
+from torch.distributed._composable.replicate import replicate
+from torch.distributed._tensor import Replicate, Shard
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper as ptd_checkpoint_wrapper
+from torch.distributed.tensor.parallel import (
+    ColwiseParallel,
+    PrepareModuleInput,
+    PrepareModuleOutput,
+    RowwiseParallel,
+    SequenceParallel,
+    parallelize_module
+)
+
+from fla.modules.fused_linear_cross_entropy import LinearLossParallel
+from fla.modules.mlp import SwiGLULinearParallel
+from fla.modules.parallel import PrepareModuleWeight
+from torchtitan.config_manager import TORCH_DTYPE_MAP, JobConfig
+from torchtitan.distributed.parallel_dims import ParallelDims
+from torchtitan.tools.logging import logger
+
+
+def parallelize_fla(
+    model: nn.Module,
+    world_mesh: DeviceMesh,
+    parallel_dims: ParallelDims,
+    job_config: JobConfig,
+):
+    """
+    Apply tensor parallelism, activation checkpointing, torch.compile, and data
+    parallelism to the model.
+
+    NOTE: The passed-in model preferably should be on meta device. Otherwise,
+    the model must fit on GPU or CPU memory.
+    """
+
+    if parallel_dims.tp_enabled:
+        if (
+            job_config.experimental.enable_async_tensor_parallel
+            and not job_config.training.compile
+        ):
+            raise RuntimeError("Async TP requires --training.compile")
+        enable_float8_linear = "float8" in job_config.model.converters
+        apply_tp(
+            model,
+            world_mesh["tp"],
+            loss_parallel=parallel_dims.loss_parallel_enabled,
+            enable_float8=enable_float8_linear,
+            enable_async_tp=job_config.experimental.enable_async_tensor_parallel,
+        )
+
+    if job_config.activation_checkpoint.mode != "none":
+        apply_ac(model, job_config.activation_checkpoint)
+
+    # turn on per-block compile after AC wrapping and before FSDP
+    if job_config.training.compile:
+        apply_compile(model)
+
+    if (
+        parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled
+    ):  # apply FSDP or HSDP, potentially with Context Parallel
+        if parallel_dims.dp_replicate_enabled:
+            dp_mesh_dim_names = ("dp_replicate", "dp_shard_cp")
+        else:
+            dp_mesh_dim_names = ("dp_shard_cp",)
+
+        apply_fsdp(
+            model,
+            world_mesh[tuple(dp_mesh_dim_names)],
+            param_dtype=TORCH_DTYPE_MAP[job_config.training.mixed_precision_param],
+            reduce_dtype=TORCH_DTYPE_MAP[job_config.training.mixed_precision_reduce],
+            pp_enabled=parallel_dims.pp_enabled,
+            cpu_offload=job_config.training.enable_cpu_offload,
+            reshard_after_forward_policy=job_config.training.fsdp_reshard_after_forward,
+        )
+
+        if parallel_dims.dp_replicate_enabled:
+            logger.info("Applied HSDP to the model")
+        else:
+            logger.info("Applied FSDP to the model")
+
+        if parallel_dims.cp_enabled:
+            logger.info("Applied Context Parallel to the model")
+
+        if job_config.training.enable_cpu_offload:
+            logger.info("Applied CPU Offloading to the model")
+    elif parallel_dims.dp_replicate_enabled:
+        if world_mesh.ndim > 1:
+            raise RuntimeError("DDP has not supported > 1D parallelism")
+        apply_ddp(
+            model,
+            world_mesh,
+            enable_compile=job_config.training.compile,
+            enable_compiled_autograd=job_config.experimental.enable_compiled_autograd,
+        )
+
+
+class TPPlan:
+    def __init__(
+        self,
+        model=None,
+        loss_parallel=False,
+        enable_float8=False,
+    ):
+        self.model = model
+        self.loss_parallel = loss_parallel
+        self.enable_float8 = enable_float8
+        self.base_model_prefix = getattr(model, "base_model_prefix", "model")
+
+        # TODO(vkuzo): once float8 configuration supports delayed scaling,
+        # add a check here to enforce supported float8 all-gather configurations
+        # TODO(vkuzo): add the items below to __init__.py of torchao.float8 and import from there
+        try:
+            from torchao.float8.float8_tensor_parallel import (
+                Float8ColwiseParallel,
+                Float8RowwiseParallel,
+                PrepareFloat8ModuleInput
+            )
+        except ImportError:
+            Float8ColwiseParallel = None
+            Float8RowwiseParallel = None
+            PrepareFloat8ModuleInput = None
+        if self.enable_float8 and Float8ColwiseParallel is not None:
+            self.rowwise_parallel = Float8RowwiseParallel
+            self.colwise_parallel = Float8ColwiseParallel
+            self.prepare_module_input = PrepareFloat8ModuleInput
+            self.prepare_module_output = PrepareModuleOutput
+        else:
+            self.rowwise_parallel = RowwiseParallel
+            self.colwise_parallel = ColwiseParallel
+            self.prepare_module_input = PrepareModuleInput
+            self.prepare_module_output = PrepareModuleOutput
+
+    @property
+    def model_plan(self):
+        plans = {
+            f"{self.base_model_prefix}.embeddings": RowwiseParallel(
+                input_layouts=Replicate(),
+                output_layouts=Shard(1),
+            ),
+            f"{self.base_model_prefix}.norm": SequenceParallel(),
+        }
+        if self.loss_parallel:
+            plans.update(
+                {
+                    "lm_head": ColwiseParallel(
+                        input_layouts=Shard(1),
+                        output_layouts=Shard(-1) if self.loss_parallel else Replicate(),
+                        use_local_output=not self.loss_parallel,
+                    ),
+                }
+            )
+        else:
+            plans.update(
+                {
+                    "lm_head": PrepareModuleWeight(layouts=Replicate()),
+                    "criterion": LinearLossParallel(),
+                }
+            )
+        return plans
+
+    @property
+    def layer_plan(self):
+        return {
+            "attn_norm": SequenceParallel(),
+            **self.attn_plan,
+            "mlp_norm": SequenceParallel(),
+            **self.mlp_plan,
+        }
+
+    @property
+    def attn_plan(self):
+        raise NotImplementedError(
+            f"TP plans for token mixing layers of {self.model.config.model_type} not implemented"
+        )
+
+    @property
+    def mlp_plan(self):
+        return {
+            "mlp": self.prepare_module_input(
+                input_layouts=(Shard(1),),
+                desired_input_layouts=(Replicate(),),
+            ),
+            "mlp.gate_proj": self.colwise_parallel(),
+            "mlp.up_proj": self.colwise_parallel(),
+            "mlp.down_proj": self.rowwise_parallel(output_layouts=Shard(1)),
+            "mlp.swiglu_linear": SwiGLULinearParallel(output_layouts=Shard(1)),
+        }
+
+
+class TransformerTPPlan(TPPlan):
+
+    @property
+    def attn_plan(self):
+        return {
+            "attn": self.prepare_module_input(
+                input_kwarg_layouts={"hidden_states": Shard(1)},
+                desired_input_kwarg_layouts={"hidden_states": Replicate()},
+            ),
+            "attn.q_proj": self.colwise_parallel(),
+            "attn.k_proj": self.colwise_parallel(),
+            "attn.v_proj": self.colwise_parallel(),
+            "attn.o_proj": self.rowwise_parallel(output_layouts=Shard(1)),
+        }
+
+
+class GLATPPlan(TPPlan):
+
+    @property
+    def attn_plan(self):
+        return {
+            "attn": self.prepare_module_input(
+                input_kwarg_layouts={"hidden_states": Shard(1)},
+                desired_input_kwarg_layouts={"hidden_states": Replicate()},
+            ),
+            "attn.q_proj": self.colwise_parallel(),
+            "attn.k_proj": self.colwise_parallel(),
+            "attn.v_proj": self.colwise_parallel(),
+            "attn.g_proj": self.colwise_parallel(),
+            "attn.gk_proj.0": PrepareModuleWeight(layouts=Replicate()),
+            "attn.gk_proj.1": self.colwise_parallel(),
+            "attn.g_norm": SequenceParallel(sequence_dim=-1),
+            "attn.o_proj": self.rowwise_parallel(output_layouts=Shard(1)),
+        }
+
+
+TP_PLAN_MAP = {"transformer": TransformerTPPlan, "gla": GLATPPlan}
+
+
+def apply_tp(
+    model: nn.Module,
+    tp_mesh: DeviceMesh,
+    loss_parallel: bool,
+    enable_float8: bool,
+    enable_async_tp: bool,
+):
+    """Apply tensor parallelism."""
+    # 1. Parallelize the embedding and shard its outputs (which are the first
+    # transformer block's inputs)
+    # 2. Parallelize the root norm layer over the sequence dim
+    # 3. Parallelize the final linear output layer
+    tp_plan = TP_PLAN_MAP[model.config.model_type](
+        model, loss_parallel=loss_parallel, enable_float8=enable_float8
+    )
+    parallelize_module(model, tp_mesh, tp_plan.model_plan)
+
+    blocks = get_blocks(model)
+    if blocks is None:
+        logger.warning("No block found for tensor parallelism")
+    else:
+        for _, block in enumerate(blocks):
+            parallelize_module(
+                module=block,
+                device_mesh=tp_mesh,
+                parallelize_plan=tp_plan.layer_plan,
+            )
+
+    if enable_async_tp:
+        from torch.distributed._symmetric_memory import enable_symm_mem_for_group
+
+        torch._inductor.config._micro_pipeline_tp = True
+        enable_symm_mem_for_group(tp_mesh.get_group().group_name)
+
+    logger.info(
+        f"Applied {'Float8 ' if enable_float8 else ''}{'Async ' if enable_async_tp else ''}"
+        "Tensor Parallelism to the model"
+    )
+
+
+# for selective op activation checkpointing
+_save_list = {
+    torch.ops.aten.mm.default,
+    torch.ops.aten._scaled_dot_product_efficient_attention.default,
+    torch.ops.aten._scaled_dot_product_flash_attention.default,
+    torch.ops._c10d_functional.reduce_scatter_tensor.default,
+    # for low precision training, it's useful to always save
+    # the result of max, since the absolute maximum is
+    # used to compute the scaling factor for quantization.
+    torch.ops.aten.max.default,
+}
+
+
+def _apply_ac_to_block(module: nn.Module, ac_config):
+    valid_ac_modes = ("full", "selective")
+    if ac_config.mode not in valid_ac_modes:
+        raise ValueError(
+            f"Invalid AC mode: {ac_config.mode}. Valid modes: {valid_ac_modes}"
+        )
+
+    if ac_config.mode == "full":
+        return ptd_checkpoint_wrapper(module, preserve_rng_state=False)
+
+    assert ac_config.mode == "selective", f"{ac_config.mode}"
+    use_op_sac = ac_config.selective_ac_option == "op"
+    use_layer_sac = ac_config.selective_ac_option.isdigit()
+    if not use_op_sac and not use_layer_sac:
+        raise ValueError(
+            f"Invalid selective AC option: {ac_config.selective_ac_option}. "
+            f"Valid options: 'op' or a positive int representing layer frequency"
+        )
+    if use_op_sac:
+        from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts
+
+        def _get_custom_policy(meta):
+            def _custom_policy(ctx, func, *args, **kwargs):
+                mode = "recompute" if ctx.is_recompute else "forward"
+                mm_count_key = f"{mode}_mm_count"
+                if func == torch.ops.aten.mm.default:
+                    meta[mm_count_key] += 1
+                # Saves output of all compute ops, except every second mm
+                to_save = func in _save_list and not (
+                    func == torch.ops.aten.mm.default and meta[mm_count_key] % 2 == 0
+                )
+                return (
+                    CheckpointPolicy.MUST_SAVE
+                    if to_save
+                    else CheckpointPolicy.PREFER_RECOMPUTE
+                )
+
+            return _custom_policy
+
+        def selective_checkpointing_context_fn():
+            meta = defaultdict(int)
+            return create_selective_checkpoint_contexts(_get_custom_policy(meta))
+
+        return ptd_checkpoint_wrapper(
+            module,
+            context_fn=selective_checkpointing_context_fn,
+            preserve_rng_state=False,
+        )
+    elif use_layer_sac:
+        # Checkpoint every `ac_freq` of the modules passed to this function
+        ac_freq = int(ac_config.selective_ac_option)
+        ptd_checkpoint_wrapper.__dict__.setdefault("_count", 0)
+        ptd_checkpoint_wrapper._count += 1
+        if not ac_freq or ptd_checkpoint_wrapper._count % ac_freq == 0:
+            return ptd_checkpoint_wrapper(module, preserve_rng_state=False)
+        else:
+            return module
+
+
+def apply_ac(model: nn.Module, ac_config):
+    """Apply activation checkpointing to the model."""
+    blocks = get_blocks(model)
+    if blocks is None:
+        logger.warning("No block found for activation checkpointing")
+        return
+
+    for layer_id, block in blocks.named_children():
+        block = _apply_ac_to_block(block, ac_config)
+        blocks.register_module(layer_id, block)
+
+    logger.info(f"Applied {ac_config.mode} activation checkpointing to the model")
+
+
+def apply_compile(model: nn.Module):
+    """
+    Apply torch.compile to each block, which makes compilation efficient due to
+    repeated structure. Alternatively one can compile the whole model (after applying DP).
+    """
+
+    blocks = get_blocks(model)
+    if blocks is None:
+        logger.warning("No block found for torch.compile")
+    else:
+        for layer_id, block in blocks.named_children():
+            block = torch.compile(block)
+            blocks.register_module(layer_id, block)
+        logger.info("Compiling each block with torch.compile")
+
+    real_model = get_model(model)
+
+    logger.info("Compiling the embedding, norm, and lm_head layers with torch.compile")
+    embeddings_key = get_components_name(real_model, "tok_embeddings")
+    if embeddings_key is not None:
+        embeddings = torch.compile(getattr(real_model, embeddings_key), fullgraph=True)
+        real_model.register_module(embeddings_key, embeddings)
+
+    norm_key = get_components_name(real_model, "norm")
+    if norm_key is not None:
+        norm = torch.compile(getattr(real_model, norm_key), fullgraph=True)
+        real_model.register_module(norm_key, norm)
+
+    lm_head_key = get_components_name(model, "lm_head")
+    if lm_head_key is not None:
+        lm_head = torch.compile(getattr(model, lm_head_key), fullgraph=True)
+        model.register_module(lm_head_key, lm_head)
+
+    logger.info("Compiling the entire model with torch.compile")
+    model = torch.compile(model)
+
+
+def apply_fsdp(
+    model: nn.Module,
+    dp_mesh: DeviceMesh,
+    param_dtype: torch.dtype,
+    reduce_dtype: torch.dtype,
+    pp_enabled: bool,
+    cpu_offload: bool = False,
+    reshard_after_forward_policy: str = "default",
+):
+    """
+    Apply data parallelism (via FSDP2) to the model.
+
+    Args:
+        model (nn.Module): The model to apply data parallelism to.
+        dp_mesh (DeviceMesh): The device mesh to use for data parallelism.
+        param_dtype (torch.dtype): The data type to use for model parameters.
+        reduce_dtype (torch.dtype): The data type to use for reduction operations.
+        pp_enabled (bool): Whether pipeline parallelism is enabled.
+        cpu_offload (bool, optional): Whether to offload model parameters to CPU. Defaults to False.
+        reshard_after_forward_policy (str, optional):
+            The policy to use for resharding after forward pass. Defaults to "default".
+            Other options: "never", "always".
+            - "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.
+
+    """
+    mp_policy = MixedPrecisionPolicy(param_dtype=param_dtype, reduce_dtype=reduce_dtype)
+    fsdp_config = {"mesh": dp_mesh, "mp_policy": mp_policy}
+    if cpu_offload:
+        fsdp_config["offload_policy"] = CPUOffloadPolicy()
+
+    blocks = get_blocks(model)
+    if blocks is None:
+        logger.warning("No block found for FSDP")
+    else:
+        total_blocks = len(blocks)
+        for layer_id, block in enumerate(blocks):
+            if reshard_after_forward_policy == "always":
+                reshard_after_forward = True
+            elif reshard_after_forward_policy == "never":
+                reshard_after_forward = False
+            elif reshard_after_forward_policy == "default":
+                if pp_enabled:
+                    # For PP, do not reshard after forward to avoid per-microbatch
+                    # all-gathers, which can be expensive and non-overlapped
+                    reshard_after_forward = False
+                else:
+                    # As an optimization, do not reshard after forward for the last
+                    # transformer block since FSDP would prefetch it immediately
+                    reshard_after_forward = int(layer_id) < total_blocks - 1
+            else:
+                raise ValueError(
+                    f"Invalid reshard_after_forward_policy: {reshard_after_forward_policy}."
+                )
+            fully_shard(
+                block,
+                **fsdp_config,
+                reshard_after_forward=reshard_after_forward,
+            )
+
+    fully_shard(model, **fsdp_config, reshard_after_forward=not pp_enabled)
+
+
+def apply_ddp(
+    model: nn.Module,
+    dp_mesh: DeviceMesh,
+    enable_compile: bool,
+    enable_compiled_autograd: bool,
+):
+    if enable_compile:
+        if enable_compiled_autograd:
+            torch._dynamo.config.optimize_ddp = (
+                "python_reducer_without_compiled_forward"
+            )
+        else:
+            torch._dynamo.config.optimize_ddp = "ddp_optimizer"
+
+    replicate(model, device_mesh=dp_mesh, bucket_cap_mb=100)
+
+    logger.info("Applied DDP to the model")
+
+
+def get_model(model):
+    base_model_prefix = getattr(model, "base_model_prefix", "model")
+    if not hasattr(model, base_model_prefix):
+        return None
+    model = getattr(model, base_model_prefix)
+    return model
+
+
+def get_blocks(model):
+    # TODO[flame]: adapt for network not using 'layers' attribute
+    model = get_model(model)
+    if not hasattr(model, "layers"):
+        logger.warning('no "layers" in model can be found')
+        return None
+    return model.layers
+
+
+def get_components_name(model, component_name):
+    """
+    We try to catch tok_embeddings, norm layers and lm_head layers
+    We do not catch the layer names in the blocks, for blocks see `get_blocks`
+    We assume the model has the following structure:
+    LlamaForCausalLM:
+        Model:
+            embed_tokens,
+            layers,
+            norm,
+        lm_head
+    ***
+    so, to search 'tok_embeddings' and 'norm' we need to pass `get_model(model)`
+    and for 'lm_head' we need to pass `model`
+    ***
+    """
+
+    if component_name == "tok_embeddings":
+        if hasattr(model, "tok_embeddings"):
+            return "tok_embeddings"
+        elif hasattr(model, "embed_tokens"):
+            return "embed_tokens"
+        elif hasattr(model, "embeddings"):
+            return "embeddings"
+        else:
+            logger.warning("No tok_embeddings found in model")
+            return None
+
+    elif component_name == "norm":
+        if hasattr(model, "norm"):
+            return "norm"
+        elif hasattr(model, "norms"):
+            return "norms"
+        elif hasattr(model, "layernorm"):
+            return "layernorm"
+        else:
+            logger.warning("No norm found in model")
+            return None
+
+    elif component_name == "lm_head":
+        if hasattr(model, "lm_head"):
+            return "lm_head"
+        else:
+            logger.warning("No lm_head found in model")
+            return None

flame/models/pipeline_fla.py

@@ -0,0 +1,162 @@
+# 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.
+
+# This file applies the PT-D pipeline parallelism to the Llama model.
+
+import copy
+from typing import Callable, Optional, Union
+
+import torch
+import torch.nn as nn
+from torch.distributed import DeviceMesh
+from torch.distributed.pipelining import PipelineStage
+from torch.distributed.pipelining.schedules import ScheduleZBVZeroBubble, _PipelineSchedule, get_schedule_class
+from transformers import PretrainedConfig
+
+from flame.models.parallelize_fla import get_blocks, get_components_name, get_model
+from torchtitan.config_manager import JobConfig
+from torchtitan.distributed.parallel_dims import ParallelDims
+from torchtitan.distributed.pipeline import build_pipeline_schedule, generate_split_points, stage_ids_this_rank
+from torchtitan.tools.logging import logger
+
+DeviceType = Union[int, str, torch.device]
+
+
+def pipeline_fla(
+    model: nn.Module,
+    pp_mesh: DeviceMesh,
+    parallel_dims: ParallelDims,
+    job_config: JobConfig,
+    device: DeviceType,
+    model_config: PretrainedConfig,
+    loss_fn: Callable[..., torch.Tensor],
+) -> tuple[_PipelineSchedule, list[nn.Module], bool, bool]:
+    stages, models = pipeline_fla_manual_split(
+        model, pp_mesh, parallel_dims, job_config, device, model_config
+    )
+
+    pp_schedule = build_pipeline_schedule(job_config, stages, loss_fn)
+
+    # This is used in the train loop to determine whether to pass in the input_ids and labels
+    has_first_stage = False
+    has_last_stage = False
+    for stage in stages:
+        if stage.is_first:
+            has_first_stage = True
+        if stage.is_last:
+            has_last_stage = True
+
+    return pp_schedule, models, has_first_stage, has_last_stage
+
+
+def pipeline_fla_manual_split(
+    whole_model: nn.Module,
+    pp_mesh: DeviceMesh,
+    parallel_dims: ParallelDims,
+    job_config: JobConfig,
+    device: DeviceType,
+    model_config: PretrainedConfig,
+) -> tuple[list[PipelineStage], list[nn.Module]]:
+    """
+    This API extracts one torch.nn.Module objects for the part of the model configured to run inside this stage.
+
+    It wraps the model chunk in a ManualPipelineStage object and returns both the stage and model objects.
+
+    The stage object is used to create a pipeline schedule, and the model object can be used for applying SPMD
+    parallelism.
+    """
+    pp_rank = pp_mesh.get_local_rank()
+    pp_size = pp_mesh.size()
+
+    splits = (
+        job_config.experimental.pipeline_parallel_split_points
+        or generate_split_points(
+            job_config, parallel_dims.pp, model_config.num_hidden_layers
+        )
+    )
+
+    def _build_stage(
+        stage_idx: int,
+        start_layer: Optional[str],
+        stop_layer: Optional[str],
+        is_first: bool = False,
+        is_last: bool = False,
+    ) -> tuple[PipelineStage, nn.Module]:
+        model = copy.deepcopy(whole_model)
+        if not is_first:
+            # we do `model.tok_embeddings = None` here
+            real_model = get_model(model)
+            tok_embeddings_name = get_components_name(real_model, "tok_embeddings")
+            setattr(real_model, tok_embeddings_name, None)
+
+        drop_layers = start_layer is not None
+        # Get module dictionary from get_blocks(model)
+        # and Create a list of keys before modifying dictionary
+        module_dict = get_blocks(model)._modules  # Store reference
+        layer_names = list(module_dict.keys())
+
+        # Iterate over the list of keys instead of `_modules.items()`
+        for name in layer_names:
+            # Dynamically determine prefix (blocks.* or layers.*)
+            prefix = start_layer.split(".")[0] if start_layer else "layers"
+            layer_name = f"{prefix}.{name}"  # Construct the correct name format
+
+            # Ensure `drop_layers` activation is based on actual naming
+            if layer_name == start_layer:
+                drop_layers = False
+            if layer_name == stop_layer:
+                drop_layers = True
+
+            # Delete layer if drop_layers is active
+            if drop_layers:
+                del module_dict[name]  # Safe deletion from stored dictionary
+
+        if not is_last:
+            # we do `model.norm = None` and `model.output = None`
+            real_model = get_model(model)
+            norm_name = get_components_name(real_model, "norm")
+            setattr(real_model, norm_name, None)
+
+            head_name = get_components_name(model, "lm_head")
+            setattr(model, head_name, None)
+
+        stage = PipelineStage(
+            model,
+            stage_idx,
+            num_stages,
+            device,
+            group=pp_mesh.get_group("pp"),
+        )
+        return stage, model
+
+    num_stages = len(splits) + 1
+    stage_idx = pp_rank
+
+    stages = []
+    models = []
+
+    schedule_class = get_schedule_class(
+        job_config.experimental.pipeline_parallel_schedule
+    )
+    style = "v" if schedule_class == ScheduleZBVZeroBubble else "loop"
+
+    for stage_idx in stage_ids_this_rank(pp_rank, pp_size, num_stages, style=style):
+        start_layer = splits[stage_idx - 1] if stage_idx > 0 else None
+        stop_layer = splits[stage_idx] if stage_idx < num_stages - 1 else None
+        stage, model_chunk = _build_stage(
+            stage_idx,
+            start_layer,
+            stop_layer,
+            is_first=stage_idx == 0,
+            is_last=stage_idx == num_stages - 1,
+        )
+        logger.info(
+            f"PP rank {pp_rank} is building stage_idx {stage_idx}"
+            f" with start_layer {start_layer}, stop_layer {stop_layer}"
+        )
+        stages.append(stage)
+        models.append(model_chunk)
+    return stages, models

flame/tools/utils.py

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

flame/train.py

@@ -0,0 +1,637 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import os
+import time
+from datetime import timedelta
+
+import fla  # noqa
+import torch
+from fla.modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss
+from fla.ops.utils import prepare_position_ids
+from torch.distributed.elastic.multiprocessing.errors import record
+from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.ft import FTParallelDims, init_ft_manager
+from torchtitan.components.loss import build_cross_entropy_loss
+from torchtitan.components.lr_scheduler import build_lr_schedulers
+from torchtitan.components.metrics import build_device_memory_monitor, build_metrics_processor, ensure_pp_loss_visible
+from torchtitan.components.optimizer import build_optimizers
+from torchtitan.distributed import ParallelDims
+from torchtitan.distributed import utils as dist_utils
+from torchtitan.protocols.model_converter import build_model_converters
+from torchtitan.protocols.train_spec import TrainSpec, get_train_spec, register_train_spec
+from torchtitan.tools import utils
+from torchtitan.tools.logging import init_logger, logger
+from torchtitan.tools.profiling import maybe_enable_memory_snapshot, maybe_enable_profiling
+
+import custom_models
+from flame.components.checkpoint import TrainState
+from flame.config_manager import JobConfig
+from flame.data import build_dataloader, build_dataset
+from flame.models.parallelize_fla import parallelize_fla
+from flame.models.pipeline_fla import pipeline_fla
+from flame.tools.utils import get_nparams_and_flops
+
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from fla.models import HamiltonForCausalLM as NewModelForCausalLM, HamiltonConfig as NewConfig
+# from fla.models import GLAForCausalLM as NewModelForCausalLM, GLAConfig as NewConfig
+MODEL_TYPE = NewConfig.model_type
+
+
+def build_tokenizer(job_config: JobConfig) -> AutoTokenizer:
+    return AutoTokenizer.from_pretrained(job_config.model.tokenizer_path)
+
+
+register_train_spec(
+    TrainSpec(
+        name="fla",
+        cls=AutoModelForCausalLM,
+        config=AutoConfig,
+        parallelize_fn=parallelize_fla,
+        pipelining_fn=pipeline_fla,
+        build_optimizers_fn=build_optimizers,
+        build_lr_schedulers_fn=build_lr_schedulers,
+        build_dataloader_fn=build_dataloader,
+        build_tokenizer_fn=build_tokenizer,
+        build_loss_fn=build_cross_entropy_loss,
+    )
+)
+
+
+# Enable debug tracing on failure: https://pytorch.org/docs/stable/elastic/errors.html
+@record
+def main(job_config: JobConfig):
+    logger.info(f"Starting job: {job_config.job.description}")
+    logger.info(f"Registering model type: {MODEL_TYPE}")
+
+    if job_config.experimental.custom_model_path:
+        utils.import_module_from_path(job_config.experimental.custom_model_path)
+
+    # used for colorful printing
+    color = utils.NoColor if job_config.metrics.disable_color_printing else utils.Color
+
+    if job_config.job.print_args:
+        logger.info(
+            f"{color.green}{json.dumps(job_config.to_dict(), indent=2, sort_keys=True)}{color.reset}"
+        )
+
+    # take control of garbage collection to avoid stragglers
+    gc_handler = utils.GarbageCollection(gc_freq=job_config.training.gc_freq)
+
+    device_module, device_type = utils.device_module, utils.device_type
+    device = torch.device(f"{device_type}:{int(os.environ['LOCAL_RANK'])}")
+    # Device has to be set before creating TorchFT manager.
+    device_module.set_device(device)
+    ft_manager = init_ft_manager(job_config)
+
+    # init distributed
+    world_size = int(os.environ["WORLD_SIZE"])
+    if not ft_manager.enabled:
+        parallel_dims = ParallelDims(
+            dp_shard=job_config.training.data_parallel_shard_degree,
+            dp_replicate=job_config.training.data_parallel_replicate_degree,
+            cp=job_config.experimental.context_parallel_degree,
+            tp=job_config.training.tensor_parallel_degree,
+            pp=job_config.experimental.pipeline_parallel_degree,
+            world_size=world_size,
+            enable_loss_parallel=not job_config.training.disable_loss_parallel,
+        )
+    else:
+        parallel_dims = FTParallelDims(
+            dp_shard=job_config.training.data_parallel_shard_degree,
+            dp_replicate=job_config.training.data_parallel_replicate_degree,
+            cp=job_config.experimental.context_parallel_degree,
+            tp=job_config.training.tensor_parallel_degree,
+            pp=job_config.experimental.pipeline_parallel_degree,
+            world_size=world_size,
+            enable_loss_parallel=not job_config.training.disable_loss_parallel,
+            ft_manager=ft_manager,
+        )
+    dist_utils.init_distributed(job_config)
+    # initialize device memory monitor and get peak flops for MFU calculation
+    device_memory_monitor = build_device_memory_monitor()
+    gpu_peak_flops = utils.get_peak_flops(device_memory_monitor.device_name)
+    logger.info(f"Peak FLOPS used for computing MFU: {gpu_peak_flops:.3e}")
+
+    # build meshes
+    world_mesh = parallel_dims.build_mesh(device_type=device_type)
+    if parallel_dims.dp_enabled:
+        dp_mesh = world_mesh["dp"]
+        dp_degree, dp_rank = dp_mesh.size(), dp_mesh.get_local_rank()
+    else:
+        dp_degree, dp_rank = 1, 0
+
+    if parallel_dims.pp_enabled:
+        raise NotImplementedError(
+            "Pipeline parallelism is not supported in this version"
+        )
+        """
+        ! TODO[flame]: We need to fix the pipeline parallelism for flame
+        [x] Match the key of models' components with the actual naming
+        [ ] Fix the post-init and tie-embedding for pipeline parallelism, HF's transformer automatically
+            forces to tie if head is None, we need to handle this case
+        [ ]
+        """
+        pp_mesh = world_mesh["pp"]
+
+    # Set random seed, and maybe enable deterministic mode (mainly for debugging, expect perf loss)
+    dist_utils.set_determinism(
+        world_mesh, device, job_config.training.seed, job_config.training.deterministic
+    )
+    train_spec = get_train_spec(job_config.model.name)
+
+    logger.info("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        job_config.model.tokenizer_path,
+        trust_remote_code=True,
+        model_max_length=int(1e10),
+    )
+    logger.info(f"{tokenizer}")
+    logger.info(
+        f"Loading dataset {job_config.training.dataset}"
+        f":{job_config.training.dataset_name}"
+        if job_config.training.dataset_name is not None
+        else ""
+    )
+    dataset = build_dataset(
+        dataset=job_config.training.dataset,
+        dataset_name=job_config.training.dataset_name,
+        dataset_split=job_config.training.dataset_split,
+        data_dir=job_config.training.data_dir,
+        data_files=job_config.training.data_files,
+        data_probs=job_config.training.data_probs,
+        streaming=job_config.training.streaming,
+        dp_degree=dp_degree,
+        num_workers=job_config.training.num_workers,
+        seed=job_config.training.seed,
+    )
+
+    logger.info("Building dataloader...")
+    dataloader = build_dataloader(
+        dataset=dataset,
+        tokenizer=tokenizer,
+        rank=dp_rank,
+        world_size=dp_degree,
+        batch_size=job_config.training.batch_size,
+        seq_len=job_config.training.seq_len,
+        context_len=job_config.training.context_len,
+        varlen=job_config.training.varlen,
+        num_workers=job_config.training.num_workers,
+        pin_memory=job_config.training.pin_memory,
+        persistent_workers=job_config.training.persistent_workers,
+        snapshot_every_n_steps=job_config.checkpoint.interval,
+    )
+
+
+    logger.info(f"Loading model config from {job_config.model.config}")
+    logger.info(f"Registering model type: {MODEL_TYPE}")
+    AutoConfig.register(MODEL_TYPE, NewConfig)   # important!
+    AutoModelForCausalLM.register(NewConfig, NewModelForCausalLM)  # important!
+    model_config = AutoConfig.from_pretrained(job_config.model.config)
+    # set the model configs from training inputs:
+    # 1. norm type to decide which norm layer to use
+    # 2. disable fused norm if TP is enabled
+    # 3. vocab size from tokenizer
+    # 4. context_len base on inputs
+    if parallel_dims.tp_enabled:
+        if model_config.fuse_norm:
+            logger.warning(
+                f"{color.red}"
+                f"Fused norm is not compatible with tensor parallelism. "
+                f"Disabling it for now."
+                f"{color.reset}"
+            )
+            model_config.fuse_norm = False
+    if parallel_dims.loss_parallel_enabled:
+        if model_config.fuse_linear_cross_entropy:
+            logger.warning(
+                f"{color.red}"
+                f"Loss parallel enabled. Disabling fused cross entropy for now."
+                f"{color.reset}"
+            )
+            model_config.fuse_linear_cross_entropy = False
+    model_config.vocab_size = max(tokenizer.vocab_size, model_config.vocab_size)
+
+
+
+    logger.info(
+        f"Building model from the config\n{color.green}{model_config}{color.reset}"
+    )
+    with torch.device("meta"):
+        model = AutoModelForCausalLM.from_config(model_config)
+        if (
+            getattr(model_config, "fuse_linear_cross_entropy", False)
+            and FusedLinearCrossEntropyLoss is not None
+        ):
+            model.criterion = FusedLinearCrossEntropyLoss(
+                num_chunks=8 // parallel_dims.tp
+            )
+        # defer weight initialization until after parallelisms are applied
+        model.apply(lambda m: setattr(m, "_is_hf_initialized", False))
+    logger.info(f"{color.blue}\n{model}{color.reset}\n")
+
+    logger.info("Applying model converters...")
+
+    # Build the collection of model converters. No-op if `model.converters` empty
+    model_converters = build_model_converters(job_config, parallel_dims)
+    model_converters.convert(model)
+
+    # calculate model size and flops per token
+    model_param_count, num_flops_per_token = get_nparams_and_flops(
+        model, model_config, job_config.training.context_len
+    )
+
+    # move sharded model to CPU/GPU and initialize weights via DTensor
+    if job_config.checkpoint.create_seed_checkpoint:
+        init_device = "cpu"
+    elif job_config.training.enable_cpu_offload:
+        init_device = "cpu"
+    else:
+        init_device = device_type
+
+    # apply parallelisms and initialization
+    if parallel_dims.pp_enabled:
+        # apply PT-D Pipeline Parallel
+        (
+            pp_schedule,
+            model_parts,
+            has_first_stage,
+            has_last_stage,
+        ) = train_spec.pipelining_fn(
+            model,
+            pp_mesh,
+            parallel_dims,
+            job_config,
+            device,
+            model_config,
+            train_spec.loss_fn,
+        )
+        # when PP is enabled, `model` obj is no longer used after this point, model_parts is used instead
+        del model
+
+        # For PP with looped schedules, each item in model_parts is one stage-model-chunk.
+        # We need to iterate through model_parts to apply SPMD parallelisms, compilation,
+        # optimizer, and checkpointing
+        for m in model_parts:
+            # apply SPMD-style PT-D techniques
+            train_spec.parallelize_fn(m, world_mesh, parallel_dims, job_config)
+            m.to_empty(device=init_device)
+            with torch.no_grad():
+                m.post_init()
+            m.train()
+
+        # confirm that user will be able to view loss metrics on the console
+        ensure_pp_loss_visible(parallel_dims, job_config, color)
+    else:
+        # apply PT-D Tensor Parallel, activation checkpointing, torch.compile, Data Parallel
+        train_spec.parallelize_fn(model, world_mesh, parallel_dims, job_config)
+        model.to_empty(device=init_device)
+        with torch.no_grad():
+            model.post_init()
+        model.train()
+
+        model_parts = [model]
+
+    device_mem_stats = device_memory_monitor.get_peak_stats()
+    logger.info(
+        f"{device_type.upper()} memory usage for model: "
+        f"{device_mem_stats.max_reserved_gib:.2f}GiB"
+        f"({device_mem_stats.max_reserved_pct:.2f}%)"
+    )
+
+    # build optimizer after applying parallelisms to the model
+    optimizers = train_spec.build_optimizers_fn(model_parts, job_config, ft_manager)
+    lr_schedulers = train_spec.build_lr_schedulers_fn(optimizers, job_config)
+    # Post optimizer step model converters hook.
+    # e.g. calculate float8 dynamic amax/scale for all-parameter for FSDP2
+    # where it issues a single all-reduce for all parameters at once for better performance
+    optimizers.register_step_post_hook(
+        lambda *args, **kwargs: model_converters.post_optimizer_hook(model_parts)
+    )
+
+    train_state = TrainState()
+
+    # load initial checkpoint
+    checkpoint = CheckpointManager(
+        dataloader=dataloader,
+        model_parts=model_parts,
+        optimizers=optimizers,
+        lr_schedulers=lr_schedulers,
+        states={"train_state": train_state},
+        job_config=job_config,
+        ft_manager=ft_manager,
+    )
+
+    if job_config.checkpoint.create_seed_checkpoint:
+        assert world_size == 1, (
+            "Must create seed checkpoint using a single device, to disable sharding"
+        )
+        assert job_config.checkpoint.enable_checkpoint, (
+            "Must enable checkpointing when creating a seed checkpoint"
+        )
+        checkpoint.save(curr_step=0, force=True)
+        logger.info("Created seed checkpoint")
+        return
+
+    checkpoint.load(step=job_config.checkpoint.load_step)
+    metric_logger = build_metrics_processor(job_config, parallel_dims)
+    # Set dependent attributes for metric_logger
+    metric_logger.num_flops_per_token = num_flops_per_token
+    metric_logger.optimizers = optimizers  # Pass optimizers if needed by logger logic
+    metric_logger.lr_schedulers = (
+        lr_schedulers  # Pass schedulers if needed by logger logic
+    )
+
+    # plot losses loaded from checkpoint (if any) to TensorBoard
+    # NOTE: Loss info after the last log step before checkpoint saving will not be ploted.
+    #       This can be avoided by setting checkpoint.interval to be a multiple of metrics.log_freq
+    if train_state.step > 0 and len(metric_logger.data_loading_times) > 0:
+        for idx, step in enumerate(train_state.log_steps):
+            metric_logger.log(
+                step,
+                global_avg_loss=train_state.global_avg_losses[idx],
+                global_max_loss=train_state.global_max_losses[idx],
+            )
+
+    data_iterator = iter(dataloader)
+
+    train_context = dist_utils.get_train_context(
+        parallel_dims.loss_parallel_enabled,
+        job_config.experimental.enable_compiled_autograd,
+    )
+    maybe_enable_amp = dist_utils.maybe_enable_amp(
+        parallel_dims,
+        job_config.training.mixed_precision_param,
+        device_type,
+    )
+
+    # variables used to keep info for metrics logging
+    device_memory_monitor.reset_peak_stats()
+
+    global_batch_size = (
+        job_config.training.batch_size
+        * dp_degree
+        * job_config.training.gradient_accumulation_steps
+    )
+    num_tokens_per_step = global_batch_size * job_config.training.seq_len
+    # train loop
+    logger.info(f"{color.red}***** Running training *****{color.reset}")
+    logger.info(f"{color.green}  Training starts at step {train_state.step + 1}")
+    logger.info(
+        f"{color.green}  Number of tokens per sequence = {job_config.training.seq_len:,}"
+    )
+    logger.info(
+        f"{color.green}  Gradient Accumulation steps = {job_config.training.gradient_accumulation_steps}"
+    )
+    logger.info(
+        f"{color.green}  Instantaneous batch size (per device) = {job_config.training.batch_size:,}"
+    )
+    logger.info(
+        f"{color.green}  Global batch size (w. parallel, distributed & accumulation) = {global_batch_size:,}"
+        f" ({num_tokens_per_step:,} tokens)"
+    )
+    logger.info(
+        f"{color.green}  Total optimization steps = {job_config.training.steps:,} "
+        f"({job_config.training.steps * num_tokens_per_step:,} tokens)"
+    )
+    logger.info(
+        f"{color.green}  Warmup steps = {job_config.lr_scheduler.warmup_steps:,}"
+        f" ({job_config.lr_scheduler.warmup_steps * num_tokens_per_step:,} tokens)"
+    )
+    logger.info(
+        f"{color.green}  Number of parameters = {model_param_count:,} {color.reset}"
+    )
+
+    with (
+        maybe_enable_profiling(
+            job_config, global_step=train_state.step
+        ) as torch_profiler,
+        maybe_enable_memory_snapshot(
+            job_config, global_step=train_state.step
+        ) as memory_profiler,
+    ):
+        while train_state.step < job_config.training.steps:
+            train_state.step += 1
+            gc_handler.run(train_state.step)
+
+            optimizers.zero_grad()
+
+            losses = []
+            # do gradient accumulation if enabled
+            for _ in range(job_config.training.gradient_accumulation_steps):
+                # get batch
+                data_load_start = time.perf_counter()
+                batch = next(data_iterator)
+                input_ids, labels = batch["input_ids"], batch["labels"]
+
+                # Update metrics processor state before forward/backward
+                metric_logger.ntokens_since_last_log += labels.numel()
+                metric_logger.data_loading_times.append(
+                    time.perf_counter() - data_load_start
+                )
+
+                input_ids = input_ids.to(device_type)
+
+                """
+                TODO[flame]: We need to carefully handle the position_ids for TP/CP
+                Depending on the Models'PE, the position_ids might be different.
+
+                e.g. for TP
+                    For RoPE, all ranks have the same position_ids. [FOR HF model]
+                    For sinusoidal, each rank has the coresponding chunked  position_ids. [FOR HF model]
+
+                e.g. for CP, [optional_context_parallel_ctx shoudl automatically distbute the position_ids]
+                    Each rank has the coresponding chunked position_ids. [FOR All model]
+
+                """
+                labels = labels.to(device_type)
+                cu_seqlens = (
+                    batch["cu_seqlens"].to(device_type)
+                    if "cu_seqlens" in batch
+                    else None
+                )
+                if cu_seqlens is not None:
+                    position_ids = prepare_position_ids(cu_seqlens).to(torch.int32)
+                else:
+                    position_ids = (
+                        torch.arange(0, input_ids.shape[1], device=device_type)
+                        .repeat(input_ids.shape[0], 1)
+                        .to(torch.int32)
+                    )
+                # apply context parallelism if cp is enabled
+                # ensure CP handles the separate freqs_cis buffer for each pp stage
+                optional_context_parallel_ctx = (
+                    dist_utils.create_context_parallel_ctx(
+                        cp_mesh=world_mesh["cp"],
+                        cp_buffers=[input_ids, labels, position_ids],
+                        cp_seq_dims=[1, 1, 1],
+                        cp_no_restore_buffers={input_ids, labels, position_ids},
+                        cp_rotate_method=job_config.experimental.context_parallel_rotate_method,
+                    )
+                    if parallel_dims.cp_enabled
+                    else None
+                )
+
+                # #! TODO[flame], we should distribute the position_ids as well with CP
+                if parallel_dims.pp_enabled:
+                    raise NotImplementedError(
+                        "Pipeline parallelism is not supported in this version"
+                    )
+                    # Pipeline Parallel forward / backward inside step() call
+                    with train_context(optional_context_parallel_ctx):
+                        targets, losses = (
+                            (labels, []) if has_last_stage else (None, None)
+                        )
+
+                        if has_first_stage:
+                            pp_schedule.step(input_ids, target=targets, losses=losses)
+                        else:
+                            pp_schedule.step(target=targets, losses=losses)
+
+                    # accumulate losses across pipeline microbatches
+                    # TODO: PP+FSDP unexpectedly puts the loss back to the CPU
+                    loss = (
+                        torch.mean(torch.stack(losses)).to(device)
+                        if has_last_stage
+                        else torch.tensor([-1.0], device=device)
+                    )
+                else:
+                    # Non-PP forward / backward
+                    with train_context(optional_context_parallel_ctx):
+                        with maybe_enable_amp:
+                            output = model(
+                                input_ids=input_ids,
+                                labels=labels,
+                                position_ids=position_ids,
+                                cu_seqlens=cu_seqlens,
+                        )
+                        loss = (
+                            output.loss
+                            / job_config.training.gradient_accumulation_steps
+                        )
+                        loss.backward()
+                        # print('--------------------------')
+
+                losses.append(loss)
+                del batch
+            loss = sum(losses)
+
+            # clip gradients
+            grad_norm = dist_utils.clip_grad_norm_(
+                [p for m in model_parts for p in m.parameters()],
+                job_config.training.max_norm,
+                foreach=True,
+                pp_mesh=pp_mesh if parallel_dims.pp_enabled else None,
+            )
+
+            # optimizer step
+            checkpoint.maybe_wait_for_staging()
+            if job_config.training.skip_nan_inf and (
+                grad_norm.isnan() or grad_norm.isinf()
+            ):
+                logger.warning(
+                    f"Skipping optimizer step - detected invalid gradient norm: {grad_norm:.4f}"
+                )
+                optimizers.zero_grad()
+                train_state.skipped_step += 1
+            else:
+                optimizers.step()
+            lr_schedulers.step()
+
+            # log metrics - Use MetricsProcessor
+            if metric_logger.should_log(train_state.step):
+                if (
+                    parallel_dims.dp_replicate_enabled
+                    or parallel_dims.dp_shard_enabled
+                    or parallel_dims.cp_enabled
+                ):
+                    loss = loss.detach()
+                    # Use dist_mean/max on the accumulated loss for the step
+                    global_avg_loss, global_max_loss = (
+                        dist_utils.dist_mean(
+                            loss,
+                            world_mesh["dp_cp"],
+                        ),
+                        dist_utils.dist_max(
+                            loss,
+                            world_mesh["dp_cp"],
+                        ),
+                    )
+                else:
+                    # Scale back the loss before logging
+                    global_avg_loss = global_max_loss = loss.item()
+
+                # Update train state tokens and elapsed time
+                time_now = time.perf_counter()
+                time_delta = (
+                    time_now - metric_logger.time_last_log
+                )  # Use metric_logger's time
+                train_state.token += (
+                    metric_logger.ntokens_since_last_log  # Use tokens tracked by metric_logger
+                    * parallel_dims.world_size
+                    / parallel_dims.non_data_parallel_size
+                )
+                train_state.elapsed += timedelta(seconds=time_delta)
+                train_state.log_steps.append(train_state.step)
+                train_state.global_avg_losses.append(global_avg_loss)
+                train_state.global_max_losses.append(global_max_loss)
+
+                # Log using the metric processor
+                last_lr = lr_schedulers.schedulers[0].get_last_lr()[0]
+                eta = (
+                    train_state.elapsed
+                    * (job_config.training.steps - train_state.step)
+                    / train_state.step
+                )
+                metric_logger.log(
+                    train_state.step,
+                    global_avg_loss,
+                    global_max_loss,
+                    extra_metrics={
+                        "optimizer/lr": last_lr,
+                        "optimizer/grad_norm": grad_norm.item(),
+                        "optimizer/skipped_step": train_state.skipped_step,
+                    },
+                )
+
+                logger.info(
+                    f"{color.blue}lr: {last_lr:.4e} gnorm: {grad_norm:5.2f} "
+                    f"{color.magenta}[{str(train_state.elapsed).split('.')[0]:>8}<{str(eta).split('.')[0]:>8}]{color.reset}"
+                )
+
+            checkpoint.save(
+                train_state.step, force=(train_state.step == job_config.training.steps)
+            )
+
+            # signal the profiler that the next profiling step has started
+            if torch_profiler:
+                torch_profiler.step()
+            if memory_profiler:
+                memory_profiler.step()
+
+            # reduce timeout after first train step for faster signal
+            # (assuming lazy init and compilation are finished)
+            if train_state.step == 1:
+                dist_utils.set_pg_timeouts(
+                    timeout=timedelta(seconds=job_config.comm.train_timeout_seconds),
+                    world_mesh=world_mesh,
+                )
+
+    if torch.distributed.get_rank() == 0:
+        logger.info("Sleeping 2 seconds for other ranks to complete")
+        time.sleep(2)
+
+    metric_logger.close()
+    logger.info("Training completed")
+
+
+if __name__ == "__main__":
+    init_logger()
+    config = JobConfig()
+    config.parse_args()
+    main(config)
+    torch.distributed.destroy_process_group()

flame/train2.py

@@ -0,0 +1,625 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import os
+import time
+from datetime import timedelta
+
+import fla  # noqa
+import torch
+from fla.modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss
+from fla.ops.utils import prepare_position_ids
+from torch.distributed.elastic.multiprocessing.errors import record
+from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.ft import FTParallelDims, init_ft_manager
+from torchtitan.components.loss import build_cross_entropy_loss
+from torchtitan.components.lr_scheduler import build_lr_schedulers
+from torchtitan.components.metrics import build_device_memory_monitor, build_metrics_processor, ensure_pp_loss_visible
+from torchtitan.components.optimizer import build_optimizers
+from torchtitan.distributed import ParallelDims
+from torchtitan.distributed import utils as dist_utils
+from torchtitan.protocols.model_converter import build_model_converters
+from torchtitan.protocols.train_spec import TrainSpec, get_train_spec, register_train_spec
+from torchtitan.tools import utils
+from torchtitan.tools.logging import init_logger, logger
+from torchtitan.tools.profiling import maybe_enable_memory_snapshot, maybe_enable_profiling
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+import custom_models
+from flame.components.checkpoint import TrainState
+from flame.config_manager import JobConfig
+from flame.data import build_dataloader, build_dataset
+from flame.models.parallelize_fla import parallelize_fla
+from flame.models.pipeline_fla import pipeline_fla
+from flame.tools.utils import get_nparams_and_flops
+
+
+def build_tokenizer(job_config: JobConfig) -> AutoTokenizer:
+    return AutoTokenizer.from_pretrained(job_config.model.tokenizer_path)
+
+
+register_train_spec(
+    TrainSpec(
+        name="fla",
+        cls=AutoModelForCausalLM,
+        config=AutoConfig,
+        parallelize_fn=parallelize_fla,
+        pipelining_fn=pipeline_fla,
+        build_optimizers_fn=build_optimizers,
+        build_lr_schedulers_fn=build_lr_schedulers,
+        build_dataloader_fn=build_dataloader,
+        build_tokenizer_fn=build_tokenizer,
+        build_loss_fn=build_cross_entropy_loss,
+    )
+)
+
+
+# Enable debug tracing on failure: https://pytorch.org/docs/stable/elastic/errors.html
+@record
+def main(job_config: JobConfig):
+    logger.info(f"Starting job: {job_config.job.description}")
+
+    if job_config.experimental.custom_model_path:
+        utils.import_module_from_path(job_config.experimental.custom_model_path)
+
+    # used for colorful printing
+    color = utils.NoColor if job_config.metrics.disable_color_printing else utils.Color
+
+    if job_config.job.print_args:
+        logger.info(
+            f"{color.green}{json.dumps(job_config.to_dict(), indent=2, sort_keys=True)}{color.reset}"
+        )
+
+    # take control of garbage collection to avoid stragglers
+    gc_handler = utils.GarbageCollection(gc_freq=job_config.training.gc_freq)
+
+    device_module, device_type = utils.device_module, utils.device_type
+    device = torch.device(f"{device_type}:{int(os.environ['LOCAL_RANK'])}")
+    # Device has to be set before creating TorchFT manager.
+    device_module.set_device(device)
+    ft_manager = init_ft_manager(job_config)
+
+    # init distributed
+    world_size = int(os.environ["WORLD_SIZE"])
+    if not ft_manager.enabled:
+        parallel_dims = ParallelDims(
+            dp_shard=job_config.training.data_parallel_shard_degree,
+            dp_replicate=job_config.training.data_parallel_replicate_degree,
+            cp=job_config.experimental.context_parallel_degree,
+            tp=job_config.training.tensor_parallel_degree,
+            pp=job_config.experimental.pipeline_parallel_degree,
+            world_size=world_size,
+            enable_loss_parallel=not job_config.training.disable_loss_parallel,
+        )
+    else:
+        parallel_dims = FTParallelDims(
+            dp_shard=job_config.training.data_parallel_shard_degree,
+            dp_replicate=job_config.training.data_parallel_replicate_degree,
+            cp=job_config.experimental.context_parallel_degree,
+            tp=job_config.training.tensor_parallel_degree,
+            pp=job_config.experimental.pipeline_parallel_degree,
+            world_size=world_size,
+            enable_loss_parallel=not job_config.training.disable_loss_parallel,
+            ft_manager=ft_manager,
+        )
+    dist_utils.init_distributed(job_config)
+    # initialize device memory monitor and get peak flops for MFU calculation
+    device_memory_monitor = build_device_memory_monitor()
+    gpu_peak_flops = utils.get_peak_flops(device_memory_monitor.device_name)
+    logger.info(f"Peak FLOPS used for computing MFU: {gpu_peak_flops:.3e}")
+
+    # build meshes
+    world_mesh = parallel_dims.build_mesh(device_type=device_type)
+    if parallel_dims.dp_enabled:
+        dp_mesh = world_mesh["dp"]
+        dp_degree, dp_rank = dp_mesh.size(), dp_mesh.get_local_rank()
+    else:
+        dp_degree, dp_rank = 1, 0
+
+    if parallel_dims.pp_enabled:
+        raise NotImplementedError(
+            "Pipeline parallelism is not supported in this version"
+        )
+        """
+        ! TODO[flame]: We need to fix the pipeline parallelism for flame
+        [x] Match the key of models' components with the actual naming
+        [ ] Fix the post-init and tie-embedding for pipeline parallelism, HF's transformer automatically
+            forces to tie if head is None, we need to handle this case
+        [ ]
+        """
+        pp_mesh = world_mesh["pp"]
+
+    # Set random seed, and maybe enable deterministic mode (mainly for debugging, expect perf loss)
+    dist_utils.set_determinism(
+        world_mesh, device, job_config.training.seed, job_config.training.deterministic
+    )
+    train_spec = get_train_spec(job_config.model.name)
+
+    logger.info("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        job_config.model.tokenizer_path,
+        trust_remote_code=True,
+        model_max_length=int(1e10),
+    )
+    logger.info(f"{tokenizer}")
+    logger.info(
+        f"Loading dataset {job_config.training.dataset}"
+        f":{job_config.training.dataset_name}"
+        if job_config.training.dataset_name is not None
+        else ""
+    )
+    dataset = build_dataset(
+        dataset=job_config.training.dataset,
+        dataset_name=job_config.training.dataset_name,
+        dataset_split=job_config.training.dataset_split,
+        data_dir=job_config.training.data_dir,
+        data_files=job_config.training.data_files,
+        data_probs=job_config.training.data_probs,
+        streaming=job_config.training.streaming,
+        dp_degree=dp_degree,
+        num_workers=job_config.training.num_workers,
+        seed=job_config.training.seed,
+    )
+
+    logger.info("Building dataloader...")
+    dataloader = build_dataloader(
+        dataset=dataset,
+        tokenizer=tokenizer,
+        rank=dp_rank,
+        world_size=dp_degree,
+        batch_size=job_config.training.batch_size,
+        seq_len=job_config.training.seq_len,
+        context_len=job_config.training.context_len,
+        varlen=job_config.training.varlen,
+        num_workers=job_config.training.num_workers,
+        pin_memory=job_config.training.pin_memory,
+        persistent_workers=job_config.training.persistent_workers,
+        snapshot_every_n_steps=job_config.checkpoint.interval,
+    )
+
+
+    logger.info(f"Loading model config from {job_config.model.config}")
+    model_config = AutoConfig.from_pretrained(job_config.model.config)
+    # set the model configs from training inputs:
+    # 1. norm type to decide which norm layer to use
+    # 2. disable fused norm if TP is enabled
+    # 3. vocab size from tokenizer
+    # 4. context_len base on inputs
+    if parallel_dims.tp_enabled:
+        if model_config.fuse_norm:
+            logger.warning(
+                f"{color.red}"
+                f"Fused norm is not compatible with tensor parallelism. "
+                f"Disabling it for now."
+                f"{color.reset}"
+            )
+            model_config.fuse_norm = False
+    if parallel_dims.loss_parallel_enabled:
+        if model_config.fuse_linear_cross_entropy:
+            logger.warning(
+                f"{color.red}"
+                f"Loss parallel enabled. Disabling fused cross entropy for now."
+                f"{color.reset}"
+            )
+            model_config.fuse_linear_cross_entropy = False
+    model_config.vocab_size = max(tokenizer.vocab_size, model_config.vocab_size)
+
+
+
+    logger.info(
+        f"Building model from the config\n{color.green}{model_config}{color.reset}"
+    )
+    with torch.device("meta"):
+        model = AutoModelForCausalLM.from_config(model_config)
+        if (
+            getattr(model_config, "fuse_linear_cross_entropy", False)
+            and FusedLinearCrossEntropyLoss is not None
+        ):
+            model.criterion = FusedLinearCrossEntropyLoss(
+                num_chunks=8 // parallel_dims.tp
+            )
+        # defer weight initialization until after parallelisms are applied
+        model.apply(lambda m: setattr(m, "_is_hf_initialized", False))
+    logger.info(f"{color.blue}\n{model}{color.reset}\n")
+
+    # Build the collection of model converters. No-op if `model.converters` empty
+    model_converters = build_model_converters(job_config, parallel_dims)
+    model_converters.convert(model)
+
+    # calculate model size and flops per token
+    model_param_count, num_flops_per_token = get_nparams_and_flops(
+        model, model_config, job_config.training.context_len
+    )
+
+    # move sharded model to CPU/GPU and initialize weights via DTensor
+    if job_config.checkpoint.create_seed_checkpoint:
+        init_device = "cpu"
+    elif job_config.training.enable_cpu_offload:
+        init_device = "cpu"
+    else:
+        init_device = device_type
+
+    # apply parallelisms and initialization
+    if parallel_dims.pp_enabled:
+        # apply PT-D Pipeline Parallel
+        (
+            pp_schedule,
+            model_parts,
+            has_first_stage,
+            has_last_stage,
+        ) = train_spec.pipelining_fn(
+            model,
+            pp_mesh,
+            parallel_dims,
+            job_config,
+            device,
+            model_config,
+            train_spec.loss_fn,
+        )
+        # when PP is enabled, `model` obj is no longer used after this point, model_parts is used instead
+        del model
+
+        # For PP with looped schedules, each item in model_parts is one stage-model-chunk.
+        # We need to iterate through model_parts to apply SPMD parallelisms, compilation,
+        # optimizer, and checkpointing
+        for m in model_parts:
+            # apply SPMD-style PT-D techniques
+            train_spec.parallelize_fn(m, world_mesh, parallel_dims, job_config)
+            m.to_empty(device=init_device)
+            with torch.no_grad():
+                m.post_init()
+            m.train()
+
+        # confirm that user will be able to view loss metrics on the console
+        ensure_pp_loss_visible(parallel_dims, job_config, color)
+    else:
+        # apply PT-D Tensor Parallel, activation checkpointing, torch.compile, Data Parallel
+        train_spec.parallelize_fn(model, world_mesh, parallel_dims, job_config)
+        model.to_empty(device=init_device)
+        with torch.no_grad():
+            model.post_init()
+        model.train()
+
+        model_parts = [model]
+
+    device_mem_stats = device_memory_monitor.get_peak_stats()
+    logger.info(
+        f"{device_type.upper()} memory usage for model: "
+        f"{device_mem_stats.max_reserved_gib:.2f}GiB"
+        f"({device_mem_stats.max_reserved_pct:.2f}%)"
+    )
+
+    # build optimizer after applying parallelisms to the model
+    optimizers = train_spec.build_optimizers_fn(model_parts, job_config, ft_manager)
+    lr_schedulers = train_spec.build_lr_schedulers_fn(optimizers, job_config)
+    # Post optimizer step model converters hook.
+    # e.g. calculate float8 dynamic amax/scale for all-parameter for FSDP2
+    # where it issues a single all-reduce for all parameters at once for better performance
+    optimizers.register_step_post_hook(
+        lambda *args, **kwargs: model_converters.post_optimizer_hook(model_parts)
+    )
+
+    train_state = TrainState()
+
+    # load initial checkpoint
+    checkpoint = CheckpointManager(
+        dataloader=dataloader,
+        model_parts=model_parts,
+        optimizers=optimizers,
+        lr_schedulers=lr_schedulers,
+        states={"train_state": train_state},
+        job_config=job_config,
+        ft_manager=ft_manager,
+    )
+
+    if job_config.checkpoint.create_seed_checkpoint:
+        assert world_size == 1, (
+            "Must create seed checkpoint using a single device, to disable sharding"
+        )
+        assert job_config.checkpoint.enable_checkpoint, (
+            "Must enable checkpointing when creating a seed checkpoint"
+        )
+        checkpoint.save(curr_step=0, force=True)
+        logger.info("Created seed checkpoint")
+        return
+
+    checkpoint.load(step=job_config.checkpoint.load_step)
+    metric_logger = build_metrics_processor(job_config, parallel_dims)
+    # Set dependent attributes for metric_logger
+    metric_logger.num_flops_per_token = num_flops_per_token
+    metric_logger.optimizers = optimizers  # Pass optimizers if needed by logger logic
+    metric_logger.lr_schedulers = (
+        lr_schedulers  # Pass schedulers if needed by logger logic
+    )
+
+    # plot losses loaded from checkpoint (if any) to TensorBoard
+    # NOTE: Loss info after the last log step before checkpoint saving will not be ploted.
+    #       This can be avoided by setting checkpoint.interval to be a multiple of metrics.log_freq
+    if train_state.step > 0 and len(metric_logger.data_loading_times) > 0:
+        for idx, step in enumerate(train_state.log_steps):
+            metric_logger.log(
+                step,
+                global_avg_loss=train_state.global_avg_losses[idx],
+                global_max_loss=train_state.global_max_losses[idx],
+            )
+
+    data_iterator = iter(dataloader)
+
+    train_context = dist_utils.get_train_context(
+        parallel_dims.loss_parallel_enabled,
+        job_config.experimental.enable_compiled_autograd,
+    )
+    maybe_enable_amp = dist_utils.maybe_enable_amp(
+        parallel_dims,
+        job_config.training.mixed_precision_param,
+        device_type,
+    )
+
+    # variables used to keep info for metrics logging
+    device_memory_monitor.reset_peak_stats()
+
+    global_batch_size = (
+        job_config.training.batch_size
+        * dp_degree
+        * job_config.training.gradient_accumulation_steps
+    )
+    num_tokens_per_step = global_batch_size * job_config.training.seq_len
+    # train loop
+    logger.info(f"{color.red}***** Running training *****{color.reset}")
+    logger.info(f"{color.green}  Training starts at step {train_state.step + 1}")
+    logger.info(
+        f"{color.green}  Number of tokens per sequence = {job_config.training.seq_len:,}"
+    )
+    logger.info(
+        f"{color.green}  Gradient Accumulation steps = {job_config.training.gradient_accumulation_steps}"
+    )
+    logger.info(
+        f"{color.green}  Instantaneous batch size (per device) = {job_config.training.batch_size:,}"
+    )
+    logger.info(
+        f"{color.green}  Global batch size (w. parallel, distributed & accumulation) = {global_batch_size:,}"
+        f" ({num_tokens_per_step:,} tokens)"
+    )
+    logger.info(
+        f"{color.green}  Total optimization steps = {job_config.training.steps:,} "
+        f"({job_config.training.steps * num_tokens_per_step:,} tokens)"
+    )
+    logger.info(
+        f"{color.green}  Warmup steps = {job_config.lr_scheduler.warmup_steps:,}"
+        f" ({job_config.lr_scheduler.warmup_steps * num_tokens_per_step:,} tokens)"
+    )
+    logger.info(
+        f"{color.green}  Number of parameters = {model_param_count:,} {color.reset}"
+    )
+
+    with (
+        maybe_enable_profiling(
+            job_config, global_step=train_state.step
+        ) as torch_profiler,
+        maybe_enable_memory_snapshot(
+            job_config, global_step=train_state.step
+        ) as memory_profiler,
+    ):
+        while train_state.step < job_config.training.steps:
+            train_state.step += 1
+            gc_handler.run(train_state.step)
+
+            optimizers.zero_grad()
+
+            losses = []
+            # do gradient accumulation if enabled
+            for _ in range(job_config.training.gradient_accumulation_steps):
+                # get batch
+                data_load_start = time.perf_counter()
+                batch = next(data_iterator)
+                input_ids, labels = batch["input_ids"], batch["labels"]
+
+                # Update metrics processor state before forward/backward
+                metric_logger.ntokens_since_last_log += labels.numel()
+                metric_logger.data_loading_times.append(
+                    time.perf_counter() - data_load_start
+                )
+
+                input_ids = input_ids.to(device_type)
+
+                """
+                TODO[flame]: We need to carefully handle the position_ids for TP/CP
+                Depending on the Models'PE, the position_ids might be different.
+
+                e.g. for TP
+                    For RoPE, all ranks have the same position_ids. [FOR HF model]
+                    For sinusoidal, each rank has the coresponding chunked  position_ids. [FOR HF model]
+
+                e.g. for CP, [optional_context_parallel_ctx shoudl automatically distbute the position_ids]
+                    Each rank has the coresponding chunked position_ids. [FOR All model]
+
+                """
+                labels = labels.to(device_type)
+                cu_seqlens = (
+                    batch["cu_seqlens"].to(device_type)
+                    if "cu_seqlens" in batch
+                    else None
+                )
+                if cu_seqlens is not None:
+                    position_ids = prepare_position_ids(cu_seqlens).to(torch.int32)
+                else:
+                    position_ids = (
+                        torch.arange(0, input_ids.shape[1], device=device_type)
+                        .repeat(input_ids.shape[0], 1)
+                        .to(torch.int32)
+                    )
+                # apply context parallelism if cp is enabled
+                # ensure CP handles the separate freqs_cis buffer for each pp stage
+                optional_context_parallel_ctx = (
+                    dist_utils.create_context_parallel_ctx(
+                        cp_mesh=world_mesh["cp"],
+                        cp_buffers=[input_ids, labels, position_ids],
+                        cp_seq_dims=[1, 1, 1],
+                        cp_no_restore_buffers={input_ids, labels, position_ids},
+                        cp_rotate_method=job_config.experimental.context_parallel_rotate_method,
+                    )
+                    if parallel_dims.cp_enabled
+                    else None
+                )
+
+                # #! TODO[flame], we should distribute the position_ids as well with CP
+                if parallel_dims.pp_enabled:
+                    raise NotImplementedError(
+                        "Pipeline parallelism is not supported in this version"
+                    )
+                    # Pipeline Parallel forward / backward inside step() call
+                    with train_context(optional_context_parallel_ctx):
+                        targets, losses = (
+                            (labels, []) if has_last_stage else (None, None)
+                        )
+
+                        if has_first_stage:
+                            pp_schedule.step(input_ids, target=targets, losses=losses)
+                        else:
+                            pp_schedule.step(target=targets, losses=losses)
+
+                    # accumulate losses across pipeline microbatches
+                    # TODO: PP+FSDP unexpectedly puts the loss back to the CPU
+                    loss = (
+                        torch.mean(torch.stack(losses)).to(device)
+                        if has_last_stage
+                        else torch.tensor([-1.0], device=device)
+                    )
+                else:
+                    # Non-PP forward / backward
+                    with train_context(optional_context_parallel_ctx):
+                        with maybe_enable_amp:
+                            output = model(
+                                input_ids=input_ids,
+                                labels=labels,
+                                position_ids=position_ids,
+                                cu_seqlens=cu_seqlens,
+                        )
+                        loss = (
+                            output.loss
+                            / job_config.training.gradient_accumulation_steps
+                        )
+                        loss.backward()
+
+                losses.append(loss)
+            loss = sum(losses)
+
+            # clip gradients
+            grad_norm = dist_utils.clip_grad_norm_(
+                [p for m in model_parts for p in m.parameters()],
+                job_config.training.max_norm,
+                foreach=True,
+                pp_mesh=pp_mesh if parallel_dims.pp_enabled else None,
+            )
+
+            # optimizer step
+            checkpoint.maybe_wait_for_staging()
+            if job_config.training.skip_nan_inf and (
+                grad_norm.isnan() or grad_norm.isinf()
+            ):
+                logger.warning(
+                    f"Skipping optimizer step - detected invalid gradient norm: {grad_norm:.4f}"
+                )
+                optimizers.zero_grad()
+                train_state.skipped_step += 1
+            else:
+                optimizers.step()
+            lr_schedulers.step()
+
+            # log metrics - Use MetricsProcessor
+            if metric_logger.should_log(train_state.step):
+                if (
+                    parallel_dims.dp_replicate_enabled
+                    or parallel_dims.dp_shard_enabled
+                    or parallel_dims.cp_enabled
+                ):
+                    loss = loss.detach()
+                    # Use dist_mean/max on the accumulated loss for the step
+                    global_avg_loss, global_max_loss = (
+                        dist_utils.dist_mean(
+                            loss,
+                            world_mesh["dp_cp"],
+                        ),
+                        dist_utils.dist_max(
+                            loss,
+                            world_mesh["dp_cp"],
+                        ),
+                    )
+                else:
+                    # Scale back the loss before logging
+                    global_avg_loss = global_max_loss = loss.item()
+
+                # Update train state tokens and elapsed time
+                time_now = time.perf_counter()
+                time_delta = (
+                    time_now - metric_logger.time_last_log
+                )  # Use metric_logger's time
+                train_state.token += (
+                    metric_logger.ntokens_since_last_log  # Use tokens tracked by metric_logger
+                    * parallel_dims.world_size
+                    / parallel_dims.non_data_parallel_size
+                )
+                train_state.elapsed += timedelta(seconds=time_delta)
+                train_state.log_steps.append(train_state.step)
+                train_state.global_avg_losses.append(global_avg_loss)
+                train_state.global_max_losses.append(global_max_loss)
+
+                # Log using the metric processor
+                last_lr = lr_schedulers.schedulers[0].get_last_lr()[0]
+                eta = (
+                    train_state.elapsed
+                    * (job_config.training.steps - train_state.step)
+                    / train_state.step
+                )
+                metric_logger.log(
+                    train_state.step,
+                    global_avg_loss,
+                    global_max_loss,
+                    extra_metrics={
+                        "optimizer/lr": last_lr,
+                        "optimizer/grad_norm": grad_norm.item(),
+                        "optimizer/skipped_step": train_state.skipped_step,
+                    },
+                )
+
+                logger.info(
+                    f"{color.blue}lr: {last_lr:.4e} gnorm: {grad_norm:5.2f} "
+                    f"{color.magenta}[{str(train_state.elapsed).split('.')[0]:>8}<{str(eta).split('.')[0]:>8}]{color.reset}"
+                )
+
+            checkpoint.save(
+                train_state.step, force=(train_state.step == job_config.training.steps)
+            )
+
+            # signal the profiler that the next profiling step has started
+            if torch_profiler:
+                torch_profiler.step()
+            if memory_profiler:
+                memory_profiler.step()
+
+            # reduce timeout after first train step for faster signal
+            # (assuming lazy init and compilation are finished)
+            if train_state.step == 1:
+                dist_utils.set_pg_timeouts(
+                    timeout=timedelta(seconds=job_config.comm.train_timeout_seconds),
+                    world_mesh=world_mesh,
+                )
+
+    if torch.distributed.get_rank() == 0:
+        logger.info("Sleeping 2 seconds for other ranks to complete")
+        time.sleep(2)
+
+    metric_logger.close()
+    logger.info("Training completed")
+
+
+if __name__ == "__main__":
+    init_logger()
+    config = JobConfig()
+    config.parse_args()
+    main(config)
+    torch.distributed.destroy_process_group()

flame/train_restart.py

@@ -0,0 +1,694 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import os
+import time
+from datetime import timedelta
+
+import fla  # noqa
+import torch
+from fla.modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss
+from fla.ops.utils import prepare_position_ids
+from torch.distributed.elastic.multiprocessing.errors import record
+from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.ft import FTParallelDims, init_ft_manager
+from torchtitan.components.loss import build_cross_entropy_loss
+from torchtitan.components.lr_scheduler import build_lr_schedulers
+from torchtitan.components.metrics import build_device_memory_monitor, build_metrics_processor, ensure_pp_loss_visible
+from torchtitan.components.optimizer import build_optimizers
+from torchtitan.distributed import ParallelDims
+from torchtitan.distributed import utils as dist_utils
+from torchtitan.protocols.model_converter import build_model_converters
+from torchtitan.protocols.train_spec import TrainSpec, get_train_spec, register_train_spec
+from torchtitan.tools import utils
+from torchtitan.tools.logging import init_logger, logger
+from torchtitan.tools.profiling import maybe_enable_memory_snapshot, maybe_enable_profiling
+
+import custom_models
+from flame.components.checkpoint import TrainState
+from flame.config_manager import JobConfig
+from flame.data import build_dataloader, build_dataset
+from flame.models.parallelize_fla import parallelize_fla
+from flame.models.pipeline_fla import pipeline_fla
+from flame.tools.utils import get_nparams_and_flops
+
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from fla.models import HamiltonForCausalLM as NewModelForCausalLM, HamiltonConfig as NewConfig
+# from fla.models import GLAForCausalLM as NewModelForCausalLM, GLAConfig as NewConfig
+MODEL_TYPE = NewConfig.model_type
+
+def build_tokenizer(job_config: JobConfig) -> AutoTokenizer:
+    return AutoTokenizer.from_pretrained(job_config.model.tokenizer_path)
+
+
+register_train_spec(
+    TrainSpec(
+        name="fla",
+        cls=AutoModelForCausalLM,
+        config=AutoConfig,
+        parallelize_fn=parallelize_fla,
+        pipelining_fn=pipeline_fla,
+        build_optimizers_fn=build_optimizers,
+        build_lr_schedulers_fn=build_lr_schedulers,
+        build_dataloader_fn=build_dataloader,
+        build_tokenizer_fn=build_tokenizer,
+        build_loss_fn=build_cross_entropy_loss,
+    )
+)
+
+
+# Enable debug tracing on failure: https://pytorch.org/docs/stable/elastic/errors.html
+@record
+def main(job_config: JobConfig):
+    logger.info(f"Starting job: {job_config.job.description}")
+    logger.info(f"Registering model type: {MODEL_TYPE}")
+
+    if job_config.experimental.custom_model_path:
+        utils.import_module_from_path(job_config.experimental.custom_model_path)
+
+    # used for colorful printing
+    color = utils.NoColor if job_config.metrics.disable_color_printing else utils.Color
+
+    if job_config.job.print_args:
+        logger.info(
+            f"{color.green}{json.dumps(job_config.to_dict(), indent=2, sort_keys=True)}{color.reset}"
+        )
+
+    # take control of garbage collection to avoid stragglers
+    gc_handler = utils.GarbageCollection(gc_freq=job_config.training.gc_freq)
+
+    device_module, device_type = utils.device_module, utils.device_type
+    device = torch.device(f"{device_type}:{int(os.environ['LOCAL_RANK'])}")
+    # Device has to be set before creating TorchFT manager.
+    device_module.set_device(device)
+    ft_manager = init_ft_manager(job_config)
+
+    # init distributed
+    world_size = int(os.environ["WORLD_SIZE"])
+    if not ft_manager.enabled:
+        parallel_dims = ParallelDims(
+            dp_shard=job_config.training.data_parallel_shard_degree,
+            dp_replicate=job_config.training.data_parallel_replicate_degree,
+            cp=job_config.experimental.context_parallel_degree,
+            tp=job_config.training.tensor_parallel_degree,
+            pp=job_config.experimental.pipeline_parallel_degree,
+            world_size=world_size,
+            enable_loss_parallel=not job_config.training.disable_loss_parallel,
+        )
+    else:
+        parallel_dims = FTParallelDims(
+            dp_shard=job_config.training.data_parallel_shard_degree,
+            dp_replicate=job_config.training.data_parallel_replicate_degree,
+            cp=job_config.experimental.context_parallel_degree,
+            tp=job_config.training.tensor_parallel_degree,
+            pp=job_config.experimental.pipeline_parallel_degree,
+            world_size=world_size,
+            enable_loss_parallel=not job_config.training.disable_loss_parallel,
+            ft_manager=ft_manager,
+        )
+    dist_utils.init_distributed(job_config)
+    # initialize device memory monitor and get peak flops for MFU calculation
+    device_memory_monitor = build_device_memory_monitor()
+    gpu_peak_flops = utils.get_peak_flops(device_memory_monitor.device_name)
+    logger.info(f"Peak FLOPS used for computing MFU: {gpu_peak_flops:.3e}")
+
+    # build meshes
+    world_mesh = parallel_dims.build_mesh(device_type=device_type)
+    if parallel_dims.dp_enabled:
+        dp_mesh = world_mesh["dp"]
+        dp_degree, dp_rank = dp_mesh.size(), dp_mesh.get_local_rank()
+    else:
+        dp_degree, dp_rank = 1, 0
+
+    if parallel_dims.pp_enabled:
+        raise NotImplementedError(
+            "Pipeline parallelism is not supported in this version"
+        )
+        """
+        ! TODO[flame]: We need to fix the pipeline parallelism for flame
+        [x] Match the key of models' components with the actual naming
+        [ ] Fix the post-init and tie-embedding for pipeline parallelism, HF's transformer automatically
+            forces to tie if head is None, we need to handle this case
+        [ ]
+        """
+        pp_mesh = world_mesh["pp"]
+
+    # Set random seed, and maybe enable deterministic mode (mainly for debugging, expect perf loss)
+    dist_utils.set_determinism(
+        world_mesh, device, job_config.training.seed, job_config.training.deterministic
+    )
+    train_spec = get_train_spec(job_config.model.name)
+
+    logger.info("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        job_config.model.tokenizer_path,
+        trust_remote_code=True,
+        model_max_length=int(1e10),
+    )
+    logger.info(f"{tokenizer}")
+    logger.info(
+        f"Loading dataset {job_config.training.dataset}"
+        f":{job_config.training.dataset_name}"
+        if job_config.training.dataset_name is not None
+        else ""
+    )
+    dataset = build_dataset(
+        dataset=job_config.training.dataset,
+        dataset_name=job_config.training.dataset_name,
+        dataset_split=job_config.training.dataset_split,
+        data_dir=job_config.training.data_dir,
+        data_files=job_config.training.data_files,
+        data_probs=job_config.training.data_probs,
+        streaming=job_config.training.streaming,
+        dp_degree=dp_degree,
+        num_workers=job_config.training.num_workers,
+        seed=job_config.training.seed,
+    )
+
+    logger.info("Building dataloader...")
+    dataloader = build_dataloader(
+        dataset=dataset,
+        tokenizer=tokenizer,
+        rank=dp_rank,
+        world_size=dp_degree,
+        batch_size=job_config.training.batch_size,
+        seq_len=job_config.training.seq_len,
+        context_len=job_config.training.context_len,
+        varlen=job_config.training.varlen,
+        num_workers=job_config.training.num_workers,
+        pin_memory=job_config.training.pin_memory,
+        persistent_workers=job_config.training.persistent_workers,
+        snapshot_every_n_steps=job_config.checkpoint.interval,
+    )
+
+
+    logger.info(f"Loading model config from {job_config.model.config}")
+    logger.info(f"Registering model type: {MODEL_TYPE}")
+    AutoConfig.register(MODEL_TYPE, NewConfig)   # important!
+    AutoModelForCausalLM.register(NewConfig, NewModelForCausalLM)  # important!
+    model_config = AutoConfig.from_pretrained(job_config.model.config)
+    # set the model configs from training inputs:
+    # 1. norm type to decide which norm layer to use
+    # 2. disable fused norm if TP is enabled
+    # 3. vocab size from tokenizer
+    # 4. context_len base on inputs
+    if parallel_dims.tp_enabled:
+        if model_config.fuse_norm:
+            logger.warning(
+                f"{color.red}"
+                f"Fused norm is not compatible with tensor parallelism. "
+                f"Disabling it for now."
+                f"{color.reset}"
+            )
+            model_config.fuse_norm = False
+    if parallel_dims.loss_parallel_enabled:
+        if model_config.fuse_linear_cross_entropy:
+            logger.warning(
+                f"{color.red}"
+                f"Loss parallel enabled. Disabling fused cross entropy for now."
+                f"{color.reset}"
+            )
+            model_config.fuse_linear_cross_entropy = False
+    model_config.vocab_size = max(tokenizer.vocab_size, model_config.vocab_size)
+
+
+
+    logger.info(
+        f"Building model from the config\n{color.green}{model_config}{color.reset}"
+    )
+    with torch.device("meta"):
+        model = AutoModelForCausalLM.from_config(model_config)
+        if (
+            getattr(model_config, "fuse_linear_cross_entropy", False)
+            and FusedLinearCrossEntropyLoss is not None
+        ):
+            model.criterion = FusedLinearCrossEntropyLoss(
+                num_chunks=8 // parallel_dims.tp
+            )
+        # defer weight initialization until after parallelisms are applied
+        model.apply(lambda m: setattr(m, "_is_hf_initialized", False))
+    logger.info(f"{color.blue}\n{model}{color.reset}\n")
+
+    logger.info("Applying model converters...")
+
+    # Build the collection of model converters. No-op if `model.converters` empty
+    model_converters = build_model_converters(job_config, parallel_dims)
+    model_converters.convert(model)
+
+    # calculate model size and flops per token
+    model_param_count, num_flops_per_token = get_nparams_and_flops(
+        model, model_config, job_config.training.context_len
+    )
+
+    # move sharded model to CPU/GPU and initialize weights via DTensor
+    if job_config.checkpoint.create_seed_checkpoint:
+        init_device = "cpu"
+    elif job_config.training.enable_cpu_offload:
+        init_device = "cpu"
+    else:
+        init_device = device_type
+
+    # apply parallelisms and initialization
+    if parallel_dims.pp_enabled:
+        # apply PT-D Pipeline Parallel
+        (
+            pp_schedule,
+            model_parts,
+            has_first_stage,
+            has_last_stage,
+        ) = train_spec.pipelining_fn(
+            model,
+            pp_mesh,
+            parallel_dims,
+            job_config,
+            device,
+            model_config,
+            train_spec.loss_fn,
+        )
+        # when PP is enabled, `model` obj is no longer used after this point, model_parts is used instead
+        del model
+
+        # For PP with looped schedules, each item in model_parts is one stage-model-chunk.
+        # We need to iterate through model_parts to apply SPMD parallelisms, compilation,
+        # optimizer, and checkpointing
+        for m in model_parts:
+            # apply SPMD-style PT-D techniques
+            train_spec.parallelize_fn(m, world_mesh, parallel_dims, job_config)
+            m.to_empty(device=init_device)
+            with torch.no_grad():
+                m.post_init()
+            m.train()
+
+        # confirm that user will be able to view loss metrics on the console
+        ensure_pp_loss_visible(parallel_dims, job_config, color)
+    else:
+        # apply PT-D Tensor Parallel, activation checkpointing, torch.compile, Data Parallel
+        train_spec.parallelize_fn(model, world_mesh, parallel_dims, job_config)
+        model.to_empty(device=init_device)
+        with torch.no_grad():
+            model.post_init()
+        model.train()
+
+        model_parts = [model]
+
+    device_mem_stats = device_memory_monitor.get_peak_stats()
+    logger.info(
+        f"{device_type.upper()} memory usage for model: "
+        f"{device_mem_stats.max_reserved_gib:.2f}GiB"
+        f"({device_mem_stats.max_reserved_pct:.2f}%)"
+    )
+
+    # build optimizer after applying parallelisms to the model
+    optimizers = train_spec.build_optimizers_fn(model_parts, job_config, ft_manager)
+    lr_schedulers = train_spec.build_lr_schedulers_fn(optimizers, job_config)
+    # Post optimizer step model converters hook.
+    # e.g. calculate float8 dynamic amax/scale for all-parameter for FSDP2
+    # where it issues a single all-reduce for all parameters at once for better performance
+    optimizers.register_step_post_hook(
+        lambda *args, **kwargs: model_converters.post_optimizer_hook(model_parts)
+    )
+
+    train_state = TrainState()
+
+    # load initial checkpoint
+    checkpoint = CheckpointManager(
+        dataloader=dataloader,
+        model_parts=model_parts,
+        optimizers=optimizers,
+        lr_schedulers=lr_schedulers,
+        states={"train_state": train_state},
+        job_config=job_config,
+        ft_manager=ft_manager,
+    )
+
+    # if job_config.training.streaming and job_config.checkpoint.enable_checkpoint:
+    #     checkpoint = CheckpointManager(
+    #         dataloader=None,  # 使用变量
+    #         model_parts=model_parts,
+    #         optimizers=optimizers,
+    #         lr_schedulers=lr_schedulers,
+    #         states={"train_state": train_state},
+    #         job_config=job_config,
+    #         ft_manager=ft_manager,
+    #     )
+    #     if hasattr(checkpoint, 'states') and 'dataloader' in checkpoint.states:
+    #         print("[Fix] Manually removing 'dataloader' from checkpoint states to avoid Missing Key error.")
+    #         del checkpoint.states['dataloader']
+    # else:
+    #     checkpoint = CheckpointManager(
+    #     dataloader=dataloader,
+    #     model_parts=model_parts,
+    #     optimizers=optimizers,
+    #     lr_schedulers=lr_schedulers,
+    #     states={"train_state": train_state},
+    #     job_config=job_config,
+    #     ft_manager=ft_manager,
+    # )
+
+
+    if job_config.checkpoint.create_seed_checkpoint:
+        assert world_size == 1, (
+            "Must create seed checkpoint using a single device, to disable sharding"
+        )
+        assert job_config.checkpoint.enable_checkpoint, (
+            "Must enable checkpointing when creating a seed checkpoint"
+        )
+        checkpoint.save(curr_step=0, force=True)
+        logger.info("Created seed checkpoint")
+        return
+
+    logger.info(job_config.checkpoint)
+    checkpoint.load(step=job_config.checkpoint.load_step)
+    metric_logger = build_metrics_processor(job_config, parallel_dims)
+
+    # Set dependent attributes for metric_logger
+    metric_logger.num_flops_per_token = num_flops_per_token
+    metric_logger.optimizers = optimizers  # Pass optimizers if needed by logger logic
+    metric_logger.lr_schedulers = (
+        lr_schedulers  # Pass schedulers if needed by logger logic
+    )
+
+    # plot losses loaded from checkpoint (if any) to TensorBoard
+    # NOTE: Loss info after the last log step before checkpoint saving will not be ploted.
+    #       This can be avoided by setting checkpoint.interval to be a multiple of metrics.log_freq
+    if train_state.step > 0 and len(metric_logger.data_loading_times) > 0:
+        for idx, step in enumerate(train_state.log_steps):
+            metric_logger.log(
+                step,
+                global_avg_loss=train_state.global_avg_losses[idx],
+                global_max_loss=train_state.global_max_losses[idx],
+            )
+
+    data_iterator = iter(dataloader)
+
+    # if job_config.training.streaming and train_state.step > 0:
+    #     from tqdm import tqdm
+    #     import gc
+    #     skip_count = 14
+    #     # skip_count = train_state.step * job_config.training.gradient_accumulation_steps
+    #     local_rank = torch.distributed.get_rank()
+    #     if local_rank == 0:
+    #         logger.info(f"Streaming Resume: Skipping {skip_count} micro-batches to catch up...")
+    #     for i in tqdm(range(skip_count), 
+    #         desc="Skipping Data", 
+    #         unit="batch", 
+    #         disable=(local_rank != 0),
+    #         dynamic_ncols=True):  
+    #         try:
+    #             batch = next(data_iterator)
+    #             del batch 
+    #         except StopIteration:
+    #             if local_rank == 0:
+    #                 logger.warning(f"Data iterator exhausted before finishing skip at step {_}!")
+    #             break
+
+    #         if i % 500 == 0:
+    #             gc.collect()
+    #     gc.collect()
+
+    #     if local_rank == 0:
+    #         logger.info("Data skipping completed. Resuming training...")
+    # =================================================================
+
+
+    
+
+    train_context = dist_utils.get_train_context(
+        parallel_dims.loss_parallel_enabled,
+        job_config.experimental.enable_compiled_autograd,
+    )
+    maybe_enable_amp = dist_utils.maybe_enable_amp(
+        parallel_dims,
+        job_config.training.mixed_precision_param,
+        device_type,
+    )
+
+    # variables used to keep info for metrics logging
+    device_memory_monitor.reset_peak_stats()
+
+    global_batch_size = (
+        job_config.training.batch_size
+        * dp_degree
+        * job_config.training.gradient_accumulation_steps
+    )
+    num_tokens_per_step = global_batch_size * job_config.training.seq_len
+    # train loop
+    logger.info(f"{color.red}***** Running training *****{color.reset}")
+    logger.info(f"{color.green}  Training starts at step {train_state.step + 1}")
+    logger.info(
+        f"{color.green}  Number of tokens per sequence = {job_config.training.seq_len:,}"
+    )
+    logger.info(
+        f"{color.green}  Gradient Accumulation steps = {job_config.training.gradient_accumulation_steps}"
+    )
+    logger.info(
+        f"{color.green}  Instantaneous batch size (per device) = {job_config.training.batch_size:,}"
+    )
+    logger.info(
+        f"{color.green}  Global batch size (w. parallel, distributed & accumulation) = {global_batch_size:,}"
+        f" ({num_tokens_per_step:,} tokens)"
+    )
+    logger.info(
+        f"{color.green}  Total optimization steps = {job_config.training.steps:,} "
+        f"({job_config.training.steps * num_tokens_per_step:,} tokens)"
+    )
+    logger.info(
+        f"{color.green}  Warmup steps = {job_config.lr_scheduler.warmup_steps:,}"
+        f" ({job_config.lr_scheduler.warmup_steps * num_tokens_per_step:,} tokens)"
+    )
+    logger.info(
+        f"{color.green}  Number of parameters = {model_param_count:,} {color.reset}"
+    )
+
+    with (
+        maybe_enable_profiling(
+            job_config, global_step=train_state.step
+        ) as torch_profiler,
+        maybe_enable_memory_snapshot(
+            job_config, global_step=train_state.step
+        ) as memory_profiler,
+    ):
+        while train_state.step < job_config.training.steps:
+            train_state.step += 1
+            gc_handler.run(train_state.step)
+
+            optimizers.zero_grad()
+
+            losses = []
+            # do gradient accumulation if enabled
+            for _ in range(job_config.training.gradient_accumulation_steps):
+                # get batch
+                data_load_start = time.perf_counter()
+                batch = next(data_iterator)
+                input_ids, labels = batch["input_ids"], batch["labels"]
+
+                # Update metrics processor state before forward/backward
+                metric_logger.ntokens_since_last_log += labels.numel()
+                metric_logger.data_loading_times.append(
+                    time.perf_counter() - data_load_start
+                )
+
+                input_ids = input_ids.to(device_type)
+
+                """
+                TODO[flame]: We need to carefully handle the position_ids for TP/CP
+                Depending on the Models'PE, the position_ids might be different.
+
+                e.g. for TP
+                    For RoPE, all ranks have the same position_ids. [FOR HF model]
+                    For sinusoidal, each rank has the coresponding chunked  position_ids. [FOR HF model]
+
+                e.g. for CP, [optional_context_parallel_ctx shoudl automatically distbute the position_ids]
+                    Each rank has the coresponding chunked position_ids. [FOR All model]
+
+                """
+                labels = labels.to(device_type)
+                cu_seqlens = (
+                    batch["cu_seqlens"].to(device_type)
+                    if "cu_seqlens" in batch
+                    else None
+                )
+                if cu_seqlens is not None:
+                    position_ids = prepare_position_ids(cu_seqlens).to(torch.int32)
+                else:
+                    position_ids = (
+                        torch.arange(0, input_ids.shape[1], device=device_type)
+                        .repeat(input_ids.shape[0], 1)
+                        .to(torch.int32)
+                    )
+                # apply context parallelism if cp is enabled
+                # ensure CP handles the separate freqs_cis buffer for each pp stage
+                optional_context_parallel_ctx = (
+                    dist_utils.create_context_parallel_ctx(
+                        cp_mesh=world_mesh["cp"],
+                        cp_buffers=[input_ids, labels, position_ids],
+                        cp_seq_dims=[1, 1, 1],
+                        cp_no_restore_buffers={input_ids, labels, position_ids},
+                        cp_rotate_method=job_config.experimental.context_parallel_rotate_method,
+                    )
+                    if parallel_dims.cp_enabled
+                    else None
+                )
+
+                # #! TODO[flame], we should distribute the position_ids as well with CP
+                if parallel_dims.pp_enabled:
+                    raise NotImplementedError(
+                        "Pipeline parallelism is not supported in this version"
+                    )
+                    # Pipeline Parallel forward / backward inside step() call
+                    with train_context(optional_context_parallel_ctx):
+                        targets, losses = (
+                            (labels, []) if has_last_stage else (None, None)
+                        )
+
+                        if has_first_stage:
+                            pp_schedule.step(input_ids, target=targets, losses=losses)
+                        else:
+                            pp_schedule.step(target=targets, losses=losses)
+
+                    # accumulate losses across pipeline microbatches
+                    # TODO: PP+FSDP unexpectedly puts the loss back to the CPU
+                    loss = (
+                        torch.mean(torch.stack(losses)).to(device)
+                        if has_last_stage
+                        else torch.tensor([-1.0], device=device)
+                    )
+                else:
+                    # Non-PP forward / backward
+                    with train_context(optional_context_parallel_ctx):
+                        with maybe_enable_amp:
+                            output = model(
+                                input_ids=input_ids,
+                                labels=labels,
+                                position_ids=position_ids,
+                                cu_seqlens=cu_seqlens,
+                        )
+                        loss = (
+                            output.loss
+                            / job_config.training.gradient_accumulation_steps
+                        )
+                        loss.backward()
+                        # print('--------------------------')
+
+                losses.append(loss)
+            loss = sum(losses)
+
+            # clip gradients
+            grad_norm = dist_utils.clip_grad_norm_(
+                [p for m in model_parts for p in m.parameters()],
+                job_config.training.max_norm,
+                foreach=True,
+                pp_mesh=pp_mesh if parallel_dims.pp_enabled else None,
+            )
+
+            # optimizer step
+            checkpoint.maybe_wait_for_staging()
+            if job_config.training.skip_nan_inf and (
+                grad_norm.isnan() or grad_norm.isinf()
+            ):
+                logger.warning(
+                    f"Skipping optimizer step - detected invalid gradient norm: {grad_norm:.4f}"
+                )
+                optimizers.zero_grad()
+                train_state.skipped_step += 1
+            else:
+                optimizers.step()
+            lr_schedulers.step()
+
+            # log metrics - Use MetricsProcessor
+            if metric_logger.should_log(train_state.step):
+                if (
+                    parallel_dims.dp_replicate_enabled
+                    or parallel_dims.dp_shard_enabled
+                    or parallel_dims.cp_enabled
+                ):
+                    loss = loss.detach()
+                    # Use dist_mean/max on the accumulated loss for the step
+                    global_avg_loss, global_max_loss = (
+                        dist_utils.dist_mean(
+                            loss,
+                            world_mesh["dp_cp"],
+                        ),
+                        dist_utils.dist_max(
+                            loss,
+                            world_mesh["dp_cp"],
+                        ),
+                    )
+                else:
+                    # Scale back the loss before logging
+                    global_avg_loss = global_max_loss = loss.item()
+
+                # Update train state tokens and elapsed time
+                time_now = time.perf_counter()
+                time_delta = (
+                    time_now - metric_logger.time_last_log
+                )  # Use metric_logger's time
+                train_state.token += (
+                    metric_logger.ntokens_since_last_log  # Use tokens tracked by metric_logger
+                    * parallel_dims.world_size
+                    / parallel_dims.non_data_parallel_size
+                )
+                train_state.elapsed += timedelta(seconds=time_delta)
+                train_state.log_steps.append(train_state.step)
+                train_state.global_avg_losses.append(global_avg_loss)
+                train_state.global_max_losses.append(global_max_loss)
+
+                # Log using the metric processor
+                last_lr = lr_schedulers.schedulers[0].get_last_lr()[0]
+                eta = (
+                    train_state.elapsed
+                    * (job_config.training.steps - train_state.step)
+                    / train_state.step
+                )
+                metric_logger.log(
+                    train_state.step,
+                    global_avg_loss,
+                    global_max_loss,
+                    extra_metrics={
+                        "optimizer/lr": last_lr,
+                        "optimizer/grad_norm": grad_norm.item(),
+                        "optimizer/skipped_step": train_state.skipped_step,
+                    },
+                )
+
+                logger.info(
+                    f"{color.blue}lr: {last_lr:.4e} gnorm: {grad_norm:5.2f} "
+                    f"{color.magenta}[{str(train_state.elapsed).split('.')[0]:>8}<{str(eta).split('.')[0]:>8}]{color.reset}"
+                )
+
+            checkpoint.save(
+                train_state.step, force=(train_state.step == job_config.training.steps)
+            )
+
+            # signal the profiler that the next profiling step has started
+            if torch_profiler:
+                torch_profiler.step()
+            if memory_profiler:
+                memory_profiler.step()
+
+            # reduce timeout after first train step for faster signal
+            # (assuming lazy init and compilation are finished)
+            if train_state.step == 1:
+                dist_utils.set_pg_timeouts(
+                    timeout=timedelta(seconds=job_config.comm.train_timeout_seconds),
+                    world_mesh=world_mesh,
+                )
+
+    if torch.distributed.get_rank() == 0:
+        logger.info("Sleeping 2 seconds for other ranks to complete")
+        time.sleep(2)
+
+    metric_logger.close()
+    logger.info("Training completed")
+
+
+if __name__ == "__main__":
+    init_logger()
+    config = JobConfig()
+    config.parse_args()
+    main(config)
+    torch.distributed.destroy_process_group()

flame/utils/convert_dcp_to_hf.py

@@ -0,0 +1,74 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+import argparse
+import io
+import os
+import tempfile
+from datetime import timedelta
+
+import fla  # noqa
+import torch
+import torch.serialization
+from torch.distributed.checkpoint.format_utils import dcp_to_torch_save
+from torchtitan.tools.logging import init_logger, logger
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+from fla.models import HamiltonForCausalLM as NewModelForCausalLM, HamiltonConfig as NewConfig
+MODEL_TYPE = NewConfig.model_type
+
+# import custom_models
+
+
+@torch.inference_mode()
+def save_pretrained(
+    path: str,
+    step: int,
+    config: str,
+    tokenizer: str
+):
+    logger.info(f"Loading the config from {config}")
+    
+
+    AutoConfig.register(MODEL_TYPE, NewConfig)   # important!
+    config = AutoConfig.from_pretrained(config, trust_remote_code=True)
+
+    logger.info(f"Saving the config to {path}")
+    config.save_pretrained(path)
+    logger.info(f"Loading the tokenizer from {tokenizer}")
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer, trust_remote_code=True)
+    logger.info(f"Saving the tokenizer to {path}")
+    tokenizer.save_pretrained(path)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        checkpoint = os.path.join(path, f'checkpoint/step-{step}')
+        checkpoint_path = os.path.join(tmpdir, 'checkpoint.pt')
+        logger.info(f"Saving the distributed checkpoint to {checkpoint_path}")
+        dcp_to_torch_save(checkpoint, checkpoint_path)
+
+        logger.info(f"Initializing the model from config\n{config}")
+        # model = AutoModelForCausalLM.from_config(config)
+        AutoModelForCausalLM.register(NewConfig, NewModelForCausalLM)  # important!
+        model = AutoModelForCausalLM.from_config(config)
+
+        logger.info(model)
+        logger.info("Loading state dict from the checkpoint")
+
+        # Add datetime.timedelta and io.BytesIO to safe globals
+        torch.serialization.add_safe_globals([timedelta, io.BytesIO])
+        # torch.load now with default weights_only=True will work
+        model.load_state_dict(torch.load(checkpoint_path, map_location='cpu')['model'])
+
+        logger.info(f"Saving the model to {path}")
+        model.save_pretrained(path)
+
+
+if __name__ == "__main__":
+    init_logger()
+    parser = argparse.ArgumentParser("Convert DCP format model weights to huggingface-style.")
+    parser.add_argument("--path", type=str, required=True)
+    parser.add_argument("--step", type=int, required=True)
+    parser.add_argument("--config", type=str, required=True)
+    parser.add_argument("--tokenizer", type=str, required=True)
+    args = parser.parse_args()
+    save_pretrained(args.path, args.step, args.config, args.tokenizer)

flame/utils/convert_hf_to_dcp.py

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+import argparse
+from pathlib import Path
+
+import torch
+import torch.distributed.checkpoint as DCP
+from transformers import AutoModelForCausalLM
+
+import fla  # noqa
+from torchtitan.tools.logging import init_logger, logger
+
+
+@torch.inference_mode()
+def convert_hf_weights(model: str, checkpoint: str):
+    logger.info(f"Loading model from {model}")
+    model = AutoModelForCausalLM.from_pretrained(model)
+    state_dict = model.state_dict()
+
+    logger.info(f"Writing to DCP at '{checkpoint}'")
+    checkpoint.mkdir(parents=True, exist_ok=True)
+    storage_writer = DCP.filesystem.FileSystemWriter(checkpoint, thread_count=8)
+    DCP.save({"model": state_dict}, storage_writer=storage_writer)
+
+
+if __name__ == "__main__":
+    init_logger()
+    parser = argparse.ArgumentParser(description="Convert huggingface-style model weights to DCP format.")
+    parser.add_argument("--model", type=str, required=True)
+    parser.add_argument("--checkpoint", type=Path, required=True)
+    args = parser.parse_args()
+
+    convert_hf_weights(args.model, args.checkpoint)

flame/utils/preprocess.py

@@ -0,0 +1,122 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+import argparse
+from typing import Any, Dict, List
+
+from transformers import AutoTokenizer, PreTrainedTokenizer
+
+from flame.data import build_dataset
+from torchtitan.tools.logging import init_logger, logger
+
+
+def tokenize(
+    examples: Dict[str, List[Any]],
+    tokenizer: PreTrainedTokenizer,
+) -> Dict:
+    if 'text' in examples:
+        samples = examples['text']
+    elif 'content' in examples:
+        samples = examples['content']
+    else:
+        raise ValueError(f'No "text" or "content" field found in examples:\n{examples}')
+    input_ids = tokenizer(samples)['input_ids']
+    bits_per_token = [len(sample.encode(encoding='utf-8')) * 8 / len(input_ids[i]) for i, sample in enumerate(samples)]
+    return {'input_ids': input_ids, 'bits_per_token': bits_per_token}
+
+
+if __name__ == '__main__':
+    init_logger()
+    parser = argparse.ArgumentParser(description='Preprocess the dataset.')
+    parser.add_argument(
+        '--dataset',
+        default='HuggingFaceFW/fineweb-edu',
+        help='Dataset to use, with comma separated values',
+    )
+    parser.add_argument(
+        '--dataset_name',
+        default='sample-100BT',
+        help='The name of the dataset config, with comma separated values if provided',
+    )
+    parser.add_argument(
+        '--dataset_split',
+        default='train',
+        help='Dataset split to use, with comma separated values if provided',
+    )
+    parser.add_argument(
+        '--data_dir',
+        default=None,
+        help='Data dirs to use, with comma separated values if provided',
+    )
+    parser.add_argument(
+        '--data_files',
+        default=None,
+        help='Data files to use, with comma separated values if provided',
+    )
+    parser.add_argument(
+        '--data_probs',
+        default=None,
+        help='Data sampling probabilities, with comma separated values if provided',
+    )
+    parser.add_argument(
+        '--streaming',
+        action='store_true',
+        help='Whether to use streaming mode',
+    )
+    parser.add_argument(
+        '--num_workers',
+        type=int,
+        default=64,
+        help='Number of workers to use for preprocessing',
+    )
+    parser.add_argument(
+        '--seed',
+        type=int,
+        default=42,
+        help='Random seed for preprocessing',
+    )
+    parser.add_argument(
+        '--path',
+        default='data',
+        help='Path to save the preprocessed dataset',
+    )
+    parser.add_argument(
+        '--tokenizer',
+        default='fla-hub/transformer-1.3B-100B',
+        help='Tokenizer to use',
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=2048,
+        help="Batch size for processing"
+    )
+    args = parser.parse_args()
+
+    logger.info(f'Loading tokenizer {args.tokenizer}')
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
+    logger.info(f'{tokenizer}')
+    logger.info(f'Loading dataset {args.dataset} {args.dataset_name} {args.dataset_split}')
+    dataset = build_dataset(
+        dataset=args.dataset,
+        dataset_name=args.dataset_name,
+        dataset_split=args.dataset_split,
+        data_dir=args.data_dir,
+        data_files=args.data_files,
+        data_probs=args.data_probs,
+        streaming=args.streaming,
+        num_workers=args.num_workers,
+        seed=args.seed,
+    )
+    logger.info(f'Tokenizing and processing the dataset with batch size {args.batch_size}')
+    dataset = dataset.map(
+        lambda examples: tokenize(examples, tokenizer),
+        batched=True,
+        batch_size=args.batch_size,
+        remove_columns=list(next(iter(dataset)).keys()),
+        num_proc=args.num_workers,
+        desc="Running tokenizer on dataset"
+    )
+    logger.info(f'{dataset}')
+    logger.info(f'Saving tokenized dataset to {args.path}')
+    dataset.save_to_disk(args.path)