[WEIGHT LOADER] Add support for serverless_llm format loader
#13440
Conversation
- Introduced 'serverless_llm' as a new load format option. - Added methods to save and load Serverless LLM model states in the model runner and loader. - Updated relevant classes and configurations to accommodate the new model loading mechanism. Signed-off-by: Xinyuan Tong <[email protected]>
Signed-off-by: Xinyuan Tong <[email protected]>
JustinTong0323
requested review from
CatherineSue,
Fridge003,
Ying1123,
hnyls2002,
ispobock,
merrymercy,
slin1237,
xiezhq-hermann and
zhyncs
as code owners
Summary of ChangesHello @JustinTong0323, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request integrates support for the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for the serverless_llm format loader, enabling the system to load and save model states from ServerlessLLM checkpoint stores. The changes are well-integrated across various components, including the LoadFormat enum, Engine API, Scheduler mixin, and ModelRunner. A new ServerlessLLMModelLoader class is implemented, mirroring the behavior of existing loaders like ShardedStateLoader for handling sharded tensors and managing model state. The implementation includes necessary utility methods for filtering subtensors, managing device maps, and ensuring all parameters are correctly loaded and saved. Overall, the changes are logical and extend the system's capabilities effectively.
| def _filter_subtensors(tensors: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | ||
| """Filter out view/sub-tensors that share storage with other tensors.""" | ||
| same_storage_groups: Dict[Any, List[Tuple[str, torch.Tensor]]] = ( | ||
| collections.defaultdict(list) | ||
| ) | ||
| for key, tensor in tensors.items(): | ||
| if tensor.numel(): | ||
| ptr = tensor.untyped_storage().data_ptr() | ||
| same_storage_groups[tensor.device, ptr].append((key, tensor)) | ||
|
|
||
| def get_end_ptr(tensor: torch.Tensor) -> int: | ||
| return tensor.view(-1)[-1].data_ptr() + tensor.element_size() | ||
|
|
||
| result: Dict[str, torch.Tensor] = {} | ||
| for group in same_storage_groups.values(): | ||
| for k, t in group: | ||
| a, b = t.data_ptr(), get_end_ptr(t) | ||
| for k2, t2 in group: | ||
| if not t2.is_contiguous(): | ||
| continue | ||
| a2, b2 = t2.data_ptr(), get_end_ptr(t2) | ||
| if a < a2 or b2 < b: | ||
| continue | ||
| if a2 < a or b < b2 or not t.is_contiguous(): | ||
| break | ||
| if k2 < k: | ||
| break | ||
| else: | ||
| result[k] = t | ||
| return result |
There was a problem hiding this comment.
The _filter_subtensors method is duplicated here from ShardedStateLoader. To improve maintainability and reduce code duplication, consider extracting this common utility into a shared helper function or a base class method if applicable. 1
Rules References
Footnotes
-
Avoid code duplication to improve maintainability and reduce the risk of inconsistencies when changes are needed. ↩
Signed-off-by: Xinyuan Tong <[email protected]>
- Added try-except blocks around the import of load_dict and its usage to log errors if the import fails or if loading the model state dictionary encounters issues. - Improved logging messages to guide users on potential installation or server issues. Signed-off-by: Xinyuan Tong <[email protected]>
2 participants
Motivation
Add support for https://github.com/ServerlessLLM/ServerlessLLM
cc @future-xy @SecretSettler
Modifications
Accuracy Tests
Benchmarking and Profiling
Checklist