remove ownership_tokens logic

Author: H-Huang

torchtitan/distributed/dual_pipe_v.py

@@ -295,11 +295,11 @@ def run_forward():
             arg_mbs[forward_mb_index],
             kwarg_mbs[forward_mb_index],
         )
-        # TODO its error prone to have this logic scattered inside and outside the runtime file..
-        # this goes along with the patch to pytorch: https://github.com/pytorch/pytorch/pull/167002/
-        key = f"{forward_stage.stage_index}_{forward_mb_index}"
-        assert key not in schedule.ownership_tokens
-        schedule.ownership_tokens[key] = output.view_as(output).grad_fn
+        # # TODO its error prone to have this logic scattered inside and outside the runtime file..
+        # # this goes along with the patch to pytorch: https://github.com/pytorch/pytorch/pull/167002/
+        # key = f"{forward_stage.stage_index}_{forward_mb_index}"
+        # assert key not in schedule.ownership_tokens
+        # schedule.ownership_tokens[key] = output.view_as(output).grad_fn
         schedule._maybe_compute_loss(
             forward_stage, output, ctx.target_mbs, forward_mb_index
         )

torchtitan/distributed/pipeline_parallel.py

@@ -11,7 +11,6 @@
 
 import torch
 import torch.nn as nn
-from torch._subclasses.fake_tensor import FakeTensorMode
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.pipelining import PipelineStage
 from torch.distributed.pipelining.schedules import (

torchtitan/models/deepseek_v3/__init__.py

@@ -97,7 +97,7 @@
         qk_rope_head_dim=64,
         v_head_dim=128,
         mscale=0.70,
-        attn_type="flex",
+        # attn_type="flex",
         # attn_mask_type="block_causal",
     ),
     "236B": DeepSeekV3ModelArgs(

torchtitan/models/deepseek_v3/train_configs/debug_model.toml

@@ -65,7 +65,7 @@ export_dtype = "float32"
 async_mode = "disabled"  # ["disabled", "async", "async_with_pinned_mem"]
 
 [activation_checkpoint]
-mode = "none"  # ["none", "selective", "full"]
+mode = "selective"  # ["none", "selective", "full"]
 selective_ac_option = 'op'  # 'int' = ac every positive int layer or 'op', ac based on ops policy
 
 [compile]