[Hardware][Neuron] Refactor neuron support (#3471)

2024-03-21 18:22:17 -07:00
parent ea5f14e6ff
commit e90fc21f2e
33 changed files with 615 additions and 549 deletions
--- a/tests/lora/test_worker.py
+++ b/tests/lora/test_worker.py
@@ -33,7 +33,7 @@ def test_worker_apply_lora(sql_lora_files):
                               max_loras=32),
        distributed_init_method=f"file://{tempfile.mkstemp()[1]}",
    )
-    worker.init_model()
+    worker.init_device()
    worker.load_model()

    worker.model_runner.set_active_loras([], LoRAMapping([], []))