Minor fixes in simple_fsdp experiments (#1853)

#1850 removed `name` field in
`TrainSpec`. The experiments in simple_fsdp should also be updated.
Otherwise it won't run.

#1776 added `use_flex_attn`
field to `apply_non_moe_tp()`, which is missing in simple_fsdp
experiments

```
NGPU=8 CONFIG_FILE=./torchtitan/models/llama3/train_configs/debug_model.toml ./run_train.sh --model.name simple_fsdp.llama3 --compile.enable
```

```
NGPU=8 CONFIG_FILE=./torchtitan/models/deepseek_v3/train_configs/debug_model.toml ./run_train.sh --model.name simple_fsdp.deepseek_v3 --compile.enable
```

Author: yiming0416

torchtitan/experiments/simple_fsdp/deepseek_v3/__init__.py

@@ -21,7 +21,6 @@
 
 def get_train_spec() -> TrainSpec:
     return TrainSpec(
-        name="simple_fsdp.deepseek_v3",
         model_cls=SimpleFSDPDeepSeekV3Model,
         model_args=deepseekv3_configs,
         parallelize_fn=parallelize_deepseekv3,

torchtitan/experiments/simple_fsdp/deepseek_v3/parallelize.py

@@ -55,11 +55,13 @@ def parallelize_deepseekv3(
                 "Currently, float8 tensorwise TP is not tested for deepseekv3"
             )
 
+        use_flex_attn = getattr(model.model_args, "use_flex_attn", False)
         apply_non_moe_tp(
             model,
             world_mesh["tp"],
             loss_parallel=not job_config.parallelism.disable_loss_parallel,
             enable_float8_tensorwise_tp=False,
+            use_flex_attn=use_flex_attn,
         )
         maybe_enable_async_tp(job_config, world_mesh["tp"])
 

torchtitan/experiments/simple_fsdp/llama3/__init__.py

@@ -20,7 +20,6 @@
 
 def get_train_spec() -> TrainSpec:
     return TrainSpec(
-        name="simple_fsdp.llama3",
         model_cls=SimpleFSDPTransformer,
         model_args=llama3_configs,
         parallelize_fn=parallelize_llama,