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

fix use_reentrant #2700

Merged
merged 3 commits into from
Dec 19, 2024
Merged

fix use_reentrant #2700

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4d8c665
fix use_reentrant
Jintao-Huang Dec 19, 2024
00735c4
fix
Jintao-Huang Dec 19, 2024
640034b
fix
Jintao-Huang Dec 19, 2024
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
1 change: 1 addition & 0 deletions swift/llm/model/model/internlm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ def load_model(self):
CLIPVisionTower.load_model = load_model

model, tokenizer = get_model_tokenizer_with_flash_attn(model_dir, model_info, model_kwargs, load_model, **kwargs)
model.vit.vision_tower.gradient_checkpointing_enable()
if model is not None:
if version == 'v2' and use_flash_attn:
# fix AttributeError: no attribute 'attention_dropout'
Expand Down
11 changes: 6 additions & 5 deletions swift/llm/model/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -246,26 +246,27 @@ def safe_snapshot_download(model_id_or_path: str,
if not download_model:
ignore_file_pattern += ['*.bin', '*.safetensors']
hub = get_hub(use_hf)
model_id_or_path = model_id_or_path.split(':', 1) # get sub_folder
if len(model_id_or_path) == 1:
model_id_or_path = [model_id_or_path[0], None]
model_id_or_path, sub_folder = model_id_or_path
if model_id_or_path.startswith('~'):
model_id_or_path = os.path.abspath(os.path.expanduser(model_id_or_path))
with safe_ddp_context(hash_id=model_id_or_path):
if os.path.exists(model_id_or_path):
model_dir = model_id_or_path
sub_folder = None
else:
if model_id_or_path.startswith('/'): # startswith
raise ValueError(f"path: '{model_id_or_path}' not found")
model_id_or_path = model_id_or_path.split(':', 1) # get sub_folder
if len(model_id_or_path) == 1:
model_id_or_path = [model_id_or_path[0], None]
model_id_or_path, sub_folder = model_id_or_path
model_dir = hub.download_model(model_id_or_path, revision, ignore_file_pattern, token=hub_token, **kwargs)

logger.info(f'Loading the model using model_dir: {model_dir}')

model_dir = os.path.abspath(os.path.expanduser(model_dir))
assert os.path.isdir(model_dir), f'model_dir: {model_dir}'
if sub_folder:
model_dir = os.path.join(model_dir, sub_folder)
assert os.path.isdir(model_dir), f'model_dir: {model_dir}'
return model_dir


Expand Down
6 changes: 3 additions & 3 deletions swift/trainers/arguments.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,14 +36,14 @@ def _fix_gradient_checkpointing(self):
if hasattr(torch.utils.checkpoint, '_old_checkpoint'): # avoid double patching
return
# Consistent with the default behavior of transformers.
default_use_reentrant = (
use_reentrant_ = (
self.gradient_checkpointing_kwargs.get('use_reentrant', True)
if self.gradient_checkpointing_kwargs else True)
_old_checkpoint = torch.utils.checkpoint.checkpoint

@wraps(_old_checkpoint)
def _new_checkpoint(*args, use_reentrant=default_use_reentrant, **kwargs):
return _old_checkpoint(*args, use_reentrant=use_reentrant, **kwargs)
def _new_checkpoint(*args, use_reentrant=None, **kwargs):
return _old_checkpoint(*args, use_reentrant=use_reentrant_, **kwargs)

torch.utils.checkpoint._old_checkpoint = _old_checkpoint
torch.utils.checkpoint.checkpoint = _new_checkpoint
Expand Down
Loading