docs: Add docstrings and type hints

Author: rrutmann

src/modalities/models/gpt2/gpt2_model.py

@@ -806,7 +806,7 @@ def __init__(
             ffn_norm_config (LayerNormWrapperConfig): Config for the feed-forward network normalization module.
             lm_head_norm_config (LayerNormWrapperConfig): Config for the language model head normalization module.
             use_weight_tying (bool): Whether to use weight tying.
-            seed (int, optional): The random seed. Defaults to None.
+            seed (Optional[int]): The random seed. Defaults to None.
         """
         weight_decay_groups = {
             "linear": [".attn", ".mlp", ".lm_head.weight"],
@@ -910,6 +910,15 @@ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         ...
 
     def forward(self, inputs: dict[str, torch.Tensor] | torch.Tensor) -> dict[str, torch.Tensor] | torch.Tensor:
+        """
+        Forward pass of the GPT2LLM module.
+
+        Args:
+            inputs (dict[str, torch.Tensor] | torch.Tensor): Input data.
+
+        Returns:
+            dict[str, torch.Tensor] | torch.Tensor: Model output.
+        """
         if isinstance(inputs, dict):
             return {self.prediction_key: self.forward_impl(inputs[self.sample_key])}
         else:

src/modalities/models/model_factory.py

@@ -58,6 +58,15 @@ class ModelFactory:
 
     @staticmethod
     def _is_model_on_meta_device(model: nn.Module) -> bool:
+        """
+        Checks if all parameters and buffers of the model are on the meta device.
+
+        Args:
+            model (nn.Module): The model to check.
+
+        Returns:
+            bool: True if all parameters and buffers are on meta device, False otherwise.
+        """
         meta_counter = 0
         param_counter = 0
         for _, tensor in itertools.chain(model.named_parameters(), model.named_buffers()):

src/modalities/running_env/fsdp/device_mesh.py

@@ -84,7 +84,8 @@ def get_device_mesh(
     enable_loss_parallel: bool,
     world_size: int,
 ) -> DeviceMesh:
-    """Gets the device mesh for the specified parallelism degrees.
+    """
+    Gets the device mesh for the specified parallelism degrees.
 
     Args:
         device_type (str): The device type.

src/modalities/trainer.py

@@ -28,6 +28,24 @@ class ThroughputAggregationKeys(Enum):
 
 
 class Trainer:
+    """
+    Trainer class for model training.
+
+    Args:
+        global_rank (int): The global rank.
+        num_pipeline_parallel_ranks (int): Number of pipeline parallel ranks.
+        progress_publisher (MessagePublisher[ProgressUpdate]): Progress publisher.
+        evaluation_result_publisher (MessagePublisher[EvaluationResultBatch]): Evaluation result publisher.
+        gradient_acc_steps (int): Gradient accumulation steps.
+        global_num_tokens_per_train_step (int): Global number of tokens per train step.
+        num_seen_train_steps (int): Number of seen train steps.
+        global_num_seen_tokens (int): Global number of seen tokens.
+        num_target_steps (int): Number of target steps.
+        num_target_tokens (int): Number of target tokens.
+        gradient_clipper (GradientClipperIF): Gradient clipper.
+        mfu_calculator (Optional[MFUCalculatorABC]): MFU calculator.
+    """
+
     def __init__(
         self,
         global_rank: int,