diff --git a/library/train_util.py b/library/train_util.py index 62cd145e1..b9b4eecf8 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -2416,6 +2416,7 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth: "--mixed_precision", type=str, default="no", choices=["no", "fp16", "bf16"], help="use mixed precision / 混合精度を使う場合、その精度" ) parser.add_argument("--full_fp16", action="store_true", help="fp16 training including gradients / 勾配も含めてfp16で学習する") + parser.add_argument("--full_bf16", action="store_true", help="bf16 training including gradients / 勾配も含めてbf16で学習する") parser.add_argument( "--clip_skip", type=int, diff --git a/sdxl_train.py b/sdxl_train.py index 9cf20252d..06cbc5710 100644 --- a/sdxl_train.py +++ b/sdxl_train.py @@ -267,6 +267,14 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module): unet.to(weight_dtype) text_encoder1.to(weight_dtype) text_encoder2.to(weight_dtype) + elif args.full_bf16: + assert ( + args.mixed_precision == "bf16" + ), "full_bf16 requires mixed precision='bf16' / full_bf16を使う場合はmixed_precision='bf16'を指定してください。" + accelerator.print("enable full bf16 training.") + unet.to(weight_dtype) + text_encoder1.to(weight_dtype) + text_encoder2.to(weight_dtype) # acceleratorがなんかよろしくやってくれるらしい if args.train_text_encoder: