add model_parts ref to MetricsProcessor (#1578)

Adds a `ModelProtocol.get_extra_metrics` method for more flexible custom
metric reporting, as discussed in #1576

Probably this should be an abstract method, but I was wary of making
this a breaking change for users who inherit this commit.

The current signature is `get_extra_metrics(self, parallel_dims:
ParallelDims) -> None | dict`. I also considered adding some subset of
`JobConfig`, `TrainSpec`, and `pp_has_{first,last}_stage`; not sure what
else might be useful.

Tested via running the debugmodel with print statements. 

CC @rakkit @wwwjn

Author: garrett361

torchtitan/components/metrics.py

@@ -319,6 +319,7 @@ class MetricsProcessor:
     num_flops_per_token: int
     optimizers: OptimizersContainer | None
     lr_schedulers: LRSchedulersContainer | None
+    model_parts: list[torch.nn.Module] | None
 
     def __init__(
         self,
@@ -349,6 +350,7 @@ def __init__(
         self.num_flops_per_token = -1
         self.optimizers = None
         self.lr_schedulers = None
+        self.model_parts = None
 
     def should_log(self, step: int) -> bool:
         return step == 1 or step % self.job_config.metrics.log_freq == 0

torchtitan/train.py

@@ -294,6 +294,7 @@ def __init__(self, job_config: JobConfig):
             )
         )
         self.metrics_processor.optimizers = self.optimizers
+        self.metrics_processor.model_parts = self.model_parts
 
         # Initialize trainer states that will be saved in checkpoint.
         # These attributes must be initialized before checkpoint loading.