Skip to content

[dreambooth] low precision guard #1916

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jan 5, 2023
Merged

[dreambooth] low precision guard #1916

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7deec93
[dreambooth] low precision guard
williamberman Jan 4, 2023
3612523
fix
williamberman Jan 4, 2023
8237e21
add docs to cli args
williamberman Jan 4, 2023
3faf0f4
Update examples/dreambooth/train_dreambooth.py
williamberman Jan 4, 2023
568ceb7
style
williamberman Jan 4, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions examples/dreambooth/train_dreambooth.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,10 @@ def parse_args(input_args=None):
type=str,
default=None,
required=False,
help="Revision of pretrained model identifier from huggingface.co/models.",
help=(
"Revision of pretrained model identifier from huggingface.co/models. Trainable model components should be"
" float32 precision."
),
)
parser.add_argument(
"--tokenizer_name",
Expand Down Expand Up @@ -140,7 +143,11 @@ def parse_args(input_args=None):
parser.add_argument(
"--center_crop", action="store_true", help="Whether to center crop images before resizing to resolution"
)
parser.add_argument("--train_text_encoder", action="store_true", help="Whether to train the text encoder")
parser.add_argument(
"--train_text_encoder",
action="store_true",
help="Whether to train the text encoder. If set, the text encoder should be float32 precision.",
)
parser.add_argument(
"--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."
)
Expand Down Expand Up @@ -671,6 +678,17 @@ def main(args):
if not args.train_text_encoder:
text_encoder.to(accelerator.device, dtype=weight_dtype)

low_precision_error_string = (
"Please make sure to always have all model weights in full float32 precision when starting training - even if"
" doing mixed precision training. copy of the weights should still be float32."
)

if unet.dtype != torch.float32:
raise ValueError(f"Unet loaded as datatype {unet.dtype}. {low_precision_error_string}")

if args.train_text_encoder and text_encoder.dtype != torch.float32:
raise ValueError(f"Text encoder loaded as datatype {text_encoder.dtype}. {low_precision_error_string}")

# We need to recalculate our total training steps as the size of the training dataloader may have changed.
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
if overrode_max_train_steps:
Expand Down