use logger in ft (#1539)

Summary:
- wasn't seeing print statements getting printed
- the statements show up using the logger
- also added some logging to validate the model is being split for
diloco

Author: tushar00jain

torchtitan/components/ft/diloco/utils.py

@@ -7,6 +7,7 @@
 import torch.nn as nn
 from torchtitan.components.ft.config import FaultTolerance as FTConfig
 from torchtitan.distributed.pipeline_parallel import generate_llm_fqn_per_model_part
+from torchtitan.tools.logging import logger
 
 
 def module_split(
@@ -98,7 +99,9 @@ def _build_fragment_from_modules(
             fragment_idx,
             module_names,
         )
-        print(f"building fragment_idx {fragment_idx} " f"with modules {module_names}")
+        logger.info(
+            f"building fragment_idx {fragment_idx} " f"with modules {module_names}"
+        )
         model_fragments.append(model_fragment)
 
     return model_fragments
@@ -118,13 +121,14 @@ def fragment_llm(
 
     if module_fqns_per_model_fragment == []:
         if ft_config.num_fragments == 1:
+            logger.info("Created 1 model fragments")
             return [model]
 
         module_fqns_per_model_fragment = generate_llm_fqn_per_model_part(
             ft_config.num_fragments, n_layers, input_weight, output_weight
         )
 
     model_fragments = module_split(model, module_fqns_per_model_fragment)
-    print(f"Created {len(model_fragments)} model fragments")
+    logger.info(f"Created {len(model_fragments)} model fragments")
 
     return model_fragments

torchtitan/components/ft/manager.py

@@ -16,6 +16,7 @@
 from torch.distributed._composable.fsdp.fully_shard import FSDPModule
 from torch.distributed.distributed_c10d import ReduceOp
 from torchtitan.components.ft.config import FaultTolerance as FTConfig
+from torchtitan.tools.logging import logger
 
 if importlib.util.find_spec("torchft") is not None:
     import torchft as ft
@@ -125,6 +126,9 @@ def maybe_semi_sync_training(
         assert (
             ft_manager._manager is not None
         ), "FTManager must be enabled to use semi-sync training."
+        logger.info(
+            f"using fragment function to split model: {fragment_fn is not None}"
+        )
         if semi_sync_method.lower() == "diloco":
             if fragment_fn:
                 model_parts = fragment_fn(model, ft_config, n_layers)