From 657571f17edcaa9ff4420e935e6deee0ca338cf3 Mon Sep 17 00:00:00 2001 From: Dan Saunders Date: Tue, 6 Jan 2026 13:01:48 -0500 Subject: [PATCH] fixes --- unsloth-cli.py | 11 +++++- unsloth/context_parallel.py | 74 +++++++++++++------------------------ 2 files changed, 34 insertions(+), 51 deletions(-) diff --git a/unsloth-cli.py b/unsloth-cli.py index b1c1f21c34..c9a38f7d8f 100644 --- a/unsloth-cli.py +++ b/unsloth-cli.py @@ -31,7 +31,6 @@ Happy fine-tuning! import argparse import os -from typing import Optional def run(args): @@ -46,6 +45,13 @@ def run(args): logging.getLogger("hf-to-gguf").setLevel(logging.WARNING) + # Convert gradient checkpointing arg to expected type + use_gradient_checkpointing = args.use_gradient_checkpointing + if use_gradient_checkpointing.lower() == "true": + use_gradient_checkpointing = True + elif use_gradient_checkpointing.lower() == "false": + use_gradient_checkpointing = False + # Load model and tokenizer device_map, distributed = prepare_device_map() model, tokenizer = FastLanguageModel.from_pretrained( @@ -54,6 +60,7 @@ def run(args): dtype = args.dtype, load_in_4bit = args.load_in_4bit, device_map = device_map, + use_gradient_checkpointing = use_gradient_checkpointing, ) # Configure PEFT model @@ -72,7 +79,7 @@ def run(args): lora_alpha = args.lora_alpha, lora_dropout = args.lora_dropout, bias = args.bias, - use_gradient_checkpointing = args.use_gradient_checkpointing, + use_gradient_checkpointing = use_gradient_checkpointing, random_state = args.random_state, use_rslora = args.use_rslora, loftq_config = args.loftq_config, diff --git a/unsloth/context_parallel.py b/unsloth/context_parallel.py index 30ead77ad6..bf4695806c 100644 --- a/unsloth/context_parallel.py +++ b/unsloth/context_parallel.py @@ -239,6 +239,10 @@ class ContextParallelManager: global_tokens = local_tokens.clone() dist.all_reduce(global_tokens, op = dist.ReduceOp.SUM, group = self._cp_group) + # Cache token counts for reduce_grad_norm + self._cached_local_tokens = local_tokens.detach() + self._cached_global_tokens = global_tokens.detach() + # Weight loss by local fraction weight = local_tokens.detach() / global_tokens.detach() weighted_loss = tensor * weight @@ -254,23 +258,21 @@ class ContextParallelManager: """ Reduce gradient norm across CP group. - Each rank computes a local gradient norm (L2). The global norm is: - sqrt(sum(local_norm_i^2 for all ranks)) - - This is needed because the Trainer computes grad_norm locally, but with - CP each rank only has partial gradients for its sequence shard. + The gradients are already scaled by (local_tokens / global_tokens) from + the weighted loss in reduce_loss. We sum local norms directly - no + additional weighting needed since the scaling is already in the gradients. """ if self._cp_group is None: return grad_norm - # Square the local norm, sum across ranks, then sqrt - local_norm_sq = torch.tensor( - grad_norm**2, + # Simple sum of local norms - gradients already have the weight factor + local_norm = torch.tensor( + grad_norm, dtype = torch.float32, device = torch.device(DEVICE_TYPE_TORCH), ) - dist.all_reduce(local_norm_sq, op = dist.ReduceOp.SUM, group = self._cp_group) - return float(local_norm_sq.sqrt().item()) + dist.all_reduce(local_norm, op = dist.ReduceOp.SUM, group = self._cp_group) + return float(local_norm.item()) def patch_sft_config(): @@ -404,35 +406,15 @@ def patch_sft_trainer() -> None: ): setattr(accelerator.state, "device_mesh", mesh) - # When using pure context parallelism (dp_world_size=1), disable DDP - # to avoid gradient checkpointing compatibility issues - if manager and manager.data_parallel_world_size == 1: - try: - from accelerate.utils import DistributedType - - args = getattr(self, "args", None) - distributed_state = getattr(args, "distributed_state", None) - if ( - distributed_state is not None - and distributed_state.distributed_type == DistributedType.MULTI_GPU - ): - distributed_state.distributed_type = DistributedType.NO - except ImportError: - pass - - # Enable sync_each_batch when using CP with gradient accumulation and DDP. - # This keeps the computation graph constant for DDP + static_graph mode. + # Enable sync_each_batch when using CP to ensure consistent computation graph + # for DDP + static_graph mode. This is needed because ring attention changes + # the graph structure, and sync_each_batch keeps it constant. if ( manager - and manager.data_parallel_world_size > 1 # Only needed with actual DP and accelerator is not None and hasattr(accelerator, "gradient_state") ): - grad_accum_steps = getattr( - getattr(self, "args", None), "gradient_accumulation_steps", 1 - ) - if grad_accum_steps > 1: - accelerator.gradient_state.plugin_kwargs["sync_each_batch"] = True + accelerator.gradient_state.plugin_kwargs["sync_each_batch"] = True # Attach attention hooks for proper ring attention behavior with load balancing. # This ensures attention_mask is removed and is_causal=True for all self_attn calls. @@ -446,20 +428,17 @@ def patch_sft_trainer() -> None: manager = getattr(self, "_context_parallel_manager", None) kwargs.pop("num_items_in_batch", None) - # For context parallelism with shift_labels, prefer letting the model - # handle the pre-shifted targets when it advertises support. Otherwise - # fall back to an external loss that consumes the sharded tensors. + # For context parallelism with shift_labels, always use external loss. + # This is required because unsloth's fused CE loss pre-computes gradients + # during forward and ignores grad_output in backward. When we weight the + # loss in reduce_loss(), the chain rule requires gradients be multiplied + # by the weight, but fused loss backward returns pre-computed gradients + # without this factor. External loss uses standard autograd which correctly + # propagates the weight through the chain rule. shift_labels = inputs.get("shift_labels") use_cp_shift_labels = manager and isinstance(shift_labels, torch.Tensor) - model_supports_shift_labels = bool( - getattr( - model, - "_unsloth_supports_context_parallel_shift_labels", - False, - ) - ) - if use_cp_shift_labels and not model_supports_shift_labels: + if use_cp_shift_labels: # Remove labels so model doesn't compute loss internally saved_labels = inputs.pop("labels", None) # Also remove shift_labels from inputs (model doesn't expect it) @@ -571,11 +550,8 @@ def patch_sft_trainer() -> None: @functools.wraps(original_log) def patched_log(self, logs, start_time = None): manager = getattr(self, "_context_parallel_manager", None) - # Reduce grad_norm across CP group if present if manager and "grad_norm" in logs: - grad_norm = logs["grad_norm"] - if isinstance(grad_norm, (int, float)): - logs["grad_norm"] = manager.reduce_grad_norm(grad_norm) + logs["grad_norm"] = manager.reduce_grad_norm(logs["grad_norm"]) return original_log(self, logs, start_time) trainer_cls.__init__ = patched_init