areal-project · garrett4wade · Jul 21, 2025 · Jul 14, 2025 · Jul 14, 2025 · Jul 14, 2025
diff --git a/arealite/env/__init__.py b/arealite/env/__init__.py
diff --git a/arealite/engine/autotp_engine.py → arealite/experimental/autotp_engine.py b/arealite/engine/autotp_engine.py → arealite/experimental/autotp_engine.py
diff --git a/arealite/engine/sglang_engine.py → arealite/experimental/sglang_engine.py b/arealite/engine/sglang_engine.py → arealite/experimental/sglang_engine.py
diff --git a/arealite/tests/test_sglang_local_engine.py → ...imental/tests/test_sglang_local_engine.py b/arealite/tests/test_sglang_local_engine.py → ...imental/tests/test_sglang_local_engine.py
@@ -23,7 +23,7 @@
     SGLangConfig,
 )
 from arealite.api.io_struct import LLMRequest, LLMResponse
-from arealite.engine.sglang_engine import SGLangEngine
+from arealite.experimental.sglang_engine import SGLangEngine
 from arealite.workflow.rlvr import RLVRWorkflow
 from realhf.api.core.data_api import load_hf_tokenizer
 from realhf.base import seeding

diff --git a/arealite/reward/__init__.py b/arealite/reward/__init__.py
diff --git a/arealite/tests/test_train_engine.py b/arealite/tests/test_train_engine.py
@@ -52,8 +52,8 @@ def mock_input(
 
 
 def get_engine(engine_type: str, model_path: str):
-    from arealite.engine.autotp_engine import DeepSpeedAutoTPEngine
     from arealite.engine.fsdp_engine import FSDPEngine
+    from arealite.experimental.autotp_engine import DeepSpeedAutoTPEngine
 
     engine_cls = {"auto_tp": DeepSpeedAutoTPEngine, "fsdp": FSDPEngine}[engine_type]
 
@@ -135,7 +135,7 @@ def test_train_batch(engine, mock_input):
 
 @torch.no_grad()
 def test_hf_save_load_weights(tmp_path_factory, engine, mock_input):
-    from arealite.engine.autotp_engine import DeepSpeedAutoTPEngine
+    from arealite.experimental.autotp_engine import DeepSpeedAutoTPEngine
 
     if isinstance(engine, DeepSpeedAutoTPEngine):
         print("AutoTP engine does not support HF save/load for now.")