feat: add FSDP to standalone

Signed-off-by: Sébastien Han <[email protected]>
redhat-et · Oct 10, 2024 · d92ecfd · d92ecfd
1 parent 46a8374
commit d92ecfd
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/standalone/standalone.py b/standalone/standalone.py
@@ -129,7 +129,7 @@
  --max_batch_len=20000 \
  --seed=42 \
  --cpu_offload_optimizer \
- --sharding_strategy=FULL_SHARD \
+ --distributed_training_framework fsdp \
  --is_granite \
  --checkpoint_at_epoch
  command:
@@ -188,14 +188,14 @@
  --output_dir=/tmp/model \
  --num_epochs={epoch_num} \
  --effective_batch_size=3840 \
- --learning_rate=2e-6 \
+ --learning_rate=1e-4 \
  --num_warmup_steps=800 \
  --save_samples=0 \
  --log_level=INFO \
  --max_batch_len=20000 \
  --seed=42 \
  --cpu_offload_optimizer \
- --sharding_strategy=FULL_SHARD \
+ --distributed_training_framework fsdp \
  --is_granite \
  --checkpoint_at_epoch
  command:

diff --git a/standalone/standalone.tpl b/standalone/standalone.tpl
@@ -114,7 +114,7 @@ spec:
  --max_batch_len=20000 \
  --seed=42 \
  --cpu_offload_optimizer \
- --sharding_strategy=FULL_SHARD \
+ --distributed_training_framework fsdp \
  --is_granite \
  --checkpoint_at_epoch
  command:
@@ -173,14 +173,14 @@ spec:
  --output_dir=/tmp/model \
  --num_epochs={epoch_num} \
  --effective_batch_size=3840 \
- --learning_rate=2e-6 \
+ --learning_rate=1e-4 \
  --num_warmup_steps=800 \
  --save_samples=0 \
  --log_level=INFO \
  --max_batch_len=20000 \
  --seed=42 \
  --cpu_offload_optimizer \
- --sharding_strategy=FULL_SHARD \
+ --distributed_training_framework fsdp \
  --is_granite \
  --checkpoint_at_epoch
  command: