mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2026-07-28 19:56:13 +08:00
[v1] refactor registry plugin structure and params (#10641)
This commit is contained in:
@@ -30,13 +30,13 @@ def _apply_kernel(rank) -> None:
|
||||
if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"):
|
||||
del sys.modules[k]
|
||||
|
||||
from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_default_kernels
|
||||
from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained("llamafactory/tiny-random-qwen3")
|
||||
original_rmsnorm_forward = model.model.layers[0].input_layernorm.forward
|
||||
original_swiglu_forward = model.model.layers[0].mlp.forward
|
||||
|
||||
model = apply_default_kernels(model=model, include_kernels="npu_fused_rmsnorm")
|
||||
model = apply_kernels(model=model, config={"name": "npu_fused_rmsnorm"})
|
||||
|
||||
assert model.model.layers[0].input_layernorm.forward.__func__ is not original_rmsnorm_forward.__func__
|
||||
assert model.model.layers[0].mlp.forward.__func__ is original_swiglu_forward.__func__
|
||||
@@ -53,13 +53,13 @@ def _apply_all_kernels(rank) -> None:
|
||||
if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"):
|
||||
del sys.modules[k]
|
||||
|
||||
from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_default_kernels
|
||||
from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained("llamafactory/tiny-random-qwen3")
|
||||
original_rmsnorm_forward = model.model.layers[0].input_layernorm.forward
|
||||
original_swiglu_forward = model.model.layers[0].mlp.forward
|
||||
|
||||
model = apply_default_kernels(model=model, include_kernels=True)
|
||||
model = apply_kernels(model=model, config={"name": "auto"})
|
||||
|
||||
assert model.model.layers[0].input_layernorm.forward.__func__ is not original_rmsnorm_forward.__func__
|
||||
assert model.model.layers[0].mlp.forward.__func__ is not original_swiglu_forward.__func__
|
||||
|
||||
@@ -18,6 +18,7 @@ import torch.multiprocessing as mp
|
||||
|
||||
from llamafactory.v1.accelerator.interface import DistributedInterface
|
||||
from llamafactory.v1.config.model_args import ModelArguments
|
||||
from llamafactory.v1.config.training_args import TrainingArguments
|
||||
from llamafactory.v1.core.model_engine import ModelEngine
|
||||
from llamafactory.v1.plugins.model_plugins.parallelization.sequence_parallel import (
|
||||
SequenceParallelModelPlugin,
|
||||
@@ -33,15 +34,14 @@ def _test_sequence_parallel_loss(
|
||||
with dist_env(local_rank, world_size, master_port):
|
||||
model_args = ModelArguments(model="llamafactory/tiny-random-qwen3")
|
||||
|
||||
# Initialize distributed interface with config
|
||||
dist_config = {"cp_mode": "ulysses", "cp_size": cp_size, "dp_size": dp_size}
|
||||
DistributedInterface(dist_config)
|
||||
training_args = TrainingArguments(cp_mode="ulysses", cp_size=cp_size, dp_size=dp_size)
|
||||
DistributedInterface(training_args)
|
||||
|
||||
# Now create model engine
|
||||
model_engine = ModelEngine(model_args=model_args)
|
||||
|
||||
# Apply sequence parallel plugin
|
||||
SequenceParallelModelPlugin(dist_config.get("cp_mode", "ulysses"))(model_engine.model, dist_config)
|
||||
SequenceParallelModelPlugin(training_args.cp_mode)(model_engine.model, training_args.cp_size)
|
||||
|
||||
input_ids = torch.arange(1, batch_size * 5 + 1, dtype=torch.long).view(batch_size, 5)
|
||||
model_inputs = {
|
||||
|
||||
Reference in New Issue
Block a user