Fix param_init_fn: move if-statement out of lambda

pbelevich · web-flow · commit fb2e802cefe4 · 2024-05-05T15:32:38.000-04:00
diff --git a/src/llama_recipes/finetuning.py b/src/llama_recipes/finetuning.py
@@ -188,7 +188,7 @@ def main(**kwargs):
             device_id=device_id,
             limit_all_gathers=True,
             sync_module_states=train_config.low_cpu_fsdp,
-            param_init_fn=lambda module: module.to_empty(device=torch.device("cuda"), recurse=False)
+            param_init_fn=(lambda module: module.to_empty(device=torch.device("cuda"), recurse=False))
             if train_config.low_cpu_fsdp and rank != 0 else None,
         )
         if fsdp_config.fsdp_activation_checkpointing:

Original file line number	Diff line number	Diff line change
`@@ -188,7 +188,7 @@ def main(**kwargs):`
`188`	`188`	`device_id=device_id,`
`189`	`189`	`limit_all_gathers=True,`
`190`	`190`	`sync_module_states=train_config.low_cpu_fsdp,`
`191`		`- param_init_fn=lambda module: module.to_empty(device=torch.device("cuda"), recurse=False)`
	`191`	`+ param_init_fn=(lambda module: module.to_empty(device=torch.device("cuda"), recurse=False))`
`192`	`192`	`if train_config.low_cpu_fsdp and rank != 0 else None,`
`193`	`193`	`)`
`194`	`194`	`if fsdp_config.fsdp_activation_checkpointing:`