address comments

Author: wwwjn

torchtitan/experiments/gpt_oss/model/args.py

@@ -72,7 +72,7 @@ class GptOssModelArgs(BaseModelArgs):
     n_kv_heads: int = 8
     sliding_window_size: int = 128
     attn_mask_type: str = "causal"
-    use_flex_attn: bool = True
+    use_flex_attn: bool = True  # NOTE: gpt-oss only support FlexAttention
     # yarn
     original_seq_len: int = 4096
     rope_theta: float = 150000.0

torchtitan/experiments/gpt_oss/train_configs/debug_model.toml

@@ -41,33 +41,30 @@ seq_len = 2048
 max_norm = 1.0  # grad norm clipping
 steps = 10
 dataset = "c4_test"  # supported datasets: c4_test (2K), c4 (177M)
-seed = 0
-deterministic = true
 
 [parallelism]
 data_parallel_replicate_degree = 1
 data_parallel_shard_degree = -1
 fsdp_reshard_after_forward = "default" # default / never / always
-tensor_parallel_degree = 2
+tensor_parallel_degree = 1
 enable_async_tensor_parallel = false
 expert_parallel_degree = 1
 expert_tensor_parallel_degree = 1
 
 [checkpoint]
-enable = true
+enable = false
 folder = "checkpoint"
 interval = 10
 last_save_model_only = false
 export_dtype = "float32"
 async_mode = "disabled"  # ["disabled", "async", "async_with_pinned_mem"]
-initial_load_path = "/data/users/jianiw/torchtitan/outputs/checkpoint-seed/step-0"
 
 [activation_checkpoint]
 mode = "none"  # ["none", "selective", "full"]
 selective_ac_option = '2'  # 'int' = ac every positive int layer or 'op', ac based on ops policy
 
 [compile]
-enable=false
+enable = false
 components = ["model", "loss"]
 
 [quantize.linear.float8]