Use fake_mode when constructing ShardingOptimizer (#70)

This is particularly important for constructor nodes and for the flop estimation, otherwise they could materialize massive tensors in memory, leading to OOM. This shows up in DeepSeek.

Author: fmassa

autoparallel/api.py

@@ -191,21 +191,18 @@ def __enter__(self):
 
         self.build_model_graph()
 
-        from torch._subclasses.fake_tensor import unset_fake_temporarily
-
-        with unset_fake_temporarily():
-            rescale_grad_comm_cost_for_mp = 1.0
-            if self.mp_policy is not None:
-                param_size = self.mp_policy.param_dtype.itemsize
-                reduce_size = self.mp_policy.reduce_dtype.itemsize
-                if param_size != reduce_size:
-                    rescale_grad_comm_cost_for_mp = reduce_size / param_size
-                    # Tiebreak, favoring performing the comms in the largest
-                    # dtype
-                    rescale_grad_comm_cost_for_mp *= 1.1
-            sharding_optimizer = ShardingOptimizer(
-                self.gm, self.mesh, rescale_grad_comm_cost_for_mp
-            )
+        rescale_grad_comm_cost_for_mp = 1.0
+        if self.mp_policy is not None:
+            param_size = self.mp_policy.param_dtype.itemsize
+            reduce_size = self.mp_policy.reduce_dtype.itemsize
+            if param_size != reduce_size:
+                rescale_grad_comm_cost_for_mp = reduce_size / param_size
+                # Tiebreak, favoring performing the comms in the largest
+                # dtype
+                rescale_grad_comm_cost_for_mp *= 1.1
+        sharding_optimizer = ShardingOptimizer(
+            self.gm, self.mesh, rescale_grad_comm_cost_for_mp
+        )
 
         # makes sharding of params and gradients the same
         sharding_optimizer.add_grad_param_constraints()

autoparallel/propagation_rules.py

@@ -700,7 +700,10 @@ def reshape_rule(mesh, op_schema):
 @register_opschema_rule(torch.ops.aten.expand.default)
 def expand_rule(mesh, op_schema_):
     op = torch.ops.aten.expand.default
-    op_schema = copy.deepcopy(op_schema_)
+    from torch._subclasses.fake_tensor import unset_fake_temporarily
+
+    with unset_fake_temporarily():
+        op_schema = copy.deepcopy(op_schema_)
     input_strat = op_schema.args_schema[0]
     orig_shape = input_strat.strategies[0].output_specs.tensor_meta.shape
     dest_shape = op_schema.args_schema[1]