chore: import upstream snapshot with attribution

deepspeed/inference/v2/checkpoint/in_memory_engine.py

@@ -0,0 +1,40 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Iterable, Tuple
+import torch
+
+from .base_engine import CheckpointEngineBase
+
+
+class InMemoryModelEngine(CheckpointEngineBase):
+    """
+    This "checkpoint" engine uses the existing interface to enable loading parameters into an
+    inference model from a model already instantiated in memory. In general, this is not the
+    recommended way to use the inference engine, and should only be used when absolutely necessary.
+
+    The primary limitation of this approach is that the model must be fully instantiated in memory.
+    In a tensor parallel scenario, this means that the model is either replicated many times in host
+    memory. Currently, it is also recommended to only use this approach for models held in host memory.
+
+    In order to free the memory held by this copy of the model, we delete the model in the first call
+    to `parameters`, so it is not safe to make this call twice.
+    """
+
+    def __init__(self, model: torch.nn.Module) -> None:
+        """
+        Create virtual checkpoint engine for the provided module.
+
+        Args:
+            model (torch.nn.Module): Model to load parameters from.
+        """
+        super().__init__()
+        self.model = model
+
+    def parameters(self) -> Iterable[Tuple[str, torch.Tensor]]:
+        for name, parameter in self.model.named_parameters():
+            yield name, parameter
+
+        del self.model