VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-06 13:47:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #2267 from kohya-ss/fix-github-actions-error

Browse Source 
fix: `str is not "no"` to `str != "no"`
This commit is contained in:
Kohya S.
2026-02-16 08:01:20 +09:00
committed by GitHub
parent 573a7fa06c ef051427df
commit 3265f2edfb
1 changed files with 13 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
library/deepspeed_utils.py
View File

@@ -96,7 +96,7 @@ def prepare_deepspeed_plugin(args: argparse.Namespace):
deepspeed_plugin.deepspeed_config["train_batch_size"] = ( deepspeed_plugin.deepspeed_config["train_batch_size"] = (
args.train_batch_size * args.gradient_accumulation_steps * int(os.environ["WORLD_SIZE"]) args.train_batch_size * args.gradient_accumulation_steps * int(os.environ["WORLD_SIZE"])
) )
deepspeed_plugin.set_mixed_precision(args.mixed_precision) deepspeed_plugin.set_mixed_precision(args.mixed_precision)
if args.mixed_precision.lower() == "fp16": if args.mixed_precision.lower() == "fp16":
deepspeed_plugin.deepspeed_config["fp16"]["initial_scale_power"] = 0 # preventing overflow. deepspeed_plugin.deepspeed_config["fp16"]["initial_scale_power"] = 0 # preventing overflow.
@@ -125,18 +125,18 @@ def prepare_deepspeed_model(args: argparse.Namespace, **models):
class DeepSpeedWrapper(torch.nn.Module): class DeepSpeedWrapper(torch.nn.Module):
def __init__(self, **kw_models) -> None: def __init__(self, **kw_models) -> None:
super().__init__() super().__init__()
self.models = torch.nn.ModuleDict() self.models = torch.nn.ModuleDict()
wrap_model_forward_with_torch_autocast = args.mixed_precision is not "no" wrap_model_forward_with_torch_autocast = args.mixed_precision != "no"
for key, model in kw_models.items(): for key, model in kw_models.items():
if isinstance(model, list): if isinstance(model, list):
model = torch.nn.ModuleList(model) model = torch.nn.ModuleList(model)
if wrap_model_forward_with_torch_autocast: if wrap_model_forward_with_torch_autocast:
model = self.__wrap_model_with_torch_autocast(model) model = self.__wrap_model_with_torch_autocast(model)
assert isinstance( assert isinstance(
model, torch.nn.Module model, torch.nn.Module
), f"model must be an instance of torch.nn.Module, but got {key} is {type(model)}" ), f"model must be an instance of torch.nn.Module, but got {key} is {type(model)}"
@@ -151,7 +151,7 @@ def prepare_deepspeed_model(args: argparse.Namespace, **models):
return model return model
def __wrap_model_forward_with_torch_autocast(self, model): def __wrap_model_forward_with_torch_autocast(self, model):
assert hasattr(model, "forward"), f"model must have a forward method." assert hasattr(model, "forward"), f"model must have a forward method."
forward_fn = model.forward forward_fn = model.forward
@@ -161,20 +161,19 @@ def prepare_deepspeed_model(args: argparse.Namespace, **models):
device_type = model.device.type device_type = model.device.type
except AttributeError: except AttributeError:
logger.warning( logger.warning(
"[DeepSpeed] model.device is not available. Using get_preferred_device() " "[DeepSpeed] model.device is not available. Using get_preferred_device() "
"to determine the device_type for torch.autocast()." "to determine the device_type for torch.autocast()."
) )
device_type = get_preferred_device().type device_type = get_preferred_device().type
with torch.autocast(device_type = device_type): with torch.autocast(device_type=device_type):
return forward_fn(*args, **kwargs) return forward_fn(*args, **kwargs)
model.forward = forward model.forward = forward
return model return model
def get_models(self): def get_models(self):
return self.models return self.models
ds_model = DeepSpeedWrapper(**models) ds_model = DeepSpeedWrapper(**models)
return ds_model return ds_model