Streaming SFT support (#3101)
* working * fixes * deprecate --iterable; cleanup * pretrain_multipack_buffer_size -> streaming_multipack_buffer_size * improvements * tests * remove unused * docs, examples * nit * nit * add val_set_size validation * val * nit * min * coderabbito * cleanup * nit * add depr warning, cleanup * nit * fix test, fix quarto * fix * review comments * review comments * fix
This commit is contained in:
50
examples/streaming/README.md
Normal file
50
examples/streaming/README.md
Normal file
@@ -0,0 +1,50 @@
|
||||
# Streaming Dataset Examples
|
||||
|
||||
This directory contains example configurations for using Axolotl's streaming dataset
|
||||
functionality, which enables memory-efficient training with large datasets.
|
||||
|
||||
## Examples
|
||||
|
||||
Run the following examples with e.g. `axolotl train examples/streaming/sft.yaml`; no
|
||||
`axolotl preprocess` required!
|
||||
|
||||
### Pretraining (`pretrain.yaml`)
|
||||
|
||||
Demonstrates streaming configuration for pretraining tasks using the fineweb-edu dataset
|
||||
with SmolLM2-135M.
|
||||
|
||||
- Uses `pretraining_dataset` configuration for automatic streaming
|
||||
- Multipack attention control to prevent cross-attention between packed sequences
|
||||
- Buffer size configuration for memory management
|
||||
|
||||
### SFT (`sft.yaml`)
|
||||
|
||||
Shows how to use streaming for supervised fine-tuning with the Alpaca dataset.
|
||||
|
||||
- Explicit `streaming: true` flag for SFT datasets
|
||||
- Memory-efficient training on instruction datasets
|
||||
- Evaluation datasets are currently not streamed
|
||||
|
||||
## Key Configuration Options
|
||||
|
||||
### `streaming`
|
||||
- Enables streaming mode for standard datasets
|
||||
- Automatically enabled for `pretraining_dataset`
|
||||
|
||||
### `streaming_multipack_buffer_size`
|
||||
- Controls buffer size for sample packing (default: 10,000)
|
||||
- Larger values improve packing efficiency but use more memory
|
||||
- Adjust based on available memory
|
||||
|
||||
### `shuffle_merged_datasets`
|
||||
- Enables shuffling of streaming datasets
|
||||
- Requires additional memory for shuffle buffer
|
||||
|
||||
### `sample_packing`
|
||||
- Packs multiple samples into single sequences
|
||||
- Minimize per-step padding tokens
|
||||
|
||||
## Performance Tips
|
||||
|
||||
- Download small / frequently-used datasets locally for better performance
|
||||
- Larger buffer sizes improve packing efficiency
|
||||
57
examples/streaming/pretrain.yaml
Normal file
57
examples/streaming/pretrain.yaml
Normal file
@@ -0,0 +1,57 @@
|
||||
base_model: HuggingFaceTB/SmolLM2-135M
|
||||
|
||||
# Streaming pretraining configuration
|
||||
pretraining_dataset:
|
||||
- path: HuggingFaceFW/fineweb-edu
|
||||
name: sample-10BT
|
||||
type: pretrain
|
||||
text_column: text
|
||||
split: train
|
||||
|
||||
# Streaming-specific settings
|
||||
streaming_multipack_buffer_size: 10000
|
||||
shuffle_merged_datasets: true
|
||||
|
||||
# Training configuration
|
||||
max_steps: 1000
|
||||
output_dir: ./outputs/smollm2-135m-pretrain-streaming
|
||||
|
||||
# Sequence and packing settings
|
||||
sequence_len: 1024
|
||||
sample_packing: true
|
||||
pretrain_multipack_attn: true # Prevent cross-attention between packed sequences
|
||||
flash_attention: true
|
||||
|
||||
# Batch size settings
|
||||
gradient_accumulation_steps: 8
|
||||
micro_batch_size: 1
|
||||
|
||||
# Optimizer and scheduler
|
||||
optimizer: adamw_torch
|
||||
lr_scheduler: cosine
|
||||
learning_rate: 5e-4
|
||||
warmup_ratio: 0.1
|
||||
weight_decay: 0.01
|
||||
|
||||
# Precision and performance
|
||||
bf16: auto
|
||||
tf32: true
|
||||
|
||||
# Logging and checkpointing
|
||||
logging_steps: 10
|
||||
save_strategy: steps
|
||||
save_steps: 250
|
||||
save_total_limit: 3
|
||||
|
||||
# Weights & Biases (optional)
|
||||
wandb_project:
|
||||
wandb_entity:
|
||||
wandb_watch:
|
||||
wandb_name:
|
||||
wandb_log_model:
|
||||
|
||||
# Special tokens
|
||||
special_tokens:
|
||||
pad_token: "<|endoftext|>"
|
||||
|
||||
# save_first_step: true # uncomment this to validate checkpoint saving works with your config
|
||||
55
examples/streaming/sft.yaml
Normal file
55
examples/streaming/sft.yaml
Normal file
@@ -0,0 +1,55 @@
|
||||
base_model: HuggingFaceTB/SmolLM2-135M
|
||||
|
||||
# Dataset configuration
|
||||
datasets:
|
||||
- path: tatsu-lab/alpaca
|
||||
type: alpaca
|
||||
split: train
|
||||
|
||||
# Streaming-specific settings
|
||||
streaming: true
|
||||
streaming_multipack_buffer_size: 10000
|
||||
shuffle_merged_datasets: true
|
||||
|
||||
# Training configuration
|
||||
max_steps: 1000
|
||||
output_dir: ./outputs/smollm2-135m-sft-streaming
|
||||
|
||||
# Sequence and packing settings
|
||||
sequence_len: 1024
|
||||
sample_packing: true
|
||||
flash_attention: true
|
||||
|
||||
# Batch size settings
|
||||
gradient_accumulation_steps: 4
|
||||
micro_batch_size: 1
|
||||
|
||||
# Optimizer and scheduler
|
||||
optimizer: adamw_torch
|
||||
lr_scheduler: cosine
|
||||
learning_rate: 2e-4
|
||||
warmup_ratio: 0.1
|
||||
weight_decay: 0.0
|
||||
|
||||
# Precision and performance
|
||||
bf16: auto
|
||||
tf32: true
|
||||
|
||||
# Logging and checkpointing
|
||||
logging_steps: 10
|
||||
save_strategy: steps
|
||||
save_steps: 100
|
||||
save_total_limit: 3
|
||||
|
||||
# Weights & Biases (optional)
|
||||
wandb_project:
|
||||
wandb_entity:
|
||||
wandb_watch:
|
||||
wandb_name:
|
||||
wandb_log_model:
|
||||
|
||||
# Special tokens
|
||||
special_tokens:
|
||||
pad_token: "<|endoftext|>"
|
||||
|
||||
# save_first_step: true # uncomment this to validate checkpoint saving works with your config
|
||||
Reference in New Issue
Block a user