Rocm support

.github/workflows/test.yaml

@@ -29,6 +29,6 @@ jobs:
           pip install -e .
 
       - name: Run test
-        timeout-minutes: 10
+        timeout-minutes: 30
         run: |
-          python tests/ci/gpu_lock_exec.py --devices 6,7 python -m unittest discover -s ./tests -p "test_*.py"
+          python -m unittest discover -s ./tests -p "test_*.py" -v

requirements-rocm.txt

@@ -0,0 +1,19 @@
+# Use the PyTorch ROCm wheel index (choose the stream that matches your system)
+--extra-index-url https://download.pytorch.org/whl/rocm6.3
+
+pre-commit
+torch==2.8.0+rocm6.3
+torchaudio==2.8.0+rocm6.3
+torchvision==0.23.0+rocm6.3
+transformers==4.57.1
+qwen-vl-utils==0.0.11
+datasets
+setuptools
+tqdm
+wandb
+psutil
+numpy
+accelerate
+pydantic
+sglang[all]==0.5.4
+openai-harmony

tests/test_target_modeling/test_auto_target_model.py

@@ -24,10 +24,10 @@ def test_auto_target_model(self):
 
         set_seed(42)
         hf_model = HFLlamaForCausalLM.from_pretrained(
-            "dphn/Dolphin3.0-Llama3.1-8B", torch_dtype=torch.bfloat16
+            "unsloth/Llama-3.2-1B", torch_dtype=torch.bfloat16
         ).cuda()
         sf_model = SFLlamaForCausalLM.from_pretrained(
-            "dphn/Dolphin3.0-Llama3.1-8B", torch_dtype=torch.bfloat16
+            "unsloth/Llama-3.2-1B", torch_dtype=torch.bfloat16
         ).cuda()
 
         # create data