Update

[ghstack-poisoned]

Author: fegin

scripts/estimate/estimation.py

@@ -15,6 +15,7 @@
 from torch.distributed._tools.fsdp2_mem_tracker import FSDPMemTracker
 from torch.testing._internal.distributed.fake_pg import FakeStore
 
+from torchtitan.components.ft import init_ft_manager
 from torchtitan.components.optimizer import build_lr_schedulers, build_optimizers
 from torchtitan.config_manager import JobConfig
 from torchtitan.distributed import ParallelDims, utils as dist_utils
@@ -102,7 +103,6 @@ def estimate_memory(job_config: JobConfig):
         if not job_config.memory_estimation.disable_fake_mode
         else contextlib.nullcontext()
     ):
-
         logger.info(
             f"Building {train_spec.name} {job_config.model.flavor} with {model_config}"
         )
@@ -122,7 +122,8 @@ def estimate_memory(job_config: JobConfig):
         model.train()
 
         # build optimizer after applying parallelisms to the model
-        optimizers = build_optimizers([model], job_config)
+        ft_manager = init_ft_manager(job_config)
+        optimizers = build_optimizers([model], job_config, ft_manager)
         lr_schedulers = build_lr_schedulers(optimizers.optimizers, job_config)
         # Post optimizer step model converters hook.
         # e.g. calculate float8 dynamic amax/scale for all-parameter for FSDP2