liger-kernel-nightly 0.5.9.dev20250512213150__tar.gz → 0.5.9.dev20250515034325__tar.gz

{liger_kernel_nightly-0.5.9.dev20250512213150 → liger_kernel_nightly-0.5.9.dev20250515034325}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.9.dev20250512213150
+Version: 0.5.9.dev20250515034325
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

liger_kernel_nightly-0.5.9.dev20250515034325/examples/medusa/requirements.txt

@@ -0,0 +1,3 @@
+accelerate==1.6.0
+scikit-learn
+transformers==4.51.3

{liger_kernel_nightly-0.5.9.dev20250512213150 → liger_kernel_nightly-0.5.9.dev20250515034325}/examples/medusa/scripts/llama3_8b_medusa.sh

@@ -22,9 +22,6 @@ export MEDUSA_LR_MULTIPLIER=4.0
 accelerate launch --config_file fsdp/acc-fsdp.conf \
     --num_machines $NUM_NODES \
     --num_processes $WORLD_SIZE \
-    --main_process_ip $MASTER_ADDR \
-    --main_process_port $MASTER_PORT \
-    --machine_rank $RANK \
     train.py \
     --bf16 True \
     --output_dir $OUTPUT_DIR \
@@ -32,7 +29,7 @@ accelerate launch --config_file fsdp/acc-fsdp.conf \
     --per_device_train_batch_size $LOCAL_TRAIN_BATCH_SIZE \
     --per_device_eval_batch_size 1 \
     --gradient_accumulation_steps $GRADIENT_ACCUMULATION_STEPS \
-    --evaluation_strategy "no" \
+    --eval_strategy "no" \
     --save_strategy "no" \
     --prediction_loss_only \
     --learning_rate $LR \
@@ -53,4 +50,4 @@ accelerate launch --config_file fsdp/acc-fsdp.conf \
     --medusa_lr_multiplier $MEDUSA_LR_MULTIPLIER \
     --medusa_only_heads False \
     --medusa_return True \
-    --use_liger True
+    --use_liger True

{liger_kernel_nightly-0.5.9.dev20250512213150 → liger_kernel_nightly-0.5.9.dev20250515034325}/examples/medusa/train.py

@@ -32,21 +32,18 @@ from callback import EfficiencyCallback
 from medusa_util import add_medusa_heads
 from safetensors.torch import save_file
 from sklearn.model_selection import train_test_split
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.distributed.fsdp.fully_sharded_data_parallel import FullStateDictConfig
-from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
 from torch.utils.data import Dataset
 from transformers import Trainer
 from transformers.trainer_pt_utils import LabelSmoother
 
-from liger_kernel.transformers import apply_liger_kernel_to_llama
+from liger_kernel.transformers import AutoLigerKernelForCausalLM
 
 IGNORE_TOKEN_ID = LabelSmoother.ignore_index
 
 
 @dataclass
 class ModelArguments:
-    model_name_or_path: Optional[str] = field(default="meta-llama/Meta-Llama-3-8B")
+    model_name_or_path: Optional[str] = field(default="meta-llama/Meta-Llama-3-8B-Instruct")
 
 
 @dataclass
@@ -310,29 +307,36 @@ def train():
     print(tokenizer(["This is a test", "secondary"], padding=True))
     print(tokenizer.apply_chat_template([{"role": "user", "content": "This is a test"}]))
 
-    # Load model and tokenizer
-    model = transformers.AutoModelForCausalLM.from_pretrained(
-        model_args.model_name_or_path,
-        # config=config,
-        cache_dir=training_args.cache_dir,
-        torch_dtype=torch.bfloat16,
-    )
+    def _model_loader():
+        # we use a customized model loader to inject medusa heads to FSDP-wrapped model variables properly.
+        # see https://github.com/linkedin/Liger-Kernel/issues/309#issuecomment-2455077623 for details.
 
-    if training_args.use_liger is True:
-        apply_liger_kernel_to_llama()
+        # Load model
+        if training_args.use_liger:
+            model_builder = AutoLigerKernelForCausalLM.from_pretrained
+        else:
+            model_builder = transformers.AutoModelForCausalLM.from_pretrained
+        model = model_builder(
+            model_args.model_name_or_path,
+            cache_dir=training_args.cache_dir,
+            torch_dtype=torch.bfloat16,
+        )
 
-    # Freeze the base model
-    for param in model.base_model.parameters():
-        param.requires_grad = False
+        # Freeze the base model
+        for param in model.base_model.parameters():
+            param.requires_grad = False
+
+        # Inject Medusa heads
+        add_medusa_heads(
+            model,
+            training_args.medusa_num_heads,
+            training_args.medusa_num_layers,
+            training_args.medusa_return,
+            training_args.medusa_only_heads,
+            training_args.use_liger,
+        )
+        return model
 
-    add_medusa_heads(
-        model,
-        training_args.medusa_num_heads,
-        training_args.medusa_num_layers,
-        training_args.medusa_return,
-        training_args.medusa_only_heads,
-        training_args.use_liger,
-    )
     # Format output dir
     training_args.output_dir = f"{training_args.output_dir}_medusa_mlp_{model_args.model_name_or_path.split('/')[-1]}_medusa_{training_args.medusa_num_heads}_lr_{training_args.learning_rate}_layers_{training_args.medusa_num_layers}"
 
@@ -341,7 +345,7 @@ def train():
 
     # Start trainner
     trainer = Trainer(
-        model=model,
+        model_init=_model_loader,
         tokenizer=tokenizer,
         args=training_args,
         callbacks=[EfficiencyCallback()],
@@ -355,17 +359,11 @@ def train():
 
     if training_args.medusa_return and training_args.medusa_only_heads:
         # Save only the updated head without saving the backbone model
-        if hasattr(model, "module"):
-            lm_head = model.module.medusa_head
-        else:
-            lm_head = model.medusa_head
-
-        with FSDP.state_dict_type(
-            model,
-            StateDictType.FULL_STATE_DICT,
-            FullStateDictConfig(offload_to_cpu=True),
-        ):
-            state_dict = lm_head.state_dict()
+        state_dict = {
+            k.replace("medusa_head.", ""): v.to(torch.bfloat16)
+            for k, v in trainer.accelerator.get_state_dict(trainer.model).items()
+            if "medusa_head" in k
+        }
 
         # Save Medusa heads
         if local_rank == 0:
@@ -373,9 +371,9 @@ def train():
                 state_dict,
                 os.path.join(training_args.output_dir, "medusa_lm_head.safetensors"),
             )
+        trainer.accelerator.wait_for_everyone()
     else:
         # Save the whole model weight
-        trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
         trainer.save_model(training_args.output_dir)
 
 

{liger_kernel_nightly-0.5.9.dev20250512213150 → liger_kernel_nightly-0.5.9.dev20250515034325}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.9.dev20250512213150"
+version = "0.5.9.dev20250515034325"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

{liger_kernel_nightly-0.5.9.dev20250512213150 → liger_kernel_nightly-0.5.9.dev20250515034325}/src/liger_kernel_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.9.dev20250512213150
+Version: 0.5.9.dev20250515034325
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

liger_kernel_nightly-0.5.9.dev20250512213150/examples/medusa/requirements.txt

@@ -1,3 +0,0 @@
-accelerate==0.29.3
-scikit-learn
-transformers==4.42.4