update aoti compile (#3)

- update aoti compile (9aa2a04d546ddb4bb6deb5d4ea24e016061b33cd)
- Update optimization_utils.py (8e5bed5da4bd6b198ad33fdbb55dbe760f3ea32a)

optimization.py

@@ -122,13 +122,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.forward = combined_transformer_1
+    drain_module_parameters(pipeline.transformer)
 
-    pipeline.transformer = combined_transformer_1
-    pipeline.transformer.config = transformer_config # pyright: ignore[reportAttributeAccessIssue]
-    pipeline.transformer.dtype = transformer_dtype # pyright: ignore[reportAttributeAccessIssue]
-
-    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: {'device': tensor.device, 'dtype': tensor.dtype} 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(**meta)