Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: OpenMDW-1.1 | |
| from typing import Any, Union | |
| from torch import nn | |
| from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict, set_model_state_dict | |
| from torch.distributed.checkpoint.stateful import Stateful | |
| class ModelWrapper(Stateful): | |
| """Wrapper for model state dict handling""" | |
| def __init__(self, model_parts: Union[list[nn.Module], nn.Module]): | |
| if not isinstance(model_parts, list): | |
| model_parts = [model_parts] | |
| self.model_parts = model_parts | |
| def state_dict(self) -> dict[str, Any]: | |
| sd = {} | |
| for model in self.model_parts: | |
| sd.update(get_model_state_dict(model)) | |
| return sd | |
| def load_state_dict(self, state_dict: dict[str, Any]) -> None: | |
| for model in self.model_parts: | |
| set_model_state_dict( | |
| model, | |
| model_state_dict=state_dict, | |
| options=StateDictOptions(strict=False), | |
| ) | |