Streaming dataset from gcs is not streaming?

[denadai2]

Dear community, I am trying to use a large dataset from gcs via streaming and sft but it seems the streaming does not work... Axolotl processes the entire dataset before training, resulting in OOM. What am I missing?

thx

[2025-09-28 10:10:01,335] [INFO] [axolotl.utils.data.wrappers] [RANK:1] Loading dataset: gs://.... (hidden) with base_type: chat_template and prompt_style: None

Tokenizing Prompts (num_proc=48):   0%|          | 0/13857512 [00:00<?, ? examples/s]
Tokenizing Prompts (num_proc=48):   0%|          | 0/13857512 [00:00<?, ? examples/s][2025-09-28 10:10:01,885] [INFO] [axolotl.utils.data.wrappers] [RANK:3] Loading dataset: gs://prolific-experiments/data-endpoints/processed_data/prism-processed-data-collection-rkmeans_4x256_idlearning-llama3.2_1b_it-merged-train/hf_dataset with base_type: chat_template and prompt_style: None

Tokenizing Prompts (num_proc=48):   0%|          | 0/13857512 [00:00<?, ? examples/s]
Tokenizing Prompts (num_proc=48):   0%|          | 1000/13857512 [00:03<13:14:35, 290.64 examples/s]
Tokenizing Prompts (num_proc=48):   0%|          | 2000/13857512 [00:03<6:18:29, 610.12 examples/s] [2025-09-28 10:10:04,768] [WARNING] [axolotl.prompt_strategies.chat_template] Content end boundary is the same as start boundary for turn 1. This is likely an empty turn.

Tokenizing Prompts (num_proc=48):   0%|          | 1000/13857512 [00:03<14:54:02, 258.31 examples/s]
Tokenizing Prompts (num_proc=48):   0%|          | 3000/13857512 [00:04<4:49:34, 797.39 examples/s]
Tokenizing Prompts (num_proc=48):   0%|          | 2000/13857512 [00:04<7:03:23, 545.41 examples/s]
Tokenizing Prompts (num_proc=48):   0%|          | 4000/13857512 [00:05<3:35:27, 1071.65 examples/s]

conf:

base_model:  /mnt/llama3.2_1b

sequence_len: 4096

# Streaming-specific settings
streaming: true
streaming_multipack_buffer_size: 10000
shuffle_merged_datasets: true

dataset_keep_in_memory: false
datasets:
  - path: gs://... (hidden)
    ds_type: parquet
    type: chat_template
    field_messages: PFM_PROMPT_MESSAGES
    message_property_mappings:
      role: from
      content: value
output_dir: ./outputs/

gradient_accumulation_steps: 2
micro_batch_size: 1
num_epochs: 1
max_steps: 100

optimizer: adamw_torch
lr_scheduler: constant_with_warmup
learning_rate: 2e-5

flash_attention: true

gradient_checkpointing: true
activation_offloading: true

logging_steps: 1
saves_per_epoch: 1

warmup_ratio: 0.03

fsdp_version: 2
fsdp_config:
  offload_params: false
  auto_wrap_policy: TRANSFORMER_BASED_WRAP
  transformer_layer_cls_to_wrap: LlamaDecoderLayer
  state_dict_type: FULL_STATE_DICT
  sharding_strategy: FULL_SHARD
  reshard_after_forward: true
  activation_checkpointing: false

[denadai2]

this does not happen when i download the dataset locally and load the axolotl dataset from he local dir

[djsaunde]

Hi @denadai2, at what point does it OOM? I only see in your logs loading up to 4K examples, but your streaming_multipack_buffer_size is set to 10K.

[denadai2]

@djsaunde I did not post all the logs because the point is that the streaming is not working. The dataset is processed locally when it comes from gcs, not in streaming.

[NanoCode012]

Could you provide more trace up to that point? It would be helpful for us to debug this issue.

Some other questions including those Dan has mentioned:

• How far did it get past the tokenizing prompt? What's the memory usage at that point?
• How much RAM do you have on your system?
• Is this local, HPC or a cloud system?
• How large is this dataset (size / tokens )?

[djsaunde]

Not really sure what is meant by streaming is not working. Can you post the full logs? If your streaming_multipack_buffer_size is too large, you may OOM. Streaming multipacking requires preprocessing chunks of samples at a time

[denadai2]

it means that the pipeline processes the entire dataset before starting the training like all 13857512 examples