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

[Partition] Remove redudant model splitting, Improve Input Model Parsing #1609

Merged
merged 3 commits into from
Mar 9, 2024
+8 −3
Merged

[Partition] Remove redudant model splitting, Improve Input Model Parsing #1609

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e1facc3
[partition] Remove redundant model splitting
a-ys Mar 7, 2024
fafd6b8
[partition] Fix partitioning of safetensors-only models
a-ys Mar 7, 2024
09540b6
Run formatter
a-ys Mar 8, 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: 0 additions & 1 deletion engines/python/setup/djl_python/neuron_utils/model_loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -301,7 +301,6 @@ def partition(self, save_path: str, **kwargs):
self.model_config.save_pretrained(save_path)
self.model = self.load_hf_model()
self.load_path = self.get_load_path()
save_pretrained_split(self.model, self.load_path)
tosterberg marked this conversation as resolved.
Show resolved Hide resolved
self.model = self.load_inf2_model_from_disk()
shutil.copytree(self.load_path, self.split_model_path)

Expand Down
7 changes: 5 additions & 2 deletions serving/docker/partition/properties_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,9 +58,10 @@ def set_and_validate_model_dir(self):
if 'model_dir' in self.properties:
model_dir = self.properties['model_dir']
model_files = glob.glob(os.path.join(model_dir, '*.bin'))
model_files = glob.glob(os.path.join(model_dir, '.safetensors'))
if not model_files:
raise ValueError(
f'No .bin files found in the dir: {model_dir}')
f'No .bin or .safetensors files found in the dir: {model_dir}')
elif 'option.model_id' in self.properties:
self.properties['model_dir'] = self.properties_dir
elif 'option.s3url' in self.properties:
Expand All @@ -70,10 +71,12 @@ def set_and_validate_model_dir(self):
self.properties.pop('option.s3url')
else:
model_files = glob.glob(os.path.join(self.properties_dir, '*.bin'))
model_files = glob.glob(os.path.join(self.properties_dir, '*.safetensors'))
if model_files:
self.properties['model_dir'] = self.properties_dir
else:
raise ValueError('Please specify the model_dir or model_id')
raise ValueError(f'No .bin or .safetensors files found in the dir: {self.properties_dir}'
'\nPlease specify the model_dir or model_id')

def validate_and_correct_checkpoints_json(self):
"""
Expand Down
Loading