Skip to content

Commit

Permalink
configurator should ask about gradient checkpointing interval
Browse files Browse the repository at this point in the history
  • Loading branch information
bghira committed Dec 3, 2024
1 parent a4e5393 commit 2ba64ab
Showing 1 changed file with 12 additions and 0 deletions.
12 changes: 12 additions & 0 deletions configure.py
Original file line number Diff line number Diff line change
Expand Up @@ -543,6 +543,18 @@ def configure_env():
)
)
env_contents["--gradient_checkpointing"] = "true"
gradient_checkpointing_interval = prompt_user(
"Would you like to configure a gradient checkpointing interval? A value larger than 1 will increase VRAM usage but speed up training by skipping checkpoint creation every Nth layer",
0,
)
try:
if int(gradient_checkpointing_interval) > 0:
env_contents["--gradient_checkpointing_interval"] = int(
gradient_checkpointing_interval
)
except:
print("Could not parse gradient checkpointing interval. Not enabling.")
pass

env_contents["--caption_dropout_probability"] = float(
prompt_user(
Expand Down

0 comments on commit 2ba64ab

Please sign in to comment.