Update on "Re-enable FSDP2 Mem Tracker integration tests"

When `run_llama_train.sh` was simplified, we moved the launch of memory estimation and tracking to `run_memory_estimation.sh`. But changing the `test_runner.py ` was missed which in-turn caused the integration test to be a no-op. This PR re-enables integration tests for the `FSDP2MemTracker`. 

[ghstack-poisoned]

Author: Sanket Jayant Purandare

estimation.py

@@ -14,7 +14,6 @@
 from torch._subclasses.fake_tensor import FakeTensorMode
 from torch.distributed import destroy_process_group
 from torch.distributed._tools.fsdp2_mem_tracker import FSDPMemTracker
-from torch.distributed.tensor.parallel import loss_parallel
 from torch.testing._internal.distributed.fake_pg import FakeStore
 
 from torchtitan.config_manager import JobConfig