To download the dataset to your local disk, create a new Python file with the following content and execute it:
py
1from datasets import load_dataset
23# load fineweb-edu with parallel processing4dataset = load_dataset("HuggingFaceFW/fineweb-edu", name="default", num_proc=64, cache_dir="/your/cache/path")56# or load a subset with roughly 100B tokens, suitable for small- or medium-sized experiments7dataset = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-100BT", num_proc=64, cache_dir="/your/cache/path")
Training Recipes
Here's an example of training a 340M FLA Transformer model with a LLaMA-like architecture from scratch on a 100BT subset of the Fineweb-edu corpus in streaming mode. (Do not use streaming mode if you are concerned about resuming training.)
[!WARNING]
If the dataset is not downloaded beforehand, the streaming mode will attempt to fetch it from a remote server and download it on-the-fly, which can be highly unstable during training due to network issues.
For stable training, ensure the dataset is downloaded locally (see Dataset Preparation). Otherwise, we assume you are only testing the new corpus.
You can specify the number of GPUs by setting the environment variable NGPU, which defaults to 8. For single-GPU debugging, set NGPU=1.
We provide several config files for different models.
By default, the learning rate is set to 1e-3 with a cosine scheduler. Other schedulers, such as WSD (wsd), are also supported.
Key parameters:
--lr_scheduler.decay_ratio: The proportion of the steps allocated to the decay phase. The learning rate will remain stable after the warmup period and only start decaying during the last decay_ratio portion of the total training steps, which is known as the Warmup-Stable-Decay (WSD) schedule.
--lr_scheduler.warmup_steps: The number of steps for the learning rate warmup phase.
--training.steps: Total number of training steps.
--training.batch_size: Batch size per device, must be 1 if --training.varlen is set.
--training.seq_len: The length of each sequence in the batch, which is concatenated from multiple samples.
--training.context_len: The max allowed length of a sample. For non-varlen mode, this is equivalent to seq_len.
--training.varlen: Whether to conduct variable-length sequence training.
--training.gradient_accumulation_steps: Number of gradient accumulation steps.
[!WARNING]
The total number of tokens processed per batch, referred to as global_batch_size, is calculated as batch_size × gradient_accumulation_steps × num_gpus.
Each step processes global_batch_size * seq_len tokens.
Monitor the value of global_batch_size, warmup_steps, and steps carefully when modifying any of the hyperparameters!
For a detailed explanation of all parameters, run:
bash train.sh -h
Usage
py
1options:2-h,--help show this help message and exit
3--job.config_file JOB.CONFIG_FILE
4 Job config file5--job.dump_folder JOB.DUMP_FOLDER
6 Folder to dump job outputs
7--job.description JOB.DESCRIPTION
8 Description of the job
9--job.use_for_integration_test
10 Add this config to the integration test suite
11--job.print_args Print the args to terminal
12--model.config MODEL.CONFIG
13 Path to the model config
14--model.norm_type MODEL.NORM_TYPE
15 Type of layer normalization to use [layernorm,16 np_layernorm, rmsnorm, fused_rmsnorm]17--model.tokenizer_path MODEL.TOKENIZER_PATH
18 Tokenizer path
19--profiling.enable_profiling
20 Whether to enable pytorch profiler
21--profiling.save_traces_folder PROFILING.SAVE_TRACES_FOLDER
22 Trace files location
23--profiling.profile_freq PROFILING.PROFILE_FREQ
24 How often to collect profiler traces,in iterations
25--profiling.enable_memory_snapshot
26 Whether to dump memory snapshot
27--profiling.save_memory_snapshot_folder PROFILING.SAVE_MEMORY_SNAPSHOT_FOLDER
28 Memeory snapshot files location
29--optimizer.name OPTIMIZER.NAME
30 Optimizer to use
31--optimizer.eps OPTIMIZER.EPS
32 Epsilon value for the optimizer.33--optimizer.fused Whether the fused implementation(CUDA only)is used.34--optimizer.scheduler {wsd,cosine,linear}35 Scheduler to use. Currently supported: wsd, cosine,36and linear.37--optimizer.lr OPTIMIZER.LR
38 Learning rate to use
39--optimizer.min_lr_ratio OPTIMIZER.MIN_LR_RATIO
40 Min lr ratio for lr scheduler
41--optimizer.early_step_in_backward
42 Whether to apply optimizer in the backward. Caution,43 optimizer_in_backward isnot compatible with gradients
44 clipping, users should not call
45 register_post_accumulate_grad_hook after the optimizer
46is built.47--training.batch_size TRAINING.BATCH_SIZE
48 Batch size
49--training.seq_len TRAINING.SEQ_LEN
50 Sequence length
51--training.context_len TRAINING.CONTEXT_LEN
52 Max length allowed for each sequence
53--training.varlen Whether to take sequences of variable length asinput54--training.warmup_steps TRAINING.WARMUP_STEPS
55 Steps for lr scheduler warmup, normally 1/5 of
56--training.steps
57--training.gradient_accumulation_steps TRAINING.GRADIENT_ACCUMULATION_STEPS
58 Number of steps to accumulate gradients before
59 updating parameters
60--training.steps TRAINING.STEPS
61 How many train steps to run
62--training.max_norm TRAINING.MAX_NORM
63 Max norm for gradient clipping
64--training.skip_nan_inf
65 Skip batch updates when NaN or INF gradients are
66 encountered during training
67--training.dataset TRAINING.DATASET
68 Dataset to use,with comma separated values
69--training.dataset_name TRAINING.DATASET_NAME
70 The name of the dataset config,with comma separated
71 values if provided
72--training.dataset_split TRAINING.DATASET_SPLIT
73 Dataset split to use,with comma separated values if74 provided
75--training.data_dir TRAINING.DATA_DIR
76 Data dirs to use,with comma separated values if77 provided
78--training.data_files TRAINING.DATA_FILES
79 Data files to use,with comma separated values if80 provided
81--training.data_probs TRAINING.DATA_PROBS
82 Data sampling probabilities,with comma separated
83 values if provided
84--training.streaming Whether to load dataset in streaming mode, used for85 huge dataset
86--training.num_workers TRAINING.NUM_WORKERS
87 Number of subprocesses to use for data loading.088 means that the data will be loaded in the main
89 process.90--training.prefetch_factor TRAINING.PREFETCH_FACTOR
91 Number of batches loaded in advance by each worker.292 means there will be a total of 2* num_workers batches
93 prefetched across all workers.94--training.data_parallel_replicate_degree TRAINING.DATA_PARALLEL_REPLICATE_DEGREE
95 The `data_parallel_replicate_degree` argument
96 specifies the degree of data parallelism for weight
97 replication. When this value is greater than 1,98 weights will be replicated across
99 `data_parallel_replicate_degree` ranks. If
100 `data_parallel_shard_degree` is also greater than 1,101 the parallelism method used is HSDP (Hybrid Sharded
102 Data Parallelism). Otherwise, the parallelism method
103 used is DDP (Distributed Data Parallelism).1 means
104 disabled.105--training.data_parallel_shard_degree TRAINING.DATA_PARALLEL_SHARD_DEGREE
106 The `data_parallel_shard_degree` argument specifies
107 the degree of data parallelism for weight sharding.108 When this value is greater than 1, weights will be
109 sharded across `data_parallel_shard_degree` ranks. If
110 `data_parallel_replicate_degree` is also greater than
1111, the parallelism method used is HSDP (Hybrid Sharded
112 Data Parallelism). Otherwise, the parallelism method
113 used is FSDP (Fully Sharded Data Parallelism).-1114 means leftover ranks will be used (After
115 DP_REPLICATE/SP/PP). Note that only
116 `data_parallel_shard_degree` can be negative.1 means
117 disabled.118--training.enable_cpu_offload
119 Whether to apply CPU offloading of parameters,120 gradients,and optimizer states in FSDP
121--training.tensor_parallel_degree TRAINING.TENSOR_PARALLEL_DEGREE
122 Tensor Parallelism degree.1 means disabled.123--training.disable_loss_parallel
124 Whether to apply loss parallel when sequence parallel
125is enabled
126--training.mixed_precision_param {bfloat16,float32}127 torch dtype to use for parameters when applying mixed
128 precision via FSDP. This feature only takes effect
129 when data_parallel_shard_degree >1130--training.mixed_precision_reduce {float32}131 torch dtype to use for reductions when applying mixed
132 precision via FSDP. This feature only takes effect
133 when data_parallel_shard_degree >1134--training.compile Whether to compile the model
135--training.gc_freq TRAINING.GC_FREQ
136 Python garbage control scheduling interval,in steps
137--training.seed TRAINING.SEED
138 Choose the base RNG seed used for training
139--training.deterministic
140 Use deterministic algorithms wherever possible, may be
141 slower
142--metrics.log_freq METRICS.LOG_FREQ
143 How often to log metrics to TensorBoard,in iterations
144--metrics.enable_tensorboard
145 Whether to log metrics to TensorBoard
146--metrics.disable_color_printing
147 Whether to disable color printing in logs
148--metrics.save_tb_folder METRICS.SAVE_TB_FOLDER
149 Folder to dump TensorBoard states
150--metrics.rank_0_only
151 Whether to save TensorBoard metrics only for rank 0or152forall ranks. When pipeline_parallel_degree is>1,153 this option uses the 0th rank of the last stage
154 pipeline group, which is the only stage that computes
155 loss metrics.156--metrics.enable_wandb
157 Whether to log metrics to Weights & Biases
158--experimental.enable_async_tensor_parallel
159 Whether to applyasync tensor parallel (currently only
160 effective when compileis enabled)161--experimental.pipeline_parallel_degree EXPERIMENTAL.PIPELINE_PARALLEL_DEGREE
162 Pipeline Parallelism degree,or number of ranks.1163 means disabled. If using looped schedules, this still
164 specifies the number of physical ranks,not the number
165 of stages. Stages per rank are inferred from split
166 points degree,and schedule.167--experimental.pipeline_parallel_split_points EXPERIMENTAL.PIPELINE_PARALLEL_SPLIT_POINTS [EXPERIMENTAL.PIPELINE_PARALLEL_SPLIT_POINTS ...]168 Specify comma-separated names of modules to use as the
169 beginning of a split point. e.g."layers.0,layers.2"170 will cause the model to be split into 3 stages, the
171 first containing all the layers up to layers.0, the
172 second containing layers.0and up to layers.2, the
173 third containing layers.2andall the remaining
174 layers. Note: fully-automated splitting may be enabled
175in the future, but currently the split points must be
176 specified manually.177--experimental.pipeline_parallel_schedule EXPERIMENTAL.PIPELINE_PARALLEL_SCHEDULE
178 Specify the Pipeline Parallel schedule to use. The
179 supported schedules are: https://github.com/pytorch/py
180 torch/blob/de4c2a3b4e89d96334dc678d1c3f2ae51a6630a0/to
181 rch/distributed/pipelining/schedules.py#L2161. The182 schedule must be compatible with the split points and183 stages_per_rank. Looped schedules (e.g.184 Interleaved1F1B) require specifying
185 pipeline_parallel_degree = number of ranks,and186 split_points = number of stages -1187--experimental.pipeline_parallel_schedule_csv EXPERIMENTAL.PIPELINE_PARALLEL_SCHEDULE_CSV
188 Specify the path to the pipeline parallel schedule csv
189file to use. The pipeline_parallel_schedule argument
190 must be either PipelineScheduleSingle,191 PipelineScheduleMulti,or _PipelineScheduleRuntime.192--experimental.pipeline_parallel_microbatches EXPERIMENTAL.PIPELINE_PARALLEL_MICROBATCHES
193 How many microbatches to split the global training
194 batch into when using pipeline parallelism. The global195 training batch size must be evenly divisible by the
196 number of microbatches. The default value will be the
197 number of pipeline stages,if unspecified.198--experimental.enable_compiled_autograd
199 Enable CompiledAutograd to compile the backward.200--experimental.context_parallel_degree EXPERIMENTAL.CONTEXT_PARALLEL_DEGREE
201 Context parallelism degree.1 means disabled.202--experimental.context_parallel_rotate_method EXPERIMENTAL.CONTEXT_PARALLEL_ROTATE_METHOD
203 The collective to use in context parallel SDPA for kv
204 shards exchange.'allgather' means to all-gather all205 kv shards on ranks after the first sub-SDPA
206 computation,'alltoall' means to all-to-all shuffle
207 the kv shards. The default value is'allgather'.208--checkpoint.enable_checkpoint
209 Whether to enable checkpoint
210--checkpoint.folder CHECKPOINT.FOLDER
211 The folder to store the checkpoints. When
212 enable_checkpoint isset to true, checkpoints will be
213in{--job.dump_folder}/{--checkpoint.folder}.214--checkpoint.interval_type CHECKPOINT.INTERVAL_TYPE
215 Checkpointing interval unit of measurement ['step',216'seconds']217--checkpoint.interval CHECKPOINT.INTERVAL
218 Checkpointing interval,in steps or seconds depending
219 on --checkpoint.interval_type
220--checkpoint.model_weights_only
221 When model_weights_only=True, only model weights will
222 be saved at the end of training. With this,223 checkpoints can be loaded using `torch.load(...,224 weights_only=True)` after conversion. When
225 model_weights_only=False, the full checkpoint will be
226 saved. A full checkpoint includes model, optimizer and227 train_state, which can be used to resume training. The
228 default value is false.229--checkpoint.export_dtype {float16,bfloat16,float32}230 Converts to the specified precision when training
231 completes and model_weights_only=true. Currently
232 supports float32, float16,and bfloat16. The default
233 value is float32.234--checkpoint.create_seed_checkpoint
235 Initializes the full model without applying
236 parallelisms,and then saves it as a seed checkpoint.237 Note: requires user to call train.py without
238 specifying any parallelisms, e.g. NGPU=1. Could be
239 implemented as a separate script, but this way shares
240 more code.241--checkpoint.async_mode CHECKPOINT.ASYNC_MODE
242 Which async checkpoint mode to use. Currently there
243 are 3 different modes.1."disabled": synchronized
244 checkpointing will be used.2."async":245 torch.distributed.checkpoint.async_save will be used.2461."async_with_pinned_mem": this option utilizes a
247 dedicated pinned memory space and creates a separate
248 process for faster GPU->CPU transfer performance and249 eliminating GIL contention. The cost is increased CPU
250 memory usage. If insufficient CPU memory is available,251 performance may degrade due to memory paging. For most
252 users,"async" should suffice as the performance
253 overhead is typically small (on the order of tens of
254 seconds) compared to checkpointing frequency. This
255 mode can be employed to pursue near-zero checkpointing
256 times (e.g.,<1 second) given appropriate hardware
257 support such as ample CPU memory and fast PCIe.258"disabled"is the default mode.259--checkpoint.keep_latest_k CHECKPOINT.KEEP_LATEST_K
260 Keeps only the latest k checkpoints,and purging older
261 ones. If 0, keep all checkpoints.0is the default
262 value.263--checkpoint.load_step CHECKPOINT.LOAD_STEP
264 Load the checkpoint at the specified step. If -1, load
265 the latest checkpoint.266--float8.enable_float8_linear
267 If true, swaps `torch.nn.Linear` with `Float8Linear`.268 This feature requires you to install 'torchao' which
269 can be found here: https://github.com/pytorch/ao
270--float8.enable_fsdp_float8_all_gather
271 Whether enable float8 all-gather in FSDP
272--float8.precompute_float8_dynamic_scale_for_fsdp
273 Whether precompute float8 scales dynamically for FSDP
274--float8.scaling_type_input {dynamic,delayed}275 float8 scaling forinput, dynamic (default)or delayed
276--float8.scaling_type_weight FLOAT8.SCALING_TYPE_WEIGHT
277 float8 scaling forinput, dynamic (default)or delayed
278--float8.scaling_type_grad_output FLOAT8.SCALING_TYPE_GRAD_OUTPUT
279 float8 scaling forinput, dynamic (default)or delayed
280--comm.init_timeout_seconds COMM.INIT_TIMEOUT_SECONDS
281 Timeout for communication operations, during
282 initialization and first train step.283--comm.train_timeout_seconds COMM.TRAIN_TIMEOUT_SECONDS
284 Timeout for communication operations after the first
285 train step -- usually a tighter bound than during
286 initialization.287--comm.trace_buf_size COMM.TRACE_BUF_SIZE
288 Flight recorder ring buffer size,>0 means recording
289 by default,0 means disabled
290--memory_estimation.enabled
291 Whether to estimate memory usage for FSDP
292--memory_estimation.disable_fake_mode
293 Whether to estimate memory under FakeTensorMode
Training with variable-length inputs
When you set the --training.varlen flag, you're enabling a more efficient training method that packs multiple documents together into a single long sequence, eliminating the need for padding.
This is particularly useful when your dataset contains documents of varying lengths.
Let's break down how --training.seq_len and --training.context_len work in this mode.
--training.seq_len (Packed Sequence Length): This is the total length of the final sequence fed to the model on one device. Instead of processing one document at a time, the dataloader takes multiple documents (each split to sequences no longer than context_len), concatenates them end-to-end, and creates a single long sequence of length seq_len.
--training.context_len (Sample Length): This parameter defines the maximum number of tokens for a single document or sample. If a document from the dataset is longer than context_len, it will be truncated. For example, if --training.context_len is set to 4,096, a document with 5,000 tokens will be cut down to its first 4,096 tokens, leaving the left tokens as another independent sequence, while a document with 3000 tokens remains unchanged.
Training with torch.compile
Starting from torch 2.0, torch.compile has been introduced as a new feature to seamlessly accelerate training processes.
In flame, one can simply enable torch.compile by adding --training.compile flag to your training script.
However, fla has integrated numerous fused kernels for acceleration, which may potentially conflict with torch.compile.
We are actively working on resolving these issues to make compilation transparent to users.
In the meantime, please ensure you are using the latest dependencies.
Specifically, we recommend using torch>=2.6 and triton>=3.0.
Training with multiple datasets
If you wish to train a model with all-round capabilities (e.g., code, math, and multilingual ability), it's necessary to train on multiple datasets.
flame allows training with multiple datasets easily.
For example, you can specify the following arguments to train on 6 datasets with different proportions:
[!NOTE]
We have done this conversion automatically in the training script since our latest updates.
Once training is complete, you may want to convert the distributed checkpoints (DCPs) into the 🤗 format for broader use.
To facilitate this, we provide a straightforward conversion script:
After this, your model will be in the 🤗 format, ready to be shared or deployed.
You can then easily publish your model using the huggingface_hub for wider accessibility.
Continual training
If you wish to build upon a strong pre-trained model (in 🤗 format) and continue training, we also offer a script to convert the 🤗 format model back into DCP format.
This allows you to seamlessly resume training with flame.
Here, <path_to_dcp> is the directory where your distributed checkpoints will be stored.
The checkpoint is intentionally saved at <step-0> within the checkpoint folder to ensure it is loadable by flame during the initial training step, similar to how a seed checkpoint is handled.
Once the conversion is complete, you can proceed with training using flame as usual, continuing from where the pretrained model left off.
Multi-node training
If you have access to multi-node GPUs, consider leveraging them for optimal performance.
This process is straightforward and well-documented in the PyTorch docs.
To set up multi-node training:
Set the environment variables MASTER_ADDR=<ip> and MASTER_PORT=<port> before running the training script across all nodes.
If you're using a job scheduler like Slurm, it will handle these variables for you.
torchtitan provides a Slurm script for multi-node training, which you can use as a reference or starting point.
Custom models
flame supports custom model architectures through seamless integration with the Hugging Face transformers library. To add your own model:
Create a new model directory under custom_models/ (see custom_models/sba for a complete example)
Implement your model classes and configuration:
Define a config class inheriting from PretrainedConfig (see custom_models/sba/config_sba.py for an example)
Create model classes inheriting from PreTrainedModel (see custom_models/sba/modeling_sba.py for an example)
Register your models in __init__.py:
Import your model classes and config classes
Register your models with the AutoModelForCausalLM, AutoModel and AutoConfig classes (see custom_models/sba/__init__.py for an example)
Create a config file for your custom model, just need to specify the model_type to the one you just named for your custom model (example: configs/sba_340m.json).
Training is extremely simple, you can just use the flame.train.py script to train your custom model.
Citation
If you find flame helpful for your work, please consider citing it.
bib
1@software{yang2025flame,
2 title = {Flame: Flash Language Modeling Made Easy},
3 author = {Zhang, Yu and Yang, Songlin},
4 url = {https://github.com/fla-org/flame},
5 month = jan,
6 year = {2025}
7}
