Skip to content
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

Remove explicit copy of model.forward() and instead check/store the id to unwrap #2104

Closed
wants to merge 2 commits into from
Closed

Remove explicit copy of model.forward() and instead check/store the id to unwrap #2104

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
31df404
Working accelerate tests
muellerzr Oct 31, 2023
f371557
Oneline
muellerzr Oct 31, 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
3 changes: 1 addition & 2 deletions src/accelerate/accelerator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1357,8 +1357,8 @@ def prepare_model(self, model: torch.nn.Module, device_placement: bool = None, e
" Please rerun your script specifying `--num_processes=1` or by launching with `python {{myscript.py}}`."
)

model._accelerate_original_forward_id = id(model.forward)
if self.native_amp:
model._original_forward = model.forward
model_forward_func = model.forward.__func__ if hasattr(model.forward, "__func__") else model.forward
autocast_context = get_mixed_precision_context_manager(self.native_amp, self.autocast_handler)
new_forward = autocast_context(model_forward_func)
Expand All @@ -1372,7 +1372,6 @@ def prepare_model(self, model: torch.nn.Module, device_placement: bool = None, e
with torch.no_grad():
convert_model(model)
model._converted_to_transformer_engine = True
model._original_forward = model.forward

kwargs = self.fp8_recipe_handler.to_kwargs() if self.fp8_recipe_handler is not None else {}
if "fp8_format" in kwargs:
Expand Down
10 changes: 6 additions & 4 deletions src/accelerate/utils/other.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,14 +86,16 @@ def extract_model_from_parallel(model, keep_fp32_wrapper: bool = True):
model = model.module

if not keep_fp32_wrapper:
original_forward_id = getattr(model, "_accelerate_original_forward_id", None)
forward = getattr(model, "forward")
original_forward = model.__dict__.pop("_original_forward", None)
if original_forward is not None:
forward_func = getattr(forward, "__func__", None)
if (original_forward_id is not None) and (id(forward) != original_forward_id):
while hasattr(forward, "__wrapped__"):
forward = forward.__wrapped__
if forward == original_forward:
forward_func = getattr(forward, "__func__", None)
if id(forward) == original_forward_id or id(forward_func) == original_forward_id:
break
model.forward = MethodType(forward, model)
model.forward = MethodType(forward_func if forward_func is not None else forward, model)
if getattr(model, "_converted_to_transformer_engine", False):
convert_model(model, to_transformer_engine=False)

Expand Down
6 changes: 3 additions & 3 deletions tests/test_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,15 +132,15 @@ def test_patch_environment_key_exists(self):

def test_can_undo_convert_outputs(self):
model = RegressionModel()
model._original_forward = model.forward
model._accelerate_original_forward_id = id(model.forward)
model.forward = convert_outputs_to_fp32(model.forward)
model = extract_model_from_parallel(model, keep_fp32_wrapper=False)
_ = pickle.dumps(model)

@require_cuda
def test_can_undo_fp16_conversion(self):
model = RegressionModel()
model._original_forward = model.forward
model._accelerate_original_forward_id = id(model.forward)
model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
model.forward = convert_outputs_to_fp32(model.forward)
model = extract_model_from_parallel(model, keep_fp32_wrapper=False)
Expand All @@ -150,7 +150,7 @@ def test_can_undo_fp16_conversion(self):
@require_torch_min_version(version="2.0")
def test_dynamo(self):
model = RegressionModel()
model._original_forward = model.forward
model._accelerate_original_forward_id = id(model.forward)
model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
model.forward = convert_outputs_to_fp32(model.forward)
model.forward = torch.compile(model.forward, backend="inductor")
Expand Down
Loading