winglian commited on
Commit
5247c50
1 Parent(s): cf66547

set env for FSDP offload params (#433)

Browse files
Files changed (1) hide show
  1. src/axolotl/utils/trainer.py +2 -0
src/axolotl/utils/trainer.py CHANGED
@@ -355,6 +355,8 @@ def calculate_total_num_steps(cfg, train_dataset, tokenizer):
355
 
356
  def setup_fsdp_envs(cfg):
357
  os.environ["ACCELERATE_USE_FSDP"] = "true"
 
 
358
  if cfg.fsdp_config.fsdp_sync_module_states:
359
  os.environ["FSDP_SYNC_MODULE_STATES"] = "true"
360
  if cfg.fsdp_config.fsdp_state_dict_type:
 
355
 
356
  def setup_fsdp_envs(cfg):
357
  os.environ["ACCELERATE_USE_FSDP"] = "true"
358
+ if cfg.fsdp_config.fsdp_offload_params:
359
+ os.environ["FSDP_OFFLOAD_PARAMS"] = "true"
360
  if cfg.fsdp_config.fsdp_sync_module_states:
361
  os.environ["FSDP_SYNC_MODULE_STATES"] = "true"
362
  if cfg.fsdp_config.fsdp_state_dict_type: