mirror of
https://github.com/axolotl-ai-cloud/axolotl.git
synced 2026-10-12 02:11:20 +00:00
Wing Lian
1855055ed1
* bump deps * bump hub too * bump hf-xet also * vendor deprecated PRM trainer * fix mistral tokenizer api change * remove override patch for fa2 * fix EP handling * Add contract tests for the expert-parallel seam with transformers 5.18 Pin every attribute, kernel signature, registry hook, and plan style that the expert-parallel integration borrows from transformers, on real decorated Qwen3-MoE and Mixtral experts modules, so the next dependency bump fails in a named CPU test. Cover the sentinel routing contract for grouped_mm and batched_mm with _is_expert_parallel set, and the IndexError batched_mm raises without it. * Localize PEFT's DTensor LoRA factors in the native NVFP4 TP path PEFT 0.21.1 wraps the sharded factor as a DTensor and installs transformers' tensor-parallel forward on it when the base layer carries a TP plan. The native NVFP4 TP LoRA forward computes with local factors and does its own reductions, so unwrap the factors and drop that forward before validating them. * Drop the Triton selective_log_softmax override in favour of trl's kernel trl 1.14.1 ships a fused log-softmax and entropy kernel that matches or beats the axolotl Triton version at every shape the RL trainers use, and takes the temperature and row_mask arguments the trainers now pass. * Own EP gradient clipping in the trainer instead of patching accelerate transformers 5.18 routes clipping through a per-mesh path whenever the gradients do not share one device mesh, which pure EP never does: experts live on their own mesh and the dense weights on the EP mesh. The per-mesh norm never sums the experts held by other ranks, so the global norm is too small and the clipped update lands 5-7% too large. The trainer's _clip_grad_norm/_get_grad_norm hooks now compute the ownership-filtered EP norm directly when EP is enabled, so no accelerate clip_grad_norm_ patch is needed. * Drop the removed monkeypatch.trainer.utils module from the API reference
35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
"""Fixtures shared by the expert-parallel integration tests."""
|
|
|
|
import pytest
|
|
|
|
|
|
@pytest.fixture
|
|
def fake_ep_sharder(monkeypatch):
|
|
"""Run the real ``shard_expert_weights`` as one rank of a faked two-rank EP group.
|
|
|
|
The collective scatter is replaced by a CPU slice; everything else the sharder
|
|
does to the module (counts, offsets, flags, DDP ignores) is the production code.
|
|
"""
|
|
from axolotl.integrations.expert_parallel import shard
|
|
|
|
def run(model, rank, world_size=2):
|
|
monkeypatch.setattr(shard.dist, "get_world_size", lambda group=None: world_size)
|
|
monkeypatch.setattr(shard.dist, "get_rank", lambda group=None: rank)
|
|
monkeypatch.setattr(
|
|
shard.dist,
|
|
"all_gather_object",
|
|
lambda ranks, value: ranks.__setitem__(
|
|
slice(None), list(range(world_size))
|
|
),
|
|
)
|
|
|
|
def scatter_on_cpu(module, name, count, ranks):
|
|
shard._replace_with_slice(
|
|
module, name, ranks[rank] * count, (ranks[rank] + 1) * count
|
|
)
|
|
|
|
monkeypatch.setattr(shard, "_scatter_expert_from_rank0", scatter_on_cpu)
|
|
return shard.shard_expert_weights(model, object())
|
|
|
|
return run
|