improve vram use w gradient checkpointing (#1167) [skip ci]
Browse files
src/axolotl/utils/config.py
CHANGED
|
@@ -159,6 +159,13 @@ def normalize_config(cfg):
|
|
| 159 |
if isinstance(cfg.pretraining_dataset, dict):
|
| 160 |
cfg.pretraining_dataset = [cfg.pretraining_dataset]
|
| 161 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 162 |
log_gpu_memory_usage(LOG, "baseline", cfg.device)
|
| 163 |
|
| 164 |
|
|
|
|
| 159 |
if isinstance(cfg.pretraining_dataset, dict):
|
| 160 |
cfg.pretraining_dataset = [cfg.pretraining_dataset]
|
| 161 |
|
| 162 |
+
if (
|
| 163 |
+
cfg.gradient_checkpointing
|
| 164 |
+
and cfg.unfrozen_parameters is None
|
| 165 |
+
and cfg.gradient_checkpointing_kwargs is None
|
| 166 |
+
):
|
| 167 |
+
cfg.gradient_checkpointing_kwargs = {"use_reentrant": True}
|
| 168 |
+
|
| 169 |
log_gpu_memory_usage(LOG, "baseline", cfg.device)
|
| 170 |
|
| 171 |
|