improve vram use w gradient checkpointing (#1167) [skip ci]
This commit is contained in:
@@ -159,6 +159,13 @@ def normalize_config(cfg):
|
|||||||
if isinstance(cfg.pretraining_dataset, dict):
|
if isinstance(cfg.pretraining_dataset, dict):
|
||||||
cfg.pretraining_dataset = [cfg.pretraining_dataset]
|
cfg.pretraining_dataset = [cfg.pretraining_dataset]
|
||||||
|
|
||||||
|
if (
|
||||||
|
cfg.gradient_checkpointing
|
||||||
|
and cfg.unfrozen_parameters is None
|
||||||
|
and cfg.gradient_checkpointing_kwargs is None
|
||||||
|
):
|
||||||
|
cfg.gradient_checkpointing_kwargs = {"use_reentrant": True}
|
||||||
|
|
||||||
log_gpu_memory_usage(LOG, "baseline", cfg.device)
|
log_gpu_memory_usage(LOG, "baseline", cfg.device)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user