new api

Author: JenniferWang

src/forge/actors/generator.py

@@ -50,14 +50,10 @@
 )
 
 from forge.controller import ForgeActor, get_proc_mesh, stop_proc_mesh
-<<<<<<< HEAD:src/forge/actors/generator.py
 from forge.data_models.completion import Completion
 from forge.data_models.prompt import to_prompt
 from forge.env import TORCHSTORE_USE_RDMA
 from forge.interfaces import Policy as GeneratorInterface
-from forge.data.sharding import VLLMSharding
-from forge.data_models.completion import Completion
-from forge.data_models.prompt import to_prompt
 from forge.observability.metrics import record_metric, Reduce
 from forge.observability.perf_tracker import Tracer
 from forge.types import ProcessConfig

src/forge/actors/trainer.py

@@ -18,17 +18,6 @@
 import torch.distributed.checkpoint as dcp
 import torchstore as ts
 
-from forge.actors._torchstore_utils import (
-    DcpHandle,
-    get_dcp_whole_state_dict_key,
-    get_param_key,
-)
-
-from forge.controller import ForgeActor
-from forge.data.utils import batch_to_device
-from forge.observability.metrics import record_metric, Reduce
-from forge.observability.perf_tracker import Tracer
-
 from monarch.actor import current_rank, current_size, endpoint
 from torch import Tensor
 from torch.distributed.checkpoint._nested_dict import flatten_state_dict

tests/integration_tests/test_policy_update.py

@@ -264,7 +264,7 @@ async def test_sanity_check(self, _setup_and_teardown):
             for _, e in errs.items():
                 assert not e, f"Validation failed with exception: {e}"
 
-        await policy.update_weights.fanout(policy_version=v1)
+        await policy.update_weights.fanout(version=v1)
         all_errs = await policy._test_validate_model_params.fanout(
             _test_validate_params_all_zeros
         )
@@ -273,7 +273,7 @@ async def test_sanity_check(self, _setup_and_teardown):
                 assert not e, f"Validation failed with exception: {e}"
 
         # Reloading v0, getting back original weights
-        await policy.update_weights.fanout(policy_version=v0)
+        await policy.update_weights.fanout(version=v0)
         all_errs = await policy._test_validate_model_params.fanout(
             _test_validate_params_unchanged
         )