Merge branch 'main' into main

Author: shimizust

.github/workflows/amd-ci.yml

@@ -64,7 +64,8 @@ jobs:
       run: |
         rocm-smi
         python -m pip install --upgrade pip
-        pip install -e .[dev] --extra-index-url https://download.pytorch.org/whl/nightly/rocm${{ matrix.rocm_version }}
+        pip install -e .[dev]
+        pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm${{ matrix.rocm_version }}/
     
     - name: List Python Environments
       run: python -m pip list

.gitignore

@@ -23,4 +23,5 @@ uv.lock
 
 # Benchmark images
 benchmark/visualizations
-.vscode/
+.vscode/
+.coverage

Makefile

@@ -5,7 +5,16 @@ all: checkstyle test test-convergence
 
 # Command to run pytest for correctness tests
 test:
-	python -m pytest --disable-warnings test/ --ignore=test/convergence
+	python -m pytest --disable-warnings \
+		-n auto \
+		--dist=load \
+		--cov=src/liger_kernel \
+		--cov-report=term-missing \
+		--ignore=test/convergence \
+		test/
+	coverage combine
+	coverage report -m
+	coverage html
 
 # Command to run ruff for linting and formatting code
 checkstyle:

README.md

@@ -129,8 +129,8 @@ y = orpo_loss(lm_head.weight, x, target)
 - `triton >= 3.0.0` Install from pypi. (e.g. `pip install triton==3.0.0`)
 
 ```bash
-# Need to pass the url when installing
-pip install -e .[dev] --extra-index-url https://download.pytorch.org/whl/nightly/rocm6.2
+pip install -e .[dev]
+pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.3/
 ```
 
 ### Optional Dependencies
@@ -164,6 +164,9 @@ pip install -e .
 
 # Setup Development Dependencies
 pip install -e ".[dev]"
+
+# NOTE -> For AMD users only
+pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.3/
 ```
 
 

docs/Examples.md

@@ -239,7 +239,7 @@ from liger_kernel.transformers.trainer import LigerORPOTrainer  # noqa: F401
 
 model = AutoModelForCausalLM.from_pretrained(
     "meta-llama/Llama-3.2-1B-Instruct",
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
 )
 
 tokenizer = AutoTokenizer.from_pretrained(

examples/alignment/run_orpo.py

@@ -9,7 +9,7 @@
 
 model = AutoModelForCausalLM.from_pretrained(
     "meta-llama/Llama-3.2-1B-Instruct",
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
 )
 
 tokenizer = AutoTokenizer.from_pretrained(

examples/huggingface/training.py

@@ -48,7 +48,7 @@ def train():
             custom_args.model_name,
             trust_remote_code=True,
             use_cache=False,
-            torch_dtype=torch.bfloat16,
+            dtype=torch.bfloat16,
             # These args will get passed to the appropriate apply_liger_kernel_to_* function
             # to override the default settings
             # cross_entropy=True,
@@ -59,7 +59,7 @@ def train():
             custom_args.model_name,
             trust_remote_code=True,
             use_cache=False,
-            torch_dtype=torch.bfloat16,
+            dtype=torch.bfloat16,
         )
 
     trainer = SFTTrainer(

examples/huggingface/training_multimodal.py

@@ -56,7 +56,7 @@ def construct_model_and_processor(model_name: str, use_liger: bool) -> torch.nn.
         model = Qwen2VLForConditionalGeneration.from_pretrained(
             pretrained_model_name_or_path=model_name,
             use_cache=False,
-            torch_dtype=torch.bfloat16,
+            dtype=torch.bfloat16,
             low_cpu_mem_usage=True,
             attn_implementation="sdpa",
         )

examples/medusa/train.py

@@ -319,7 +319,7 @@ def _model_loader():
         model = model_builder(
             model_args.model_name_or_path,
             cache_dir=training_args.cache_dir,
-            torch_dtype=torch.bfloat16,
+            dtype=torch.bfloat16,
         )
 
         # Freeze the base model

pyproject.toml

@@ -24,6 +24,36 @@ pythonpath = ["src", "."]
 asyncio_mode = "auto"
 log_cli = true
 log_cli_level = "INFO"
+addopts = [
+    "-n", "auto",
+    "--dist=load",                    # use "load" to distribute tests and let pytest-cov combine coverage
+    "--cov=src/liger_kernel",
+    "--cov-report=term-missing",
+    "--cov-report=html",
+    "--cov-config=pyproject.toml",
+    "--durations=0"
+]
+python_files = "test_*.py"
+testpaths = ["test/"]
+
+[tool.coverage.run]
+branch = true
+parallel = true
+source = ["src/liger_kernel"]
+# xdist uses subprocesses; "multiprocessing" is a safe concurrency choice
+concurrency = ["multiprocessing"]
+
+[tool.coverage.paths]
+liger_kernel = [
+  "src/liger_kernel",
+  "*/site-packages/liger_kernel"
+]
+
+[tool.coverage.report]
+omit = ["test/*"]
+show_missing = true
+skip_covered = false
+
 
 [tool.ruff]
 line-length = 120