Spaces:
Running
on
Zero
Running
on
Zero
Update optimization.py
Browse files- optimization.py +2 -2
optimization.py
CHANGED
@@ -43,13 +43,13 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
|
|
43 |
# dynamic_shapes = tree_map_only((torch.Tensor, bool), lambda x: None, call.kwargs)
|
44 |
dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
|
45 |
|
46 |
-
|
47 |
|
48 |
exported = torch.export.export(
|
49 |
mod=pipeline.transformer,
|
50 |
args=call.args,
|
51 |
kwargs=call.kwargs,
|
52 |
-
dynamic_shapes=dynamic_shapes,
|
53 |
)
|
54 |
|
55 |
return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
|
|
|
43 |
# dynamic_shapes = tree_map_only((torch.Tensor, bool), lambda x: None, call.kwargs)
|
44 |
dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
|
45 |
|
46 |
+
quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
|
47 |
|
48 |
exported = torch.export.export(
|
49 |
mod=pipeline.transformer,
|
50 |
args=call.args,
|
51 |
kwargs=call.kwargs,
|
52 |
+
# dynamic_shapes=dynamic_shapes,
|
53 |
)
|
54 |
|
55 |
return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
|