diff --git a/simpletuner/helpers/training/trainer.py b/simpletuner/helpers/training/trainer.py index 52adac7c1..080bbba57 100644 --- a/simpletuner/helpers/training/trainer.py +++ b/simpletuner/helpers/training/trainer.py @@ -103,7 +103,7 @@ from simpletuner.helpers.training.sdnq_compile import configure_sdnq_compile_mode from simpletuner.helpers.training.state_tracker import StateTracker from simpletuner.helpers.training.validation import Validation, prepare_validation_prompt_list -from simpletuner.helpers.training.wrappers import unwrap_model +from simpletuner.helpers.training.wrappers import rebind_prepared_forward, unwrap_model from simpletuner.helpers.utils import ramtorch as ramtorch_utils from simpletuner.helpers.utils.checkpoint_manager import ( CHECKPOINT_GUARD_FILENAME, @@ -4574,6 +4574,7 @@ def init_prepare_models(self, lr_scheduler): accelerator_state.parallelism_config = standalone_cp_prepare_config for label, prepared in zip(prepared_labels, results): if label == "primary_model": + prepared = rebind_prepared_forward(prepared, primary_model) self.model.set_prepared_model(prepared) # If we skipped device placement for block swap, move the model incrementally if musubi_block_swap_active: @@ -4636,6 +4637,7 @@ def prepare_model_for_inference(self) -> None: self.model.before_accelerator_prepare() prepared_model = self.accelerator.prepare_model(primary_model, evaluation_mode=True) + prepared_model = rebind_prepared_forward(prepared_model, primary_model) self.model.set_prepared_model(prepared_model) prepared_model.eval() diff --git a/simpletuner/helpers/training/wrappers.py b/simpletuner/helpers/training/wrappers.py index 0fdac6d46..f98f90b1a 100644 --- a/simpletuner/helpers/training/wrappers.py +++ b/simpletuner/helpers/training/wrappers.py @@ -1,3 +1,5 @@ +from types import MethodType + from diffusers.utils.torch_utils import is_compiled_module @@ -28,6 +30,20 @@ def unwrap_model(accelerator, model, keep_fp32_wrapper: bool = True): return _unwrap_execution_wrappers(model) +def rebind_prepared_forward(prepared, original): + if original is None or prepared is original: + return prepared + target = prepared + while target is not None and target is not original: + for attribute in ("forward", "_original_forward"): + bound = target.__dict__.get(attribute) + function = getattr(bound, "__func__", None) + if function is not None and getattr(bound, "__self__", None) is original: + setattr(target, attribute, MethodType(function, target)) + target = target.__dict__.get("_modules", {}).get("module") + return prepared + + def gather_dict_of_tensors_shapes(tensors: dict) -> dict: def _shape(value): if isinstance(value, dict): diff --git a/tests/test_wrappers_rebind.py b/tests/test_wrappers_rebind.py new file mode 100644 index 000000000..3614cb95e --- /dev/null +++ b/tests/test_wrappers_rebind.py @@ -0,0 +1,227 @@ +import unittest +from types import MethodType + +import torch +import torch._dynamo +from torch import nn + +from accelerate.utils.operations import convert_outputs_to_fp32 +from accelerate.utils.other import compile_regions + +from simpletuner.helpers.training.wrappers import rebind_prepared_forward + +OptimizedModule = torch._dynamo.eval_frame.OptimizedModule + + +class _Block(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4, bias=False) + with torch.no_grad(): + self.linear.weight.copy_(torch.eye(4)) + + def forward(self, x): + return self.linear(x) + + +class _Tiny(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([_Block(), _Block()]) + self.gradient_checkpointing = False + self.executed_by = [] + + def forward(self, x): + self.executed_by.append(self) + if self.gradient_checkpointing: + x = x * 2 + for block in self.blocks: + x = block(x) + return x + + +def _install_mixed_precision_forward(model): + autocast_context = torch.autocast(device_type="cpu", dtype=torch.bfloat16) + model._original_forward = model.forward + model_forward_func = model.forward.__func__ + new_forward = autocast_context(model_forward_func) + model.forward = MethodType(new_forward, model) + model.forward = MethodType(convert_outputs_to_fp32(model.forward.__func__), model) + return model + + +def _prepare_like_accelerate(backend="eager"): + original = _Tiny() + _install_mixed_precision_forward(original) + twin = compile_regions(original, backend=backend) + return original, twin + + +class RebindPreparedForwardTests(unittest.TestCase): + def test_compile_regions_twin_inherits_forward_bound_to_the_original(self): + original, twin = _prepare_like_accelerate() + + self.assertIsNot(twin, original) + self.assertIsInstance(twin.blocks[0], OptimizedModule) + self.assertNotIsInstance(original.blocks[0], OptimizedModule) + self.assertIs(twin.__dict__["forward"].__self__, original) + self.assertIs(twin.__dict__["_original_forward"].__self__, original) + + twin.gradient_checkpointing = True + x = torch.ones(1, 4) + out = twin(x) + + self.assertIs(twin.executed_by[-1], original) + self.assertFalse(original.gradient_checkpointing) + torch.testing.assert_close(out, x) + + def test_rebind_makes_the_twin_the_executing_module(self): + original, twin = _prepare_like_accelerate() + + returned = rebind_prepared_forward(twin, original) + + self.assertIs(returned, twin) + self.assertIs(twin.__dict__["forward"].__self__, twin) + self.assertIs(twin.__dict__["_original_forward"].__self__, twin) + self.assertIs(original.__dict__["forward"].__self__, original) + + twin.gradient_checkpointing = True + x = torch.ones(1, 4) + out = twin(x) + + self.assertIs(twin.executed_by[-1], twin) + self.assertIsInstance(twin.executed_by[-1].blocks[0], OptimizedModule) + torch.testing.assert_close(out, x * 2) + + def test_rebind_applies_with_the_inductor_backend(self): + original, twin = _prepare_like_accelerate(backend="inductor") + + self.assertIs(twin.__dict__["forward"].__self__, original) + rebind_prepared_forward(twin, original) + self.assertIs(twin.__dict__["forward"].__self__, twin) + + def test_noop_when_prepared_is_original(self): + original = _install_mixed_precision_forward(_Tiny()) + before = original.__dict__["forward"] + + self.assertIs(rebind_prepared_forward(original, original), original) + self.assertIs(original.__dict__["forward"], before) + self.assertIs(original.__dict__["forward"].__self__, original) + + def test_noop_when_original_is_none(self): + original = _install_mixed_precision_forward(_Tiny()) + twin = compile_regions(original, backend="eager") + before = twin.__dict__["forward"] + + self.assertIs(rebind_prepared_forward(twin, None), twin) + self.assertIs(twin.__dict__["forward"], before) + self.assertIs(twin.__dict__["forward"].__self__, original) + + def test_noop_when_accelerate_already_rebound_the_twin(self): + original, twin = _prepare_like_accelerate() + twin.forward = MethodType(twin.__dict__["forward"].__func__, twin) + twin._original_forward = MethodType(twin.__dict__["_original_forward"].__func__, twin) + before_forward = twin.__dict__["forward"] + before_original_forward = twin.__dict__["_original_forward"] + + self.assertIs(rebind_prepared_forward(twin, original), twin) + self.assertIs(twin.__dict__["forward"], before_forward) + self.assertIs(twin.__dict__["_original_forward"], before_original_forward) + self.assertIs(twin.__dict__["forward"].__self__, twin) + self.assertIs(original.__dict__["forward"].__self__, original) + + def test_noop_without_an_instance_forward(self): + original = _Tiny() + twin = compile_regions(original, backend="eager") + + self.assertNotIn("forward", twin.__dict__) + rebind_prepared_forward(twin, original) + self.assertNotIn("forward", twin.__dict__) + self.assertNotIn("_original_forward", twin.__dict__) + + def test_noop_for_plain_torch_compile_optimized_module(self): + original = _install_mixed_precision_forward(_Tiny()) + compiled = torch.compile(original, backend="eager") + before = compiled.__dict__.get("forward") + + self.assertIsInstance(compiled, OptimizedModule) + self.assertIs(rebind_prepared_forward(compiled, original), compiled) + self.assertIs(compiled.__dict__.get("forward"), before) + self.assertIs(original.__dict__["forward"].__self__, original) + + def test_rebinds_the_inner_twin_under_ddp(self): + import torch.distributed as dist + from torch.nn.parallel import DistributedDataParallel + + dist.init_process_group("gloo", store=dist.HashStore(), rank=0, world_size=1) + try: + original = _install_mixed_precision_forward(_Tiny()) + twin = compile_regions(DistributedDataParallel(original), backend="eager") + + self.assertIsInstance(twin, DistributedDataParallel) + self.assertIsNot(twin.module, original) + self.assertIs(twin.module.__dict__["forward"].__self__, original) + + self.assertIs(rebind_prepared_forward(twin, original), twin) + self.assertIs(twin.module.__dict__["forward"].__self__, twin.module) + self.assertIs(twin.module.__dict__["_original_forward"].__self__, twin.module) + self.assertIs(original.__dict__["forward"].__self__, original) + + twin.module.gradient_checkpointing = True + x = torch.ones(1, 4) + out = twin(x) + + self.assertIs(twin.module.executed_by[-1], twin.module) + self.assertIsInstance(twin.module.executed_by[-1].blocks[0], OptimizedModule) + torch.testing.assert_close(out, x * 2) + finally: + dist.destroy_process_group() + + def test_noop_when_accelerate_already_rebound_the_inner_twin_under_ddp(self): + import torch.distributed as dist + from torch.nn.parallel import DistributedDataParallel + + dist.init_process_group("gloo", store=dist.HashStore(), rank=0, world_size=1) + try: + original = _install_mixed_precision_forward(_Tiny()) + twin = compile_regions(DistributedDataParallel(original), backend="eager") + inner = twin.module + inner.forward = MethodType(inner.__dict__["forward"].__func__, inner) + inner._original_forward = MethodType(inner.__dict__["_original_forward"].__func__, inner) + before_forward = inner.__dict__["forward"] + before_original_forward = inner.__dict__["_original_forward"] + + self.assertIs(rebind_prepared_forward(twin, original), twin) + self.assertIs(inner.__dict__["forward"], before_forward) + self.assertIs(inner.__dict__["_original_forward"], before_original_forward) + self.assertNotIn("forward", twin.__dict__) + self.assertIs(original.__dict__["forward"].__self__, original) + + inner.gradient_checkpointing = True + x = torch.ones(1, 4) + torch.testing.assert_close(twin(x), x * 2) + self.assertIs(inner.executed_by[-1], inner) + finally: + dist.destroy_process_group() + + def test_noop_for_a_module_wrapper_with_dot_module(self): + original = _install_mixed_precision_forward(_Tiny()) + + class _Wrapper(nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + def forward(self, *args, **kwargs): + return self.module(*args, **kwargs) + + wrapper = _Wrapper(original) + + self.assertNotIn("forward", wrapper.__dict__) + self.assertIs(rebind_prepared_forward(wrapper, original), wrapper) + self.assertNotIn("forward", wrapper.__dict__) + self.assertIs(original.__dict__["forward"].__self__, original) + + +if __name__ == "__main__": + unittest.main()