drain_module_parameters

optimization.py

@@ -12,8 +12,9 @@ from torchao.quantization import quantize_
 from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
 from torchao.quantization import Int8WeightOnlyConfig
 
-from optimization_utils import capture_component_call
 from optimization_utils import aoti_compile
+from optimization_utils import capture_component_call
+from optimization_utils import drain_module_parameters
 from optimization_utils import ZeroGPUCompiledModel
 
 
@@ -105,13 +106,8 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
         else:
             return cp2(*args, **kwargs)
 
-    transformer_config = pipeline.transformer.config
-    transformer_dtype = pipeline.transformer.dtype
-
-    pipeline.transformer = combined_transformer_1
-    pipeline.transformer.config = transformer_config # pyright: ignore[reportAttributeAccessIssue]
-    pipeline.transformer.dtype = transformer_dtype # pyright: ignore[reportAttributeAccessIssue]
+    pipeline.transformer.forward = combined_transformer_1
+    drain_module_parameters(pipeline.transformer)
 
-    pipeline.transformer_2 = combined_transformer_2
-    pipeline.transformer_2.config = transformer_config # pyright: ignore[reportAttributeAccessIssue]
-    pipeline.transformer_2.dtype = transformer_dtype # pyright: ignore[reportAttributeAccessIssue]
+    pipeline.transformer_2.forward = combined_transformer_2
+    drain_module_parameters(pipeline.transformer_2)

optimization_utils.py

@@ -96,3 +96,12 @@ def capture_component_call(
         except CapturedCallException as e:
             captured_call.args = e.args
             captured_call.kwargs = e.kwargs
+
+
+def drain_module_parameters(module: torch.nn.Module):
+    state_dict_meta = {name: tensor.to('meta') for name, tensor in module.state_dict().items()}
+    state_dict = {name: torch.nn.Parameter(torch.empty_like(tensor, device='cpu')) for name, tensor in module.state_dict().items()}
+    module.load_state_dict(state_dict, assign=True)
+    for name, param in state_dict.items():
+        meta = state_dict_meta[name]
+        param.data = torch.Tensor([]).to(device=meta.device, dtype=meta.dtype)