linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity
897 Commits 69 Branches 33 Tags
6b7e3dd6faa564632a3df4faefe08f72cbf23deb
Commit Graph
39 Commits
Author SHA1 Message Date
Buffett LiuandVaibhav Jindal 2120862ba5 [Megatron] Add SwiGLU integration (#1326) 
## Summary
Adds Liger's Triton SwiGLU to Megatron-Core. Megatron's dense MLP routes
its activation through `bias_swiglu_impl` in
`megatron.core.fusions.fused_bias_swiglu`, a TorchScript `@jit_fuser`
implementation. This PR replaces it with a **new fused gate-up Triton
kernel** written for Megatron's memory layout.

At the Megatron training shape (`S=2048, B=4, bf16`) on **H100**,
against `megatron-core==0.16.1`:

|  | Forward | Full (fwd+bwd) | Memory | Memory (`in_place=True`) |
|---|---|---|---|---|
| `ffn_local=4096` | **+57%** | **+146%** | +20% | **−20%** |
| `ffn_local=32768` | **+63%** | **+166%** | +20% | **−20%** |

<img width="1680" height="1200" alt="megatron_swiglu_panel_H100"
src="https://github.com/user-attachments/assets/0c0a2710-42df-4354-a0ce-bd7a25837058"
/>

### Background

Liger's existing `LigerSiLUMulFunction` takes gate and up as **two
separate tensors** (`a`, `b`), because that is how HuggingFace models
arrive: `gate_proj(x)` and `up_proj(x)` are distinct GEMMs.

Megatron is the opposite. `linear_fc1` is a single fused
`ColumnParallelLinear` producing one `[s, b, 2 * ffn_local]` tensor with
gate in columns `[0, n)` and up in `[n, 2n)`. Bridging to the two-tensor
kernel therefore requires a `torch.chunk`, and `chunk` on the last
dimension returns **non-contiguous views**, so `@ensure_contiguous`
materializes 2 full copies on every forward. We measured that bridge
before writing anything new.


| H100, forward, `ffn_local=32768` | Time | vs Megatron |
|---|---|---|
| Megatron `bias_swiglu_impl` | 0.853 ms | — |
| Liger via `chunk` bridge | 1.828 ms | **−53%** (2.1× slower) |
| **Liger fused gate-up (this PR)** | **0.536 ms** | **+59%** |

The chunk bridge is **3.4× slower than the fused kernel** and also
inflated memory by +40%. The new `LigerFusedGateUpSiLUMulFunction` reads
both halves via a column offset into the single buffer, with no copies,
no `cat`. Input row stride is `2n`, output row stride is `n`.


## How to use
### Mode 1: One-line monkey patch

```python
from liger_kernel.megatron import apply_liger_kernel_to_megatron

apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True, swiglu=True) # swiglu defaults to False
```
Patches `bias_swiglu_impl` across Megatron's core modules
(`fused_bias_swiglu`, `mlp`, and `moe.shared_experts`).

### Mode 2: Use the class directly

```python
from liger_kernel.megatron import LigerMegatronSwiGLU

swiglu = LigerMegatronSwiGLU()
out = swiglu(intermediate, bias, fp8_input_store, cpu_offload_input)
```

The signature mirrors `bias_swiglu_impl` positionally. To use Mode 2,
subclass MLP (`see examples/megatron/run_mode2_hand_spec.py`).

### Optional: `in_place=True` for a 20% memory reduction

```python
LigerMegatronSwiGLU(in_place=True) # Opt-in, off by default
```
Overwrites the linear_fc1 output buffer during backward passes to reduce
memory by 20% vs Megatron with zero speed penalty.

## Scope
- Supported: Dense MLP (`bias_swiglu_impl`) and shared experts at any TP
size with `bias=None`.
- Transparent Fallbacks to Megatron: Auto-delegates to Megatron (and
logs once) when:
  - `bias is not None` (Liger lacks bias support)
  - `fp8_input_store=True` (requires FP8 input retention)
  - `cpu_offload_input=True` (requires offloadable inputs)
- Pass `fallback_impl=None` to raise an error instead of falling back.
- Out of Scope: MoE routed experts (`weighted_bias_swiglu_impl`).

## What's added

| File | Purpose |
|---|---|
| `src/liger_kernel/ops/swiglu.py` | New
`LigerFusedGateUpSiLUMulFunction` + fwd/bwd Triton kernels for the
`[tokens, 2n]` layout. Existing two-tensor path untouched. |
| `src/liger_kernel/megatron/swiglu.py` | `LigerMegatronSwiGLU`
nn.Module, fallback policy, shape validation. |
| `src/liger_kernel/megatron/monkey_patch.py` | `swiglu=` flag,
`_patch_bias_swiglu_impl()`, `_rebind_consumer_symbol()`. |
| `src/liger_kernel/megatron/__init__.py` | Re-exports
`LigerMegatronSwiGLU`. |
| `test/megatron/test_swiglu.py` | 36 tests — 28 dependency-free + 8
against real `megatron-core`. |
| `test/megatron/test_monkey_patch.py` | +14 SwiGLU patch-mechanism
tests; stub-harness snapshot/restore rewrite. |
| `benchmark/scripts/benchmark_megatron_swiglu.py` | 4 providers:
`liger`, `liger_in_place`, `megatron`, `torch`. |
| `examples/megatron/*` | Mode 1 + Mode 2 extended to cover SwiGLU. |

## Megatron feature compatibility

**Verified:** Tensor Parallelism (TP), fp32 + bf16, 2D and 3D inputs,
distributed checkpointing (save/load round-trip in both example
scripts), real `MLP` forward+backward parity against Megatron's
TorchScript, and both consumer module paths.

**Should work, not yet E2E-verified:** multi-rank patch application
(every rank must call `apply_...` before model build — single-process
patching is verified, the multi-rank ordering is not), CP,
`use_distributed_optimizer=True`.

**Out of scope:** activation recompute and CUDA-graph capture *with
`in_place=True`* — precisely why that flag is opt-in. The default path
is recompute-safe.

## Known issues / follow-ups
1. `in_place=True` and `retain_graph=True` incompatibility: In-place
Triton kernel writes don't bump PyTorch's `_version` counter, so running
backward a second time silently yields garbage gradients without
triggering autograd guards. We added a `ctx.already_backward` flag to
explicitly raise an error when reusing the graph with `in_place=True`.
2. MoE routed experts: `weighted_bias_swiglu_impl` requires a
`per_token_scale`-aware kernel (left as a follow-up).
3. RMSNorm patch fix: Fixed a TypeError in
`_patch_local_spec_provider_layer_norm` caused by forwarding an
unsupported has_residual kwarg on `megatron-core 0.16.1`. Signature
matching is now version-agnostic via `inspect.signature`.
4. Example script fixes: Fixed three pre-existing `megatron-core 0.16.1`
compatibility issues in the example scripts (AttributeError on
`as_mlp_submodule`, Makefile build errors in pip wheels, and invalid
tokenizer metadata).

## Testing Done
- Hardware Type: NVIDIA A10G (correctness, 1× and 2×), H100 80GB HBM3
(benchmarks), B200 (benchmarks)
- [x] `pytest test/megatron/` **with** `megatron-core` installed — **175
passed, 49 skipped**
- [x] `pytest test/megatron/` **without** `megatron-core` (default CI
shape) — **168 passed, 56 skipped**
- [x] `pytest test/transformers/test_swiglu.py` — **49 passed, 4
skipped, 4 xfailed** (existing two-tensor SwiGLU not regressed by the
shared `ops/swiglu.py` change)
- [x] run `make checkstyle` to ensure code style — clean; `prek run -a`
also clean
- [x] H100 + B200 benchmark sweeps; H100 rows committed to
`all_benchmark_data.csv`
- [ ] run `make test-convergence` to ensure convergence

## Benchmark reproduction
```bash
python benchmark/scripts/benchmark_megatron_swiglu.py --overwrite

python benchmark/benchmarks_visualizer.py \
    --kernel-name megatron_swiglu --metric-name speed --overwrite
```

---------

Co-authored-by: Vaibhav Jindal <[email protected]>
 2026-08-10 22:36:34 +00:00
PrathyushaVaibhav JindalClaude Opus 4.7
 		0400d31357 [Megatron] Add cross-entropy integration (#1207) 
## Summary

This PR adds Liger's Triton cross-entropy kernel to Megatron-LM.

Megatron has two cross-entropy code paths. This PR replaces both with
Liger:

- The **fused** path: `fused_vocab_parallel_cross_entropy` (used when
`cross_entropy_loss_fusion=True`).
- The **unfused** path: `vocab_parallel_cross_entropy` (used when
`cross_entropy_loss_fusion=False`).

You get Liger's kernel no matter which path your Megatron config picks.

Liger's kernel uses online softmax. It writes gradients in place. It
never builds the full softmax tensor. The net effect is faster training
and lower memory use.

This PR also keeps the RMSNorm patch that was added in PR #1254. So one
call patches both kernels at once.

## How to use

There are two ways. Pick whichever fits your code better.

### Mode 1: One-line monkey patch

Best for users who don't want to touch their model code. Call this once
before you build the model:

```python
from liger_kernel.megatron import apply_liger_kernel_to_megatron

apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True)
```

That's the whole API. No CE-specific knobs. The defaults match
Megatron's native behavior.

### Mode 2: Use the classes directly

Best for users who build their model spec by hand or need custom config.
Import the classes and wire them in yourself:

```python
from liger_kernel.megatron import LigerMegatronRMSNorm, LigerMegatronCrossEntropy

# In your TransformerBlockSubmodules:
input_layernorm = LigerMegatronRMSNorm
pre_mlp_layernorm = LigerMegatronRMSNorm

# In a GPTModel subclass:
class MyModel(GPTModel):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.liger_ce = LigerMegatronCrossEntropy(
            ignore_index=-100,
            label_smoothing=0.1,
        )

    def compute_language_model_loss(self, labels, logits):
        labels_sb = labels.transpose(0, 1).contiguous()
        loss_sb = self.liger_ce(logits, labels_sb, self.pg_collection.tp)
        return loss_sb.transpose(0, 1).contiguous()
```

The class accepts the same knobs Liger's standalone CE accepts
(`ignore_index`, `label_smoothing`, etc.). Use Mode 2 when you need
them.

## Scope

This PR supports `tensor_model_parallel_size=1` only.

Why? At TP>1, each GPU holds a slice of the vocabulary, not the full
vocab. Cross-entropy then needs cross-rank all-reduces to get a correct
loss. Liger's kernel does not do those all-reduces. So the math would be
wrong at TP>1.

To stop you from training with a broken loss, the patch raises
`RuntimeError` if it sees TP>1. The check runs at patch time and again
at call time, so misconfigurations fail loudly.

**[TODO]**: TP>1 support is planned as follow-up work (a separate
"vocab-parallel" Triton kernel).

You can still use this PR with pipeline parallelism, data parallelism,
or single-GPU runs. Those are all TP=1 cases.


## Testing

All unit tests live in `test/megatron/`.

### Run all megatron tests

```bash
python -m pytest test/megatron/ --no-cov
```

That runs three files:

- `test/megatron/test_rms_norm.py` — RMSNorm class correctness (forward,
backward, dtypes).
- `test/megatron/test_cross_entropy.py` — `LigerMegatronCrossEntropy`
class correctness. Covers shape sweeps, `ignore_index`,
`label_smoothing`, combined config, out-of-bounds targets,
`torch.no_grad()` path, and gradient parity vs. PyTorch's
`F.cross_entropy`.
- `test/megatron/test_monkey_patch.py` —
`apply_liger_kernel_to_megatron()` plumbing. Covers symbol replacement
for both CE paths and both RMSNorm targets, idempotency,
missing-megatron errors, TP>1 guard, dispatch behavior, and end-to-end
correctness through the patched symbols.

You should see all tests pass (currently 105 tests in the megatron
suite).

### Run a single test file

```bash
python -m pytest test/megatron/test_cross_entropy.py -v --no-cov
```

### What the tests do NOT need

Megatron-LM is **not** a test dependency. The monkey-patch tests use
stub modules in `sys.modules`, so they run on CPU without a real
Megatron install. You only need PyTorch + Triton + Liger itself.

The correctness tests (RMSNorm, CE) need a GPU because they run the
actual Triton kernel.

## Benchmarking

### Step 1: Run the benchmark

```bash
python benchmark/scripts/benchmark_megatron_cross_entropy.py --overwrite
```

This compares four providers across vocab sizes 4K → 128K (with S=2048,
B=4, BF16):

- **liger** — `LigerMegatronCrossEntropy`
- **torch** — vanilla `F.cross_entropy`
- **megatron** — Megatron's fused CE (`@jit_fuser` via TorchScript)
- **megatron-unfused** — Megatron's unfused CE (eager Python)

The script needs a CUDA/ROCm GPU. If Megatron-LM is not installed, the
two megatron providers are silently skipped and you only get `liger` vs
`torch`.

Results land in `benchmark/data/all_benchmark_data.csv`. Each row is
tagged with `kernel_name="megatron_cross_entropy"`.

### Step 2: Render plots

```bash
python benchmark/benchmarks_visualizer.py \
    --kernel-name megatron_cross_entropy --metric-name speed

python benchmark/benchmarks_visualizer.py \
    --kernel-name megatron_cross_entropy --metric-name memory
```

PNGs land in `benchmark/visualizations/`.

### Memory savings
<img width="1000" height="600"
alt="megatron_cross_entropy_memory_full_token_length"
src="https://github.com/user-attachments/assets/d08d7ee7-f6de-4e8e-a548-98b50da066ea"
/>

### Backward speedup
<img width="1000" height="600"
alt="megatron_cross_entropy_speed_backward_token_length"
src="https://github.com/user-attachments/assets/f8b524e2-658f-4206-b5a5-9d4d5e52e816"
/>

### Forward speedup
<img width="1000" height="600"
alt="megatron_cross_entropy_speed_forward_token_length"
src="https://github.com/user-attachments/assets/771115ff-143a-4c63-b112-c50bc0dddd84"
/>

### Backward+Forward speedup
<img width="1000" height="600"
alt="megatron_cross_entropy_speed_full_token_length"
src="https://github.com/user-attachments/assets/2152656c-c13e-4a6f-8b3d-577d6365a26e"
/>


## Example training scripts

Two end-to-end examples in `examples/megatron/`. Both train a tiny GPT
model with mock data for 5 iterations. They print which classes filled
which norm slots and which function backs each CE entry point, so you
can see Liger took over.

### Prerequisites

- A working Megatron-Core install (`pip install megatron-core`).
- `liger-kernel` installed.
- `psutil` (Megatron's async checkpoint worker pool uses it).
- At least 2 GPUs.

### Example 1: Mode 1 (monkey patch)

```bash
torchrun --nproc_per_node=2 \
    --master_addr=127.0.0.1 --master_port=29500 \
    examples/megatron/run_mode1_monkey_patch.py
```

One `apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True)`
call at the top. Everything else is stock Megatron code.

### Example 2: Mode 2 (hand-built spec)

```bash
torchrun --nproc_per_node=2 \
    --master_addr=127.0.0.1 --master_port=29500 \
    examples/megatron/run_mode2_hand_spec.py
```

No monkey patch. The script builds a `TransformerBlockSubmodules` with
`LigerMegatronRMSNorm` slots and subclasses `GPTModel` to use
`LigerMegatronCrossEntropy` in `compute_language_model_loss`.

### What you should see

For both:

- 5 lines of `[modeN] iter <i> loss=<float>` with the loss going down.
- A printed module tree with `LigerMegatronRMSNorm` in 5 of 5 norm slots
(4 per-layer + 1 block-level final).
- The CE symbols print as `liger_fused_vocab_parallel_cross_entropy` and
`liger_vocab_parallel_cross_entropy` (Mode 1) or
`LigerMegatronCrossEntropy` (Mode 2).
- `Successfully loaded the model` after a distributed checkpoint
round-trip.

If Apex or TransformerEngine isn't installed, you'll see harmless
warnings. Megatron falls back to the local backend, which is where Liger
plugs in.

## Testing Done

- Hardware Type: H100
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

---------

Co-authored-by: Vaibhav Jindal <[email protected]>
Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>
 2026-06-10 20:51:13 +00:00
Vaibhav Jindal e06ee2ce44 Add Megatron-Core integration for RMSNorm (#1254) 
## Summary

Adds first-class Megatron-Core support for Liger's Triton RMSNorm with
**two adoption modes**, both reaching **5/5 norm-slot coverage** on a
standard GPT-style model. No changes are required on Megatron-LM — this
PR works against upstream `megatron-core`.

**Note** - The kernel works for TP, and should ideally not break for SP,
but the behavior needs to verified end-to-end. Marked as a TODO in
`rms_norm.py`.

### Background — Megatron has *two* norm slots, not one

A GPT-style Megatron decoder has `2 × num_layers + 1` RMSNorm modules:

- `decoder.layers[i].input_layernorm` and
`decoder.layers[i].pre_mlp_layernorm` — populated from the per-layer
`TransformerLayerSubmodules` dataclass.
- `decoder.final_layernorm` — populated from a *separate*
`TransformerBlockSubmodules.layer_norm` field, which defaults to a
module-level global `LayerNormImpl` in `transformer_block.py` (TE / Apex
/ `WrappedTorchNorm`, chosen at import time).

To cover all norms, the integration must hit both code paths. Both modes
below do.

### Mode 1 — monkey-patch (one line, all GPT-family models, 5/5 norms)

```python
from liger_kernel.megatron import apply_liger_kernel_to_megatron

apply_liger_kernel_to_megatron(rms_norm=True)

# Then build your model normally — every RMSNorm slot is now Liger.
gpt_model = GPTModel(
    config=cfg,
    transformer_layer_spec=get_gpt_layer_local_spec(normalization="RMSNorm"),
    ...,
)
```

The call rewrites two sites:

1. `megatron.core.models.backends.LocalSpecProvider.layer_norm` →
returns `LigerMegatronRMSNorm` when `rms_norm=True`. Covers
`input_layernorm`, `pre_mlp_layernorm`, and (if `qk_layernorm=True`)
`q_layernorm` / `k_layernorm`.
2. `megatron.core.transformer.transformer_block.LayerNormImpl` →
replaced with a small adapter that routes to `LigerMegatronRMSNorm` when
`config.normalization == "RMSNorm"` and falls back to the original
`WrappedTorchNorm` otherwise. Covers `final_layernorm` without
disturbing LayerNorm users.

Both patches:
- Only displace the **`WrappedTorchNorm` fallback**. Users on TE or Apex
chose those deliberately — replacing them would undo TE's
`TELayerNormColumnParallelLinear` fusion or surprise Apex users.

Because both patches target the `BackendSpecProvider` / `LayerNormImpl`
abstractions — not individual model files — **a single call covers every
model that routes through Megatron's standard spec system** (GPT-3,
Llama-N, Mixtral, Nemotron, MLA, any `TransformerLayer`-based model). No
per-model code, in contrast to HF where every `modeling_<x>.py` needs
its own `apply_liger_kernel_to_<x>()`.

### Mode 2 — hand-assembled `TransformerBlockSubmodules` (slot-level
surgical control, 5/5 norms)

```python
from megatron.core.transformer.spec_utils import ModuleSpec
from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
from liger_kernel.megatron import LigerMegatronRMSNorm

per_layer = ModuleSpec(
    module=TransformerLayer,
    submodules=TransformerLayerSubmodules(
        input_layernorm=LigerMegatronRMSNorm,    # Liger
        self_attention=...,                      # TE here if you want
        pre_mlp_layernorm=LigerMegatronRMSNorm,  # Liger
        mlp=...,
    ),
)
block_spec = TransformerBlockSubmodules(
    layer_specs=[per_layer] * num_layers,
    layer_norm=LigerMegatronRMSNorm,             # ← covers final_layernorm
)
gpt_model = GPTModel(config=cfg, transformer_layer_spec=block_spec, ...)
```

Useful when mixing Liger with TransformerEngine per slot (e.g. "TE for
everything except norms"). Verbose but maximally explicit. Passing a
`TransformerBlockSubmodules` instead of a bare per-layer `ModuleSpec` is
the only way to reach `final_layernorm` without monkey-patching —
Megatron auto-promotes a per-layer spec via its module-level
`LayerNormImpl` global, which would otherwise leave the trailing norm on
PyTorch's `nn.RMSNorm`.

## What's added

| File | Purpose |
|---|---|
| `src/liger_kernel/megatron/rms_norm.py` | `LigerMegatronRMSNorm` —
conforms to Megatron's `LayerNormBuilder` protocol (`config,
hidden_size, eps`). Reads `config.normalization` (asserts `"RMSNorm"`),
`config.sequence_parallel`, `config.layernorm_zero_centered_gamma`. Sets
`weight.sequence_parallel` for the dist optimizer. `in_place=False`
hardcoded for Megatron's recompute / CUDA-graph safety. No
`megatron-core` import — duck-typed config. |
| `src/liger_kernel/megatron/monkey_patch.py` |
`apply_liger_kernel_to_megatron(rms_norm=True)` — patches both
`LocalSpecProvider.layer_norm` and `transformer_block.LayerNormImpl`.
Idempotent. |
| `src/liger_kernel/megatron/__init__.py` | Re-exports
`LigerMegatronRMSNorm`, `apply_liger_kernel_to_megatron`. |
| `examples/megatron/run_mode1_monkey_patch.py` | End-to-end runnable
script demonstrating Mode 1. |
| `examples/megatron/run_mode2_hand_spec.py` | End-to-end runnable
script demonstrating Mode 2. |
| `examples/megatron/README.md` | Prerequisites + run instructions +
expected output. |
| `test/megatron/test_rms_norm.py` | 10 unit tests against
`torch.nn.RMSNorm` reference. |

## How to reproduce

Two self-contained scripts live in this PR under
[`examples/megatron/`](examples/megatron). Each is an end-to-end
training run (TP=2, mock GPT data, 5 iterations, distributed-checkpoint
round-trip) that prints the full module tree and the resolved norm class
at every slot so reviewers can verify Liger took over. Both adapt
Megatron's own `examples/run_simple_mcore_train_loop.py`.

```bash
# Mode 1 — apply_liger_kernel_to_megatron(rms_norm=True)
torchrun --nproc_per_node=2 --master_addr=127.0.0.1 --master_port=29500 \
    examples/megatron/run_mode1_monkey_patch.py

# Mode 2 — hand-assembled TransformerBlockSubmodules
torchrun --nproc_per_node=2 --master_addr=127.0.0.1 --master_port=29500 \
    examples/megatron/run_mode2_hand_spec.py
```

The Liger-specific lines in each script are flagged with `↓↓ Mode N ↓↓`
/ `↑↑ ↑↑` comment fences so reviewers can scan straight to the
integration code without reading the boilerplate.

See [`examples/megatron/README.md`](examples/megatron/README.md) for
prerequisites and expected output (5 lines of loss + 5 lines of
`LigerMegatronRMSNorm` slots + `Successfully loaded the model`).

## Megatron feature compatibility

**Verified:** TP (tested at TP=2), distributed checkpointing (save/load
round-trip in the harness), fp32 + bf16 (unit-tested), recompute paths
(`in_place=False` keeps activation references valid), `qk_layernorm`
slots, and `final_layernorm` (covered in both modes — mode 1 via
`transformer_block.LayerNormImpl`, mode 2 via
`TransformerBlockSubmodules.layer_norm`).

**Should work, not yet E2E-verified:** SP (wrapper sets
`weight.sequence_parallel`; notably the local baseline
`WrappedTorchNorm` asserts SP off, so this is a strict improvement), CP,
`use_distributed_optimizer=True`, CUDA-graph capture. Positive tests for
these are tracked as follow-up.

**Out of scope:** FP8 / FP4 (TE owns that path), the TE backend, and the
Apex backend — the patches only displace the `WrappedTorchNorm`
fallback. Users on TE or Apex keep their existing norms.
## Known issues / follow-ups

1. **SP / CP / dist-opt / CUDA-graph not end-to-end verified.** Should
work by construction; need positive tests before production claims.
2. **Pre-existing torch 2.11 compatibility issue in
`liger_kernel.ops.rms_norm`.** It references
`torch.distributed.tensor.DTensor` without forcing the submodule import.
Worked around in our wrapper with an explicit `import
torch.distributed.tensor` — proper fix is a one-line change to
`ops/rms_norm.py`, intentionally not bundled here.
3. **Performance not measured.** This PR establishes correctness,
integration shape, and full norm coverage. Throughput / memory
benchmarks on a realistic hidden size (Llama-7B class, hidden=4096+) are
the natural next step.
4. **Convergence tests need to be added**
## Testing Done

- Hardware Type: NVIDIA H100 80GB (2×)
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
 2026-06-09 20:25:03 +00:00
Tcc0403 02d27648fd Replace all torch_dtype with dtype (#881) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Fix #880 
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

---------

Signed-off-by: Tcc0403 <[email protected]>
 2025-09-16 16:11:23 -07:00
Chiwan ParkandShao Tang aabda78ccd Fix empty Medusa head tensors (#698) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->

This PR fixes #309.

The original implementation injects Medusa heads to FSDP-unwrapped model
variables, which are not accessed from the state_dict (FSDP-wrapped). To
fix this bug, I've changed the method to inject Medusa heads with model
loader, and the method to access the state_dict through APIs from
accelerate.

Minor updates:
- Update dependencies to the latest.
- Fix scripts to run Medusa example.

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

![Screenshot_20250504_021259](https://github.com/user-attachments/assets/1fe55300-87fc-4ff6-b4a7-489746d274ee)

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: A100-80G-NVLINK 8 cards
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

Co-authored-by: Shao Tang <[email protected]>
 2025-05-14 20:43:01 -07:00
KareemandShivam Sahni af1fe54551 fix zip bug (#702) 
## Summary
fixes #701. `strict=False` is default either way
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: NVIDIA GeForce MX230
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

Co-authored-by: Shivam Sahni <[email protected]>
 2025-05-05 12:20:38 -07:00
Shivam SahniAustin LiuShao Tanghebiao064
 		1cc9c5083a Revert "Improve Hugging Face SFT Script (#539)" and Fix TVD Test for Intel #580 (#578) 
This reverts commit 5476b20d3b.
The PR has an undefined `prepare_dataset` function which causes
checkstyle to fail
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

---------

Signed-off-by: Austin Liu <[email protected]>
Co-authored-by: Austin Liu <[email protected]>
Co-authored-by: Shao Tang <[email protected]>
Co-authored-by: hebiao064 <[email protected]>
 2025-02-21 15:02:19 -08:00
Parag EkboteYun DaiShao Tang
 		5476b20d3b Improve Hugging Face SFT Script (#539) 
Fixes #487 

I've chosen to remove the deprecated parameter as previously mentioned
in the issue. The sequence length for the training dataset can be
specified using the HF datasets library as mentioned
[here](https://huggingface.co/docs/trl/main/en/sft_trainer#datasets).

Please let me know if any further rectification is required and I will
make the necessary changes.

cc: @Tcc0403

---------

Co-authored-by: Yun Dai <[email protected]>
Co-authored-by: Shao Tang <[email protected]>
 2025-02-20 17:11:20 -08:00
Shivam Sahni bb84c0d2cf fix qwen2vl and mllama test to pass failing tests (#571) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
1. Fix Qwen2VL mrope tests. Transformers removed BC with
Qwen2VLRotaryEmbedding which is why we need this change
2. Fix mllama tests to not use cross entropy. Using Cross Entropy causes
logits mismatch which is expected as we do inplace operation in cross
entropy kernel.
3. Runs make checkstyle --- there seems to be a bug in checkstyle
workflow where unformatted checkstyle seems to pass tests. Changed the
makefile to throw error in such cases.
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
 2025-02-20 11:54:47 -08:00
Shivam Sahni 77ff1a9b7d Move the checkstyle to [Ruff](https://docs.astral.sh/ruff/) (#483) 2024-12-22 20:21:17 -08:00
Byron Hsu 3a0cbdff4b Fix liger orpo trainer import error (#459) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

To fix https://github.com/linkedin/Liger-Kernel/issues/458

1. Move orpo trainer to trainer/
2. Remove trl from `[dev]` now so the test represents the actual user
environment.
3. Disable Qwen VL test for transformers>=4.47.0 for now and create
issue https://github.com/linkedin/Liger-Kernel/issues/461
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
 2024-12-10 01:29:20 -08:00
Shivam Sahni 6cb001803b Add ORPO Trainer + support HF metrics directly from chunked loss functions + fixes to avoid torch compile recompilations (#429) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
This PR adds support for the following:
1. LigerORPOTrainer: a wrapper on top of [HuggingFace ORPO
Trainer](https://github.com/huggingface/trl/blob/main/trl/trainer/orpo_trainer.py)
to use LigerORPOLoss module.
2. We also provide an example for using LigerORPOTrainer in
`examples/alignment/run_orpo.py`
3. Change FusedLinearPreference base class' forward function to return
additional metrics to align our implementation with HF ORPO Trainer
4. Additional refactor to avoid torch compile recompilations --
accumulate_chunk function now calls accumulate_helper which is torch
compiled solely and input_chunk/target_chunk/target dimension 1 (seq
len) is explicitly marked as dynamic to avoid recompilations

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
 2024-12-06 10:09:11 -08:00
Golam RabbaniandShao Tang 7e3683e23f Xpu support (#407) 
## Summary
Replica of #396 
Adds xpu support so all tests, benchmarks etc. run on XPUs or Intel
GPUs.

## Details
infer_device() function is moved to a separate file and in any file
where previously "cuda" was needed, infer_device is imported and "cuda"
is replaced with return value of a call to infer_device()

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
A100 80GB PCIe, RTX 3060, Intel Data Center GPU Max 1550
<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence

---------

Co-authored-by: Shao Tang <[email protected]>
 2024-11-22 20:03:06 -08:00
Tyler Romero 81d98ea895 Add script to reproducibly run examples on Modal (#397) 
## Summary
Add a script that will allow users to run the examples on Modal, even if
they don't have access to a 4xA100 80Gb node at work.

Also fix hparams for tuning qwen2-vl on a 4xA100 80Gb node

## Testing Done
- Hardware Type:
Modal launched from macbook, script runs on modal on 4x80Gb A100s
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
 2024-11-19 14:31:47 -08:00
Tyler RomeroandByron Hsu 9b24c61b90 Qwen2-VL Training Example w/ Liger (#389) 
## Summary
Example demonstrating how to use `SFTTrainer` to finetune Qwen2-VL on a
multimodal dataset using Liger-Kernel.

## Testing Done
- Hardware Type: 2x RTX 4090
Ran on two RTX 4090s, however I hit CUDA OOMs after about 10 steps (even
with per-device batch_size=2, adafactor, shorter max_seq_length...)

Would love some help from a LinkedIn employee with access to A100s/H100s
at work to help verify.

- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence

---------

Co-authored-by: Byron Hsu <[email protected]>
 2024-11-18 21:44:37 -08:00
Chiwan Park 659d7d7856 Fix incorrect training of first and last Medusa heads (#325) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Currently, there are two errors on Medusa training examples:

1. When we use Liger Kernel, the first head (`model.medusa_head[0]`) is
not trained.
2. When we don't use Liger Kernel, the logits of the last head
(`medusa_logits[-1]`) is ignored.

This PR fixes these errors.

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: A100 80GB 8 GPUs
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
 2024-10-31 16:24:06 -07:00
Chiwan Park 8e2bd26f67 Fix sharing a ResBlock layer for each head in Medusa example (#269) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
There is a bug of incorrect weight sharing between layers for each
Medusa head. Since `nn.Module` is Python reference, the original source
code creates a list containing references to the same weights. This PR
fixes the bug.

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: A100-80G-PCIe
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
 2024-09-25 08:34:57 -07:00
Steven Shimizu b42a27bd70 Added HF use-case benchmark script (#223) 
## Summary
- Added Hugging Face training benchmarking script used for tech report
- Writes files to
`/results/${MODEL_TYPE}_use_liger_${USE_LIGER}_batch_size_${BATCH_SIZE}_rep_${i}.log`

## Testing Done
- Ran benchmarking script

- Hardware Type: A100
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
 2024-09-06 14:16:41 -07:00
Qingquan Song 3c71ea4ffa Switch Lightning 1 GPU example to Qwen2 0.5B instruct model with 1024 max seq length (#193) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Switch to Qwen2 0.5B instruct model with 1024 max seq length for L40
training (control max memory to < 48G)

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
 2024-09-03 11:03:58 -07:00
Jason (Siyu) Zhuandjaszhu 1129508967 Add gemma 7b it benchmark (#166) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Add gemma 7b it benchmark
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
Add gemma 7b it benchmark
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
NA
<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence

---------

Co-authored-by: jaszhu <[email protected]>
 2024-08-29 22:19:45 +00:00
Jason (Siyu) Zhuandjaszhu 9c7b9d02a1 Add missing tf_keras to req.txt (#161) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Add missing tf_keras to req.txt 
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
Add missing tf_keras to req.txt , otherwise pip install -r
requirement.txt and then run the script will result in

```
ValueError: Your currently installed version of Keras is Keras 3, but this is not yet supported in Transformers. Please install the backwards-compatible tf-keras package with `pip install tf-keras`.
```

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
Done
<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence

Co-authored-by: jaszhu <[email protected]>
 2024-08-29 09:13:05 -07:00
Qingquan SongShao TangByron Hsu
 		8089f1e871 Add gemma lightning example for single L40 GPU (#120) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->

Add gemma lightning example for single L40 GPU

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

---------

Co-authored-by: Shao Tang <[email protected]>
Co-authored-by: Byron Hsu <[email protected]>
 2024-08-28 10:45:17 -07:00
Jason (Siyu) Zhuandjaszhu cec2e99806 Add code to save medusa heads and model (#130) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Add code to save medusa heads and model

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
Add code to save medusa heads and model, when medusa and
medusa_only_head are both enabled, we save only heads. Otherwise the
whole model will be saved.

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
Head module is saved:

```
medusa_lm_head.safetensors
```

Whole model is saved:
```

config.json                 model-00001-of-00005.safetensors  model-00004-of-00005.safetensors  special_tokens_map.json  training_args.bin
generation_config.json      model-00002-of-00005.safetensors  model-00005-of-00005.safetensors  tokenizer.json
  model-00003-of-00005.safetensors  model.safetensors.index.json      tokenizer_config.json
```

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence

---------

Co-authored-by: jaszhu <[email protected]>
 2024-08-27 16:46:16 -07:00
Steven Shimizu e3fd9eeab9 Adding AutoLigerKernelForCausalLM (#115) 
## Summary
- Adding AutoLigerKernelForCausalLM to automatically apply relevant
liger kernels to corresponding model to address #70
- Args passed to auto model that are part of
`apply_liger_kernel_to_{model_type}` supported args are passed through
- Refactored the monkey_patch.py so the mapping and patching functions
are all in the same file so it's easier to not miss updating. Also added
test to ensure the patching functions match the auto model mapping.

Example:
```python
model = AutoLigerKernelForCausalLM.from_pretrained(
        custom_args.model_name,
        trust_remote_code=True,
        use_cache=False,
        torch_dtype=torch.bfloat16,
        cross_entropy=True,
        fused_linear_cross_entropy=False,
    )

> INFO:liger_kernel.transformers.monkey_patch:Applying Liger kernels for model type: llama with kwargs: {'cross_entropy': True, 'fused_linear_cross_entropy': False}
```

## Testing Done
- Run huggingface example

```python
assert modeling_llama.LlamaRMSNorm != LigerRMSNorm
model = AutoLigerKernelForCausalLM.from_pretrained(
    custom_args.model_name,
    trust_remote_code=True,
    use_cache=False,
    torch_dtype=torch.bfloat16,
    # These args will get passed to the appropriate apply_liger_kernel_to_* function
    # to override the default settings
    cross_entropy=True,
    fused_linear_cross_entropy=False,
)
assert modeling_llama.LlamaRMSNorm == LigerRMSNorm
```

- Hardware Type: A100
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence


```
$ make checkstyle
flake8 .; flake8_status=$?; \
isort .; isort_status=$?; \
black .; black_status=$?; \
if [ $flake8_status -ne 0 ] || [ $isort_status -ne 0 ] || [ $black_status -ne 0 ]; then \
        exit 1; \
fi
Skipped 3 files
All done! ✨ 🍰 ✨
54 files left unchanged.

$ make test
pytest --disable-warnings test/ --ignore=test/convergence
==================================================================================================================================== test session starts ====================================================================================================================================
platform linux -- Python 3.10.14, pytest-7.1.2, pluggy-1.0.0
rootdir: /home/jobuser/Liger-Kernel
plugins: lipy-config-base-32.0.27, lipy-fabric-35.2.19, lipy-test-8.0.66, datadir-1.3.1
collected 117 items                                                                                                                                                                                                                                                                         

test/transformers/test_auto_model.py .                                                                                                                                                                                                                                                [  0%]
test/transformers/test_cross_entropy.py ..........................................................                                                                                                                                                                                    [ 50%]
test/transformers/test_fused_linear_cross_entropy.py ......                                                                                                                                                                                                                           [ 55%]
test/transformers/test_geglu.py ........                                                                                                                                                                                                                                              [ 62%]
test/transformers/test_monkey_patch.py .....                                                                                                                                                                                                                                          [ 66%]
test/transformers/test_rms_norm.py ................                                                                                                                                                                                                                                   [ 80%]
test/transformers/test_rope.py ............                                                                                                                                                                                                                                           [ 90%]
test/transformers/test_swiglu.py ........                                                                                                                                                                                                                                             [ 97%]
test/transformers/test_trainer_integration.py .                                                                                                                                                                                                                                       [ 98%]
test/triton/test_triton_monkey_patch.py ..                                                                                                                                                                                                                                            [100%]

=================================================================================================================================== 117 passed in 50.75s ====================================================================================================================================

$ make test-convergence
HF_DATASETS_OFFLINE=1 pytest --disable-warnings test/convergence
==================================================================================================================================== test session starts ====================================================================================================================================
platform linux -- Python 3.10.14, pytest-7.1.2, pluggy-1.0.0
rootdir: /home/jobuser/Liger-Kernel
plugins: lipy-config-base-32.0.27, lipy-fabric-35.2.19, lipy-test-8.0.66, datadir-1.3.1
collected 8 items                                                                                                                                                                                                                                                                           

test/convergence/test_mini_models.py ......                                                                                                                                                                                                                                           [ 75%]
test/convergence/test_mini_models_no_logits.py ..                                                                                                                                                                                                                                     [100%]

==================================================================================================================================== 8 passed in 47.97s =====================================================================================================================================

```
 2024-08-27 14:07:41 -07:00
Tyler Romero a272dad3f8 [Bug fix] Efficiency callback missing dataclass entries (#116) 
## Summary
Fix EfficiencyCallback missing entries for the Memory dataclass

## Testing Done
Ran a training run using the callback.
 2024-08-26 20:02:18 -07:00
Jason (Siyu) Zhuandjaszhu b48df3ccfd Explicitly add dependencies in req.txt for medusa example (#110) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Explicitly add dependencies in req.txt for medusa example to ensure the
example can be run without issue

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
Explicitly add dependencies in req.txt for medusa example to ensure the
example can be run without issue, fix this issue
https://github.com/linkedin/Liger-Kernel/issues/97

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
Run training successfully with pip install -r requirement.txt

```
  return fn(*args, **kwargs)
{'loss': 75.1831, 'grad_norm': 9.36429500579834, 'learning_rate': 9.931806517013612e-06, 'epoch': 1.0, 'num_input_tokens_seen': 65536, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18801.85, 'avg_tokens_per_second': 18801.85}
{'loss': 74.5029, 'grad_norm': 8.809615135192871, 'learning_rate': 9.729086208503174e-06, 'epoch': 1.5, 'num_input_tokens_seen': 98304, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18835.9, 'avg_tokens_per_second': 18818.86}
{'loss': 75.81, 'grad_norm': 9.360036849975586, 'learning_rate': 9.397368756032445e-06, 'epoch': 2.0, 'num_input_tokens_seen': 131072, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18920.33, 'avg_tokens_per_second': 18852.56}
{'loss': 71.8666, 'grad_norm': 8.996702194213867, 'learning_rate': 8.94570254698197e-06, 'epoch': 2.5, 'num_input_tokens_seen': 163840, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18781.19, 'avg_tokens_per_second': 18834.67}
{'loss': 75.8349, 'grad_norm': 8.850849151611328, 'learning_rate': 8.386407858128707e-06, 'epoch': 3.0, 'num_input_tokens_seen': 196608, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18852.58, 'avg_tokens_per_second': 18838.25}
{'loss': 75.7597, 'grad_norm': 8.883410453796387, 'learning_rate': 7.734740790612137e-06, 'epoch': 3.5, 'num_input_tokens_seen': 229376, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18763.97, 'avg_tokens_per_second': 18825.83}
{'loss': 74.5994, 'grad_norm': 8.877114295959473, 'learning_rate': 7.008477123264849e-06, 'epoch': 4.0, 'num_input_tokens_seen': 262144, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18728.62, 'avg_tokens_per_second': 18811.88}
{'loss': 74.2385, 'grad_norm': 9.293025970458984, 'learning_rate': 6.227427435703997e-06, 'epoch': 4.5, 'num_input_tokens_seen': 294912, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18774.71, 'avg_tokens_per_second': 18807.23}
{'loss': 74.3044, 'grad_norm': 8.572715759277344, 'learning_rate': 5.412896727361663e-06, 'epoch': 5.0, 'num_input_tokens_seen': 327680, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18790.52, 'avg_tokens_per_second': 18805.37}
{'loss': 76.4343, 'grad_norm': 9.21961784362793, 'learning_rate': 4.587103272638339e-06, 'epoch': 5.5, 'num_input_tokens_seen': 360448, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18792.55, 'avg_tokens_per_second': 18804.08}
{'loss': 73.8025, 'grad_norm': 9.030799865722656, 'learning_rate': 3.7725725642960047e-06, 'epoch': 6.0, 'num_input_tokens_seen': 393216, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18790.54, 'avg_tokens_per_second': 18802.85}
{'loss': 75.0139, 'grad_norm': 9.177617073059082, 'learning_rate': 2.991522876735154e-06, 'epoch': 6.5, 'num_input_tokens_seen': 425984, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18761.38, 'avg_tokens_per_second': 18799.39}
{'loss': 73.9404, 'grad_norm': 9.219317436218262, 'learning_rate': 2.265259209387867e-06, 'epoch': 7.0, 'num_input_tokens_seen': 458752, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18751.01, 'avg_tokens_per_second': 18795.66}
{'loss': 75.2985, 'grad_norm': 9.104233741760254, 'learning_rate': 1.6135921418712959e-06, 'epoch': 7.5, 'num_input_tokens_seen': 491520, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18741.49, 'avg_tokens_per_second': 18791.78}
{'loss': 73.0278, 'grad_norm': 9.329939842224121, 'learning_rate': 1.0542974530180327e-06, 'epoch': 8.0, 'num_input_tokens_seen': 524288, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18768.68, 'avg_tokens_per_second': 18790.24}
{'loss': 72.3754, 'grad_norm': 9.397867202758789, 'learning_rate': 6.026312439675553e-07, 'epoch': 8.5, 'num_input_tokens_seen': 557056, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18729.28, 'avg_tokens_per_second': 18786.42}
{'loss': 76.0987, 'grad_norm': 8.823853492736816, 'learning_rate': 2.7091379149682683e-07, 'epoch': 9.0, 'num_input_tokens_seen': 589824, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18703.4, 'avg_tokens_per_second': 18781.51}
{'loss': 73.5105, 'grad_norm': 9.082978248596191, 'learning_rate': 6.819348298638839e-08, 'epoch': 9.5, 'num_input_tokens_seen': 622592, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18727.9, 'avg_tokens_per_second': 18778.53}
{'loss': 73.9587, 'grad_norm': 9.229247093200684, 'learning_rate': 0.0, 'epoch': 10.0, 'num_input_tokens_seen': 655360, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18721.68, 'avg_tokens_per_second': 18775.53}
{'train_runtime': 40.0755, 'train_samples_per_second': 12.476, 'train_steps_per_second': 0.499, 'train_loss': 74.53130264282227, 'epoch': 10.0, 'num_input_tokens_seen': 655360, 'step_peak_memory_allocated_MB': 26279.24, 'total_peak_memory_allocated_MB': 26279.24, 'step_peak_memory_reserved_MB': 47172.0, 'total_peak_memory_reserved_MB': 47172.0, 'step_tokens_per_second': 18721.68, 'avg_tokens_per_second': 18775.53}
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [00:40<00:00,  2.00s/it]


```
<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence

Co-authored-by: jaszhu <[email protected]>
 2024-08-26 12:30:19 -07:00
Shivam Sahniandshisahni 03fd60320e Add Qwen benchmarks (#108) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
This PR adds qwen2 7b benchmarks
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: A100-80G-PCIe
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence

---------

Co-authored-by: shisahni <[email protected]>
 2024-08-26 12:06:00 -07:00
Qingquan Song e8a58aa64a Remove duplicate images (#107) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Same problem still appears after #78 


<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
 2024-08-26 11:49:05 -07:00
Qingquan Song 10666fcf22 Fix example bugs (#88) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Fix example bugs: 
1) lightning data loading
original loaded data has format below so should be flattened to extract
nested features.
```
DatasetDict({
    train: Dataset({
        features: ['train'],
        num_rows: 99842
    })
})
```

2) medusa sklearn learn dependency: public scikit-learn should use
full-name

misc: readme clarify test A100 clip HBM size, add deepspeed dependency
in lightning requirements


Thanks to https://github.com/t-vi who found out the issues. 
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
 2024-08-25 20:22:47 -07:00
Peter Chng 45cfc510a7 Remove unused images in examples/medusa/docs/images/ (#78) 
## Summary
- There are four images I believe are unused in
`examples/medusa/docs/images/`
- These images were replaced by updated benchmark results in this PR:
https://github.com/linkedin/Liger-Kernel/pull/59/
- Furthermore these images have filenames that only differ by one letter
case, causing issues on cloning the repo on case-insensitive filesystems
like those used on Windows and MacOS. (Linux not affected) You would see
an error like this:

```
git clone [email protected]:pchng/Liger-Kernel.git
Cloning into 'Liger-Kernel'...
remote: Enumerating objects: 735, done.
remote: Counting objects: 100% (342/342), done.
remote: Compressing objects: 100% (213/213), done.
remote: Total 735 (delta 211), reused 213 (delta 121), pack-reused 393 (from 1)
Receiving objects: 100% (735/735), 8.18 MiB | 11.39 MiB/s, done.
Resolving deltas: 100% (306/306), done.
warning: the following paths have collided (e.g. case-sensitive paths
on a case-insensitive filesystem) and only one from the same
colliding group is in the working tree:

  'examples/medusa/docs/images/Memory_Stage1_num_head_3.png'
  'examples/medusa/docs/images/memory_Stage1_num_head_3.png'
  'examples/medusa/docs/images/Memory_Stage1_num_head_5.png'
  'examples/medusa/docs/images/memory_Stage1_num_head_5.png'
  'examples/medusa/docs/images/Memory_Stage2_num_head_3.png'
  'examples/medusa/docs/images/memory_Stage2_num_head_3.png'
  'examples/medusa/docs/images/Memory_Stage2_num_head_5.png'
  'examples/medusa/docs/images/memory_Stage2_num_head_5.png'
```

## Testing Done
- Ran `grep -r "memory_Stage" ./*`, didn't find references to the
deleted image files.
- The relevant `README.md` file still looks OK after removing these
unused images:
https://github.com/pchng/Liger-Kernel/blob/50345cfb881388ce4de44210e08bd8c43cf1ed04/examples/medusa/README.md
 2024-08-24 21:18:33 -07:00
Jason (Siyu) Zhuandjaszhu 5fb483174d Update benchmark result of Medusa for batch size = 6 setup (#59) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->

Update benchmark result of Medusa for batch size = 6 setup
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
Update benchmark result of Medusa for batch size = 6 setup, and link the
plot for throughput and memory

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
Done
<!-- 
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- [x ] run `make test` to ensure correctness
- [x ] run `make checkstyle` to ensure code style
- [x ] run `make test-convergence` to ensure convergence

Co-authored-by: jaszhu <[email protected]>
 2024-08-23 20:27:21 -07:00
Steven Shimizu 2c43eba98b Updated test tolerances for H100 (#55) 
## Summary
- Updated test tolerances to pass on H100

## Testing Done

- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence

```
$ make checkstyle
flake8 . --exclude=.venv; flake8_status=$?; \
isort .; isort_status=$?; \
black .; black_status=$?; \
if [ $flake8_status -ne 0 ] || [ $isort_status -ne 0 ] || [ $black_status -ne 0 ]; then \
        exit 1; \
fi
Skipped 2 files
All done! ✨ 🍰 ✨
51 files left unchanged.

$ make test
pytest --disable-warnings test/ --ignore=test/convergence
================================================================================ test session starts ================================================================================
platform linux -- Python 3.10.14, pytest-7.1.2, pluggy-1.0.0
rootdir: /home/jobuser/Liger-Kernel
plugins: lipy-config-base-32.0.27, lipy-fabric-35.2.19, lipy-test-8.0.66, datadir-1.3.1
collected 114 items                                                                                                                                                                 

test/transformers/test_cross_entropy.py ..........................................................                                                                            [ 50%]
test/transformers/test_fused_linear_cross_entropy.py ......                                                                                                                   [ 56%]
test/transformers/test_geglu.py ........                                                                                                                                      [ 63%]
test/transformers/test_rms_norm.py ................                                                                                                                           [ 77%]
test/transformers/test_rope.py ............                                                                                                                                   [ 87%]
test/transformers/test_swiglu.py ........                                                                                                                                     [ 94%]
test/transformers/test_trainer_integration.py ...                                                                                                                             [ 97%]
test/transformers/test_transformers_monkey_patch.py .                                                                                                                         [ 98%]
test/triton/test_triton_monkey_patch.py ..                                                                                                                                    [100%]

=============================================================================== 114 passed in 35.25s ================================================================================
(.venv) jobuser [ ~/Liger-Kernel ]$ make test-convergence
HF_DATASETS_OFFLINE=1 pytest --disable-warnings test/convergence
================================================================================ test session starts ================================================================================
platform linux -- Python 3.10.14, pytest-7.1.2, pluggy-1.0.0
rootdir: /home/jobuser/Liger-Kernel
plugins: lipy-config-base-32.0.27, lipy-fabric-35.2.19, lipy-test-8.0.66, datadir-1.3.1
collected 6 items                                                                                                                                                                   

test/convergence/test_mini_models.py ....                                                                                                                                     [ 66%]
test/convergence/test_mini_models_no_logits.py ..                                                                                                                             [100%]

================================================================================ 6 passed in 23.04s =================================================================================
```
 2024-08-23 10:42:12 -07:00
Jason (Siyu) Zhu b418557a5a Replace model / data with public HF path, update readme (#53) 2024-08-22 22:03:31 -07:00
ByronHsu ce7735d542 format 2024-08-19 00:04:07 +00:00
Yun Dai 9cb9dc0b41 Add deepspeed to lightning example (#36) 2024-08-16 15:31:01 -07:00
Zain Merchant 37356e9a36 remove mfu and non used methods (#35) 
## Summary
Remove mfu and non used methods

## Testing Done
Ran hugging face training job
 2024-08-16 14:37:24 -07:00
Byron Hsu 8d88214133 use correct fsdp act ckpt & redo benchmark (#32) 2024-08-15 23:40:49 -07:00
Yun Dai afe4d8c224 Add lightning trainer and HF trainer fine-tuning example (#17) 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
This PR adds example of using liger kernel for finetuning with lightning
trainer.

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
Worked on on-prem but need some lightning studio credits to test the
deps and configs outside :-0
 2024-08-15 19:01:13 -07:00
Jason (Siyu) ZhujaszhuByronHsu
 		1ed2a1f180 Add medusa patching code and example job with memory efficient Liger Kernel (#11) 
# Summary

Add medusa patching code and example job with memory efficient Liger
Kernel

Medusa is a simple framework that democratizes the acceleration
techniques for LLM generation with multiple decoding heads.
[[repo](https://arxiv.org/abs/2401.10774)],
[[paper](https://arxiv.org/abs/2401.10774)]

During training, Medusa requires adding \(k\) decoding heads to the
hidden states right before the regular LM head \(h_t\). The \(k\)-th
head is used to predict the token in the \((t + k + 1)\)-th position of
the next tokens (the original language model head is used to predict the
\((t + 1)\)-th position).

The Liger fused CE kernel is highly effective in this scenario,
eliminating the need to materialize logits for each head, which usually
consumes a large volume of memory due to the extensive vocabulary size
(e.g., for LLaMA-3, the vocabulary size is 128k). The introduction of
multiple heads can easily lead to OOM (Out of Memory) issues. However,
thanks to the efficient Liger fused CE, which calculates the gradient in
place and doesn't materialize the logits, we have observed very
effective results. This efficiency opens up more opportunities for
multi-token prediction research and development.


# List of changes

* Add self-contained medusa training flow with corresponding fsdp config
* Add README.md with benchmark result
* Add patching code to apply multi-head to backbone model


### TODO ([JasonZhu1313](https://github.com/JasonZhu1313))
* [P0] Clean up the code and remove hard-code coefficient in patching
code
* [P1] parallelize the llm head inference code

# Testing done
Convergence test

Liger kernel is very effective in memory usage for medusa training: 

![memory_Stage1_num_head_3](https://github.com/user-attachments/assets/d80f5e6d-87a2-4480-bbd5-eae83e1e15a1)

![memory_Stage1_num_head_5](https://github.com/user-attachments/assets/744bd04b-524e-42b9-b3e4-da1e9a48352c)

![memory_Stage2_num_head_3](https://github.com/user-attachments/assets/0ae916eb-6396-4c75-81e3-1b40da2c054c)

![memory_Stage2_num_head_5](https://github.com/user-attachments/assets/b862d614-8dd5-47d4-b8f4-ed90531480ae)

---------

Co-authored-by: jaszhu <[email protected]>
Co-authored-by: ByronHsu <[email protected]>
 2024-08-15 10:01:25 -07:00