Fixes to align with latest autoparallel

Author: wconstab

torchtitan/experiments/auto_parallel/parallelize_llama.py

@@ -56,7 +56,7 @@ def input_fn():
     # model = model_fn()
     # return model
 
-    autop = AutoParallel(model, input_fn, world_mesh, device=world_mesh.device_type)
+    autop = AutoParallel(model, input_fn, world_mesh)
     autop.add_parameter_memory_constraint(low=None, high=None)
 
     x_sharding = (Shard(0), Replicate())

torchtitan/train.py

@@ -23,7 +23,7 @@
 )
 from torchtitan.config_manager import ConfigManager, JobConfig
 from torchtitan.distributed import ParallelDims, utils as dist_utils
-
+from torch.distributed.tensor import DTensor
 from torchtitan.protocols.model_converter import build_model_converters
 from torchtitan.tools import utils
 from torchtitan.tools.logging import init_logger, logger
@@ -158,10 +158,13 @@ def param(name):
             from torchtitan.models.llama3.model import precompute_freqs_cis
 
             model.buffers_.get_buffer("freqs_cis").copy_(
-                precompute_freqs_cis(
-                    model_args.dim // model_args.n_heads,
-                    model_args.max_seq_len,
-                    model_args.rope_theta,
+                DTensor.from_local(
+                    precompute_freqs_cis(
+                        model_args.dim // model_args.n_heads,
+                        model_args.max_seq_len,
+                        model_args.rope_theta,
+                    ),
+                    device_mesh=model.buffers_.get_buffer("freqs_cis").device_mesh,
                 )
             )