feat: add torch.compile CLI option for training

- Add --compile flag (default/reduce-overhead/max-autotune)
- Apply torch.compile in _before_wrap before DDP/FSDP wrapping
- Profiling shows MFU 85.5% -> 88.5% (+3%), time -3.2%, memory -7.9%

astrai/config/train_config.py

@@ -45,6 +45,12 @@ class TrainConfig(BaseConfig):
         default_factory=list,
         metadata={"help": "Module types to enable activation checkpointing for."},
     )
+    compile_mode: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "torch.compile mode: 'default', 'reduce-overhead', 'max-autotune', or None to disable."
+        },
+    )
 
     # checkpoint setting
     start_epoch: int = field(default=0, metadata={"help": "Start epoch for training."})