How do I use only a portion of a dataset?

[mashdragon]

I noticed that I cannot set a fractional num_epochs value. How do I instruct axolotl to use only a small portion of my dataset?

[NanoCode012]

I noticed that I cannot set a fractional num_epochs value.

I just looked at Trainer docs, and it seems that num_epochs can actually be a float. Currently in Axolotl, we expect it to be an integer. Is it useful for it to be a float?

How do I instruct axolotl to use only a small portion of my dataset?

axolotl/docs/config.qmd

Line 90 in 8fb72cb

shards: # Optional[int] number of shards to split data into

You can use this config in your dataset section. For example, the below would split the data and use only 10% (10 piece = 100%).

datasets:
  - path: fozziethebeat/alpaca_messages_2k_test
    type: chat_template
    shards: 10

[mashdragon]

This is a FFT of Llama-3.2-3B, since I hit GPU OOM trying the fft-8b-liger-fdsp config on Llama 3.1 8B. RAM/VRAM is 96GB RAM/2x24GB VRAM.

With 30GB swap added, it saved successfully on my mini-run.

You could do save_steps: 5 to check if you can successfully save early.

Oh, I wasn't aware, that's great. Due to the sample packing issue, I had to wait an hour for the mini-run to complete, so maybe next time this will help avoid that.

(And by the way, thank you so much for the help! You have already saved me many hours of struggling or trying to locate things within the code on my own!)

[NanoCode012]

Hm, that's weird. 96GB RAM should be enough for 3B model, though I haven't really benchmarked how much RAM liger+fsdp offload uses.

With 30GB swap added, it saved successfully on my mini-run.

That's good to hear. How's the current vram/ram % usage?

(And by the way, thank you so much for the help! You have already saved me many hours of struggling or trying to locate things within the code on my own!)

Happy to help!

[mashdragon]

Yes, I think having users share VRAM/RAM usage is really helpful since I get frequently disappointed by OOM errors on my hardware. It would be nice if the docs listed some benchmarks. Under this configuration:

config.yml

base_model: ./Llama-3.2-3B

plugins:
  - axolotl.integrations.liger.LigerPlugin
liger_rope: true
liger_rms_norm: true
liger_glu_activation: true
liger_fused_linear_cross_entropy: true

strict: false

chat_template: llama3
datasets:
  - path: my-dataset.jsonl
    type: completion
    split: train[60%:70%]
    text_column: text

dataset_prepared_path: last_run_prepared
val_set_size: 0.02
output_dir: ./3b-fft

sequence_len: 4096
sample_packing: true
pad_to_sequence_len: true

wandb_project:
wandb_entity:
wandb_watch:
wandb_name:
wandb_log_model:

gradient_accumulation_steps: 4
micro_batch_size: 3
num_epochs: 1
optimizer: adamw_torch
lr_scheduler: cosine
learning_rate: 2e-5

train_on_inputs: false
group_by_length: false
bf16: auto
fp16:
tf32: false

# From warning, When using FSDP full shard, instead of using `gradient_checkpointing` in TrainingArguments, please use `activation_checkpointing` in `fsdp_config`. The former introduces a redundant AllGather operation in backward pass. Reference: https://github.com/huggingface/transformers/issues/30404
gradient_checkpointing: false
gradient_checkpointing_kwargs:
  use_reentrant: false
early_stopping_patience:
resume_from_checkpoint:
logging_steps: 1
xformers_attention:
flash_attention: true

warmup_steps: 100
evals_per_epoch: 2
eval_table_size:
saves_per_epoch: 1
debug:
deepspeed:
weight_decay: 0.0
fsdp:
  - full_shard
  - auto_wrap
fsdp_config:
  fsdp_limit_all_gathers: true
  fsdp_sync_module_states: true
  fsdp_offload_params: true
  fsdp_use_orig_params: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_backward_prefetch: BACKWARD_PRE
  activation_checkpointing: true
special_tokens:
  pad_token: <|finetune_right_pad_id|>
  eos_token: <|eot_id|>

See the micro batch size of 3. A size of 4 would OOM for me. Also, see in my comment that I got a warning to set activation_checkpointing in fsdp_config rather than gradient_checkpointing. Maybe we should change Axolotl's FSDP examples.

The current VRAM usage during training is 22300 MB/24576 MB for each of the 2 GPUs. System RAM is harder for me to interpret, but free says Used 33G/Free 2.5G/Shared 33G/Buff-cache 62.5G/Available 30G, and top shows this for the two Python workers:

  VIRT    RES    SHR
 69.5g  30.2g  16.9g
 69.0g  30.2g  16.9g

So I think the 33G shared comes from the two workers. So something like 66GB total (including base system)?

I wonder if I could reduce the context size, hopefully not by too much, and possibly still fit in the Llama 3.1 8B finetune with a micro batch size of 1 (for 4096 context it went OOM from what I remember). That would probably be a lot more desired than the 3B model.

[mashdragon]

One pull request is out. I need to wait for my training to complete to test the other one: #2281

[NanoCode012]

Also, see in my comment that I got a warning to set activation_checkpointing in fsdp_config rather than gradient_checkpointing. Maybe we should change Axolotl's FSDP examples.

Could you elaborate on the warning?

[mashdragon]

The split parameter is enabled for local file datasets via #2281.

So you can add something like split: train[:10%] to your dataset's config and preprocess and load only 10% of the dataset that way. Alternatively you can set shard as discussed by @NanoCode012 here.

You can set a fractional epoch via #2282.

[NanoCode012]

Thanks for the PR!