set env for FSDP offload params (#433)
Browse files
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:
|