UlyssesSPAttentionHF with PEFT model

[shon-otmazgin-wix]

System Info

- `Accelerate` version: 1.12.0
- Platform: Linux-6.8.0-1029-aws-x86_64-with-glibc2.39
- `accelerate` bash location: /home/ubuntu/shon/remake-site2site/.venv/bin/accelerate
- Python version: 3.12.3
- Numpy version: 2.2.6
- PyTorch version: 2.9.0+cu129
- PyTorch accelerator: CUDA
- System RAM: 372.73 GB
- GPU type: NVIDIA L40S
- `Accelerate` default config:
	Not found

Information

• The official example scripts
• My own modified scripts

Tasks

• One of the scripts in the examples/ folder of Accelerate or an officially supported no_trainer script in the examples folder of the transformers repo (such as run_no_trainer_glue.py)
• My own task or dataset (give details below)

Reproduction

while calling

                mpu = UlyssesSPAttentionHF.register_with_transformers(
                    model_name_or_path=model,
                    sequence_parallel_size=sp_size,
                    seq_length=sp_handler.sp_seq_length,
                    seq_length_is_variable=sp_handler.sp_seq_length_is_variable,
                    core_attn_implementation=sp_handler.sp_attn_implementation,
                    micro_batch_size=batch_size_per_device,
                )

with PEFT model,

this code snippet (from deepspeed )is failing

        from transformers import PreTrainedModel
        if isinstance(model_name_or_path, PreTrainedModel):
            # we already have the model
            hf_model_config = model_name_or_path.config
        else:
            # if we don't have the model yet at this stage
            hf_model_config = AutoConfig.from_pretrained(model_name_or_path)

since PEFT model is not instance of PreTrainedModel

Expected behavior

accept also PEFT models

[akshatvishu]

Hey @shon-otmazgin-wix I think this is actually a DeepSpeed bug, not an Accelerate one!
DeepSpeed’s UlyssesSPAttentionHF.register_with_transformers() does:

if isinstance(model, PreTrainedModel):   # too strict

They use isinstance(model, PreTrainedModel) at ulysses_sp to decide whether to read model.config, but PEFT models (which have .config) are not subclasses of PreTrainedModel, so the check fails and DeepSpeed falls back to treating the object as a string path → instant TypeError.

I have filed an issue within DeepSpeed with potential fix and in the meantime, Accelerate can work around it by passing the base model instead of the PEFT if the maintainers agree with the change, although I feel that, DeepSpeed should rectify this at their end rather than in-here:

# accelerate/utils/deepspeed.py  
def _get_base_model_for_deepspeed(model: Module) -> Module:
    """
    Recursively unwraps PEFT/Distributed wrappers to find the underlying 
    Hugging Face PreTrainedModel that exposes a `.config` attribute.
    """
    # Loop to handle nested wrappers (e.g., PeftModel -> LoraModel -> Base)
    while hasattr(model, "base_model"):
        model = model.base_model
        # If the wrapper has a 'model' attribute (common in LoraModel), unwrap that too
        if hasattr(model, "model"):
            model = model.model
            
    return model

then we pass it like this at accelerator.py:

from accelerate.utils import UlyssesSPAttentionHF

from accelerate.utils.deepspeed import _get_base_model_for_deepspeed

# Unwrap the model to get the base HF model
model_for_ulysses = _get_base_model_for_deepspeed(model)

# Pass the UNWRAPPED model to DeepSpeed
UlyssesSPAttentionHF.register_with_transformers(
    model_name_or_path=model_for_ulysses,
    sequence_parallel_size=...
)

[shon-otmazgin-wix]

Hi @akshatvishu Thanks for that. agree with your suggestion. lets wait and see their response.

[shon-otmazgin-wix]

@akshatvishu doing below also valid workaround:

                mpu = UlyssesSPAttentionHF.register_with_transformers(
                    model_name_or_path=model.config.name_or_path,
                    sequence_parallel_size=sp_size,
                    seq_length=sp_handler.sp_seq_length,
                    seq_length_is_variable=sp_handler.sp_seq_length_is_variable,
                    core_attn_implementation=sp_handler.sp_attn_implementation,
                    micro_batch_size=batch_size_per_device,
                )

[akshatvishu]

Hey @shon-otmazgin-wix , I think you deleted your earlier comment but you were right to spot that DeepSpeed Ulysses isn’t active, so you’re falling back to standard data parallelism.

Total examples  : 20243
Optimization steps: 5061
→ effective batch size = 20243 / 5061 = 4

That matches data-parallel:
4 GPUs × 1 batch/GPU = total batch size 4
With Ulysses sequence parallelism all 4 GPUs collaborate on one sequence, so the total batch size should be 1.

The reason this is happening is because SFTTrainer (via SFTConfig) does not auto-load the parallelism_config from your Accelerate config file.
Accelerate launch only sets up the process group; the trainer decides how to split work from the Python config you pass.
Because parallelism_config was omitted, it defaults to plain data parallelism.

# Thus, you'll need to pass `ParallelismConfig` like this to the `SFTConfig`
from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,                 # all 4 GPUs on one sequence
    dp_shard_size=1,           # no DP sharding
    dp_replicate_size=1,       # no DP replication
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2",  # required for packed seqs + position_ids
    ),
)

training_args = SFTConfig(
    per_device_train_batch_size=1,
    gradient_accumulation_steps=1,
    parallelism_config=parallelism_config,  # ← this was missing
    bf16=True,
    logging_steps=1,
    # if you OOM with 80 k tokens
    # gradient_checkpointing=True,
)

trainer = SFTTrainer(model=model, args=training_args, ...)

Sanity check:
dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 → matches your GPU count.
After the fix expect:

    per-device batch size = 1
    total batch size = 1 (4 GPUs share one sequence)
    ~20 k optimisation steps per epoch

Critical: loss aggregation under SP
Ulysses does not use data-parallel loss.
Each GPU holds a slice of the sequence, so you must aggregate across SP ranks with a differentiable all_gather, weighted by valid tokens (shift_labels != -100).
The default SFTTrainer does not do this; add the snippet below to your training loop:

if parallelism_config.sp_size > 1:
    sp_group = accelerator.device_mesh["sp"].get_group()
    sp_world_size = parallelism_config.sp_size

    losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
    good_tokens = (batch["shift_labels"] != -100).view(-1).sum()
    good_tokens_per_rank = torch.distributed.nn.functional.all_gather(good_tokens, group=sp_group)

    total_loss = sum(
        losses_per_rank[r] * good_tokens_per_rank[r]
        for r in range(sp_world_size)
        if good_tokens_per_rank[r] > 0
    )
    total_good_tokens = sum(good_tokens_per_rank)
    loss = total_loss / max(total_good_tokens, 1)

Also note that shift_labels is consumed automatically by the model. If you use Liger Kernel (recommended for long sequences) this reduction is already fused and memory-efficient.

Memory note
With max_length=80_000 on L40S (46 GB) even sp_size=4 (20 k tokens / GPU) can OOM.
Enable gradient_checkpointing=Trueand keep flash_attention_2; SDPA can mishandle packed sequences.

Ref:

• https://huggingface.co/docs/accelerate/main/concept_guides/sequence_parallelism

[shon-otmazgin-wix]

I see another issue now. I installed "transformers_version": "5.0.0rc0", since SP works only with that version, and now my loss is None.

from the docs:

        # DeepSpeed SP automatically injects shift_labels into inputs (pre-shifted labels for SP).
        # The model's forward pass receives shift_labels via **kwargs and passes it to the loss function.
        # Both standard transformer models and Liger-patched models handle shift_labels correctly,
        # so we can directly use the computed loss from the model output.

see loss is none but we got logits so it means the shift_labels were not used.

configs:

deepspeed:

{
    "bf16": {
        "enabled": true
    },
    "zero_optimization": {
        "stage": 0
    },
    "gradient_accumulation_steps": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "seq_parallel_communication_data_type": "bf16"
}

accelerate:

compute_environment: LOCAL_MACHINE
deepspeed_config:
  deepspeed_config_file: ds_config.json
  zero3_init_flag: false
distributed_type: DEEPSPEED
machine_rank: 0
main_training_function: main
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
use_cpu: false
parallelism_config:
  parallelism_config_dp_replicate_size: 1
  parallelism_config_dp_shard_size: 1
  parallelism_config_sp_size: 4
  parallelism_config_sp_backend: deepspeed
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: flash_attention_2

sft config:

    training_args = SFTConfig(
        seed=seed,
        log_level='info',

        learning_rate=2e-4,

        num_train_epochs=1,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=1,
        ddp_find_unused_parameters=False,

        use_liger_kernel=False,

        save_strategy="steps",
        save_steps=100,

        logging_strategy="steps",
        logging_steps=1,

        remove_unused_columns=False,

        pad_to_multiple_of=4,
        packing=False,
        padding_free=False,
        dataset_kwargs={"skip_prepare_dataset": False},
        max_length=30_000,
        packing_strategy='bfd',

        completion_only_loss=True,
        report_to=['wandb'],
        fp16=False,
        bf16=True,
        bf16_full_eval=True,

        optim="adamw_torch_fused",

        gradient_checkpointing=True,
    )

[shon-otmazgin-wix]

@akshatvishu ok I found why above (loss=None) happens

DataCollatorForLanguageModeling adding labels but then ulysses_sp.py is removing labels and adds shift_labels but then in the model forward:

        loss = None
        if labels is not None:
            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)

no labels ...

this is from deepspeed ulysses_sp.py

            # because we have to gather logits from all sp ranks we have to do the loss function ourselves
            # therefore remove labels to avoid an attempt to calculate loss by transformers
            labels = batch.pop("labels")
            labels = torch.nn.functional.pad(labels, (0, 1), value=-100)
            batch["shift_labels"] = labels[..., 1:].contiguous()
            # free up temp memory
            del labels

[akshatvishu]

I think loss=None can happen due to various reasons:

• The UlyssesSPDataLoaderAdapter slices the input sequence across GPUs.
  To ensure correct loss calculation (distributed weighted sum), it deletes labels from the batch and replaces them with shift_labels.
• HF models return loss=None if labels are missing.
  The trainer then fails because it tries to log/back-propagate None.
• If you use completion_only_loss=True without the right collator, labels might never be created, so DeepSpeed skips its logic entirely.

You can try doing something like:

import torch
import torch.nn as nn
import torch.distributed as dist
from trl import SFTTrainer, SFTConfig, DataCollatorForCompletionOnlyLM
from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

#  Define Parallelism Config 
# Required for Trainer to initialize the distributed environment correctly
parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,           # Set to your number of GPUs
    dp_shard_size=1,     # No data parallelism (1 sequence across 4 GPUs)
    dp_replicate_size=1,
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2"
    )
)

# Define the Collator 
# Crucial: Creates 'labels' so DeepSpeed Ulysses adapter activates..
response_template = "### Answer:" 
collator = DataCollatorForCompletionOnlyLM(
    response_template=response_template, 
    tokenizer=tokenizer
)

# Define Custom Trainer 
# Override `compute_loss`  to handle the fact that DeepSpeed deleted 'labels'
class DeepSpeedUlyssesTrainer(SFTTrainer):
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        # forward pass 
        # DeepSpeed adapter has already sliced inputs, popped 'labels' and added 'shift_labels'
        outputs = model(**inputs)
        
        # If model returned no loss (because labels were missing), calculate it manually
        if outputs.loss is None and "shift_labels" in inputs:
            logits = outputs.logits
            labels = inputs["shift_labels"]
            
            # Standard Causal LM Loss (Flatten and compute CrossEntropy)
            shift_logits = logits.view(-1, self.model.config.vocab_size)
            shift_labels = labels.view(-1)
            
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(shift_logits, shift_labels)
        else:
            loss = outputs.loss

        # NOTE: For Sequence Parallelism > 1, the gradients are synced automatically during backward.
        # However, for accurate logging, you ideally want to aggregate this loss across the SP group.
        # This basic implementation ensures training runs without crashing.
        
        return (loss, outputs) if return_outputs else loss


training_args = SFTConfig(
    output_dir="./output",
    # CRITICAL: Disable this flag to prevent 'SFTTrainer' from fighting with our custom collator
    completion_only_loss=False, 
    # CRITICAL: Packing complicates SP (sequences must be perfectly divisible). Keep False for now.
    packing=False,
    # CRITICAL: Pass config so Trainer sets up the communication groups
    parallelism_config=parallelism_config,
    
    per_device_train_batch_size=1, 
    gradient_accumulation_steps=1,
    learning_rate=2e-4,
    bf16=True,
    gradient_checkpointing=True,
    logging_steps=1
)

trainer = DeepSpeedUlyssesTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    formatting_func=formatting_prompts_func, 
    data_collator=collator,
)

trainer.train()

Ref:

• for defining custom loss and understand why do lables needed to be pre-shifted:
  • https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-parallelism/
  • https://github.com/deepspeedai/DeepSpeedExamples/blob/uly-hf/post_training/sequence_parallelism/test_ulysses.py
• https://huggingface.co/docs/trl/sft_trainer#train-on-completion-only
• https://huggingface.co/docs/trl/sft_trainer

[shon-otmazgin-wix]

Thanks @akshatvishu for the detailed explanation

re: SP + SFTTrainer:
I will try what you suggest but also I say what transformers v5.0.0rc does

        # since DataLoader was Accelerate prepared w/o a model arg in the same call, we now have to complete the DL wrapping for ALST/UlyssesSP, after model has been prepared
        pc = getattr(self.accelerator, "parallelism_config", None)
        if pc is not None and pc.sp_backend == "deepspeed" and pc.sp_enabled:
            train_dataloader = self.accelerator.deepspeed_ulysses_dl_adapter(train_dataloader, model)

I need to do it also?

re: loss=None
I use DataCollatorForLanguageModeling with completion_only_loss=True (done by sft trainer) and it creates labels but deep speed as I mentioned delete them here:

            # because we have to gather logits from all sp ranks we have to do the loss function ourselves
            # therefore remove labels to avoid an attempt to calculate loss by transformers
            labels = batch.pop("labels")
            labels = torch.nn.functional.pad(labels, (0, 1), value=-100)
            batch["shift_labels"] = labels[..., 1:].contiguous()
            # free up temp memory
            del labels

but also deepspeed pass shift_labels in kwargs in the forward funtions BUT the model return loss only if labels is not none:

        loss = None
        if labels is not None:
            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)

[akshatvishu]

Do I need to wrap the DataLoader manually?

No. SFTTrainer (inheriting from Trainer) automatically wraps the dataloader with deepspeed_ulysses_dl_adapter if parallelism_config is present in TrainingArguments.

Reference: transformers/trainer.py (lines ~2335) checks for pc.sp_enabled and wraps it for you.

I use DataCollatorForLanguageModeling with completion_only_loss=True (done by sft trainer) and it creates labels but deep speed as I mentioned delete them here:

The UlyssesSPDataLoaderAdapter explicitly executes labels = batch.pop("labels"). It removes them to prevent the standard model from calculating an incorrect "local" loss on sharded sequences. It replaces them with shift_labels.

Reference: deepspeed/runtime/sequence_parallel/ulysses_sp.py

You can read https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-parallelism/
for further on information on "Why do labels need to be pre-shifted?".

BUT the model return loss only if labels is not none:

Yes. Standard Hugging Face models (like LlamaForCausalLM) strictly check if labels is not None: before calculating loss. Since DeepSpeed deleted them, the model returns None, causing the Trainer to crash.
Hence, if we look at trainer.py:

        # DeepSpeed SP automatically injects shift_labels into inputs (pre-shifted labels for SP).
        # The model's forward pass receives shift_labels via **kwargs and passes it to the loss function.
        # Both standard transformer models and Liger-patched models handle shift_labels correctly,
        # so we can directly use the computed loss from the model output.
        # See: https://huggingface.co/docs/accelerate/en/concept_guides/sequence_parallelism
        outputs = model(**inputs)
        loss = outputs.loss                            # <--- CRASH: This is None because 'labels' was popped!

Now, if you check https://huggingface.co/docs/accelerate/en/concept_guides/sequence_parallelism:

It is stated that "If you use Liger Kernel it already knows how to handle shift_labels so you don’t need to go through manual loss calculation, just calling model(**batch) will already get the loss calculated and done in a very memory-efficient way. "

# From the Article's example code:
loss = model.module.loss_function(
    logits=outputs.logits,
    labels=None,             # <--- Article explicitly passes None
    shift_labels=shift_labels, # <--- Article explicitly passes shift_labels
    ...
)

Hence , if you pass the use_liger_kernel=True flag like, you can skip the whole custom overriding compute_loss in the Trainer that I shared earlier !

training_args = SFTConfig(
    # ... other args ...
   
    use_liger_kernel=True,  # <--- Enable this!
    
    # Still require these for correct setup:
    completion_only_loss=False, # Disable built-in masking (conflict risk)
    packing=False,              # Disable packing for now
    parallelism_config=parallelism_config,
)

Otherwise, if we are not using use_liger_kernel=False ; then we do need to define the custom_loss like they did in https://huggingface.co/docs/accelerate/en/concept_guides/sequence_parallelism and https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-parallelism/:

import torch
import torch.nn as nn
import torch.distributed as dist
from trl import SFTTrainer, SFTConfig, DataCollatorForCompletionOnlyLM
from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

# Define Parallelism Config
parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,             # Adjust to your GPU count
    dp_shard_size=1,       # 1 = No Data Parallelism (Sequence Parallel only)
    dp_replicate_size=1,
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2"
    )
)

# Define Collator 
# We use this to guarantee 'labels' exist initially so DeepSpeed activates.
# We disable 'completion_only_loss' in config to avoid TRL conflicts.
response_template = "### Answer:" 
collator = DataCollatorForCompletionOnlyLM(
    response_template=response_template, 
    tokenizer=tokenizer
)

# --- 3. Custom Trainer ---
class DeepSpeedUlyssesTrainer(SFTTrainer):
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        # Forward Pass
        # DeepSpeed adapter has already sliced inputs, popped 'labels', and added 'shift_labels'
        outputs = model(**inputs)
        
        # . Manual Loss Calculation
        # If model returned None (because labels were popped), we calculate it here
        if outputs.loss is None and "shift_labels" in inputs:
            logits = outputs.logits
            labels = inputs["shift_labels"]
            
            # Standard CrossEntropy
            shift_logits = logits.view(-1, self.model.config.vocab_size)
            shift_labels = labels.view(-1)
            
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(shift_logits, shift_labels)
        else:
            loss = outputs.loss

        # Distributed Loss Aggregation (Code from HF sequence_parallelism article )
        # We assume SP > 1. We must weight the loss by valid tokens across GPUs.
        if self.args.parallelism_config.sp_size > 1:
            sp_group = self.accelerator.torch_device_mesh["sp"].get_group()
            
            # Gather losses and token counts from all SP ranks
            losses_per_rank = dist.nn.functional.all_gather(loss, group=sp_group)
            valid_tokens = (inputs["shift_labels"] != -100).view(-1).sum()
            valid_tokens_per_rank = dist.nn.functional.all_gather(valid_tokens, group=sp_group)
            
            # Weighted average: sum(loss_i * count_i) / sum(total_count)
            total_loss = sum(
                losses_per_rank[i] * valid_tokens_per_rank[i] 
                for i in range(len(losses_per_rank))
            )
            total_count = sum(valid_tokens_per_rank)
            
            # Avoid division by zero
            loss = total_loss / max(total_count, 1)

        return (loss, outputs) if return_outputs else loss


training_args = SFTConfig(
    output_dir="./output",
    completion_only_loss=False,  # CRITICAL: We handle masking via collator
    packing=False,               # CRITICAL: Incompatible with this specific setup
    parallelism_config=parallelism_config,
    per_device_train_batch_size=1, 
    gradient_accumulation_steps=1,
    learning_rate=2e-4,
    bf16=True,
    gradient_checkpointing=True,
    logging_steps=1
)

trainer = DeepSpeedUlyssesTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    formatting_func=formatting_prompts_func, 
    data_collator=collator,
)

trainer.train()

[shon-otmazgin-wix]

@akshatvishu amazing! thanks for the detailed explanation!

Regarding your reference about deepspeed_ulysses_dl_adapter this is done only in v5.0.0rc which is not released yet and has many breaks such as MoE. if I use stable version - v4.57.3 I need to do it on my own right?

and regarding
completion_only_loss=False, # CRITICAL: We handle masking via collator
can you explain here more? beacuse I need to set it to True...

DataCollatorForLanguageModeling also do masking...

if self.completion_only_loss and "completion_mask" in examples[0]:
    completion_mask = pad(
        completion_mask, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of
    )
    output["labels"][completion_mask == 0] = -100  # mask everything that is not in the completion

also this use case example is with the loss=none thing so it will not work

accelerate/examples/alst_ulysses_sequence_parallelism/sp-alst.py

Line 129 in db8835c

# The model automatically receives shift_labels via **kwargs and uses it for loss computation.

[shon-otmazgin-wix]

btw I tried pass parallelism_config to sft trainer and still it use DDP at least for v4.57.3

[akshatvishu]

if I use stable version - v4.57.3 I need to do it on my own right?

btw I tried pass parallelism_config to sft trainer and still it use DDP at least for v4.57.3

Trainer in current stable v4.57.3 simply treats your setup as standard DDP because it lacks the logic to look for parallelism_config or wrap the dataloader. Thus, you must override get_train_dataloader in your custom trainer to apply the UlyssesSPDataLoaderAdapter manually.

and regarding
completion_only_loss=False, # CRITICAL: We handle masking via collator
can you explain here more? beacuse I need to set it to True...

The completion_only_loss=True flag in SFTTrainer triggers internal logic to create a data collator for you. But in-here we're creating our own collator = DataCollatorForCompletionOnlyLM(...) to ensure labels exist for DeepSpeed. Hence, If you pass a custom collator AND set the flag to True, the Trainer creates a conflict (it tries to wrap your wrapper or validate arguments that don't match). Therefore, we are explicit passing the collator and setting the flag to False. Oour collator still applies the masking (setting prompt tokens to -100), so you do get completion-only loss training, effectively "bypassing" the Trainer's check.

also this use case example is with the loss=none thing so it will not work
at sp-alst.py:

# The model automatically receives shift_labels via **kwargs and uses it for loss computation.
# Both standard transformer models and Liger-patched models handle this correctly.

The standard Hugging Face models (like LlamaForCausalLM) in the current stable version (v4.57.3) do not look for shift_labels in their forward method. They return loss=None if labels are missing. The example likely assumes you are using a patched model (like Liger Kernel).

You can further look at this example, src/llamafactory/train/sft/trainer.py, they're at transformer version transformers==4.51.1. Here, they're completely bypassing the model's loss calculation logic when Sequence Parallelism (ALST) is enabled:

# src/llamafactory/train/sft/trainer.py
# In __init__, they force the Trainer to stop trying to pass loss arguments to the model.
if processor is not None:
   # avoid wrong loss under gradient accumulation
   self.model_accepts_loss_kwargs = False

In compute_loss, they don't call model(**inputs) and expect a loss back. They detect if ALST (Ulysses) is active and hand off the entire calculation to a specialized handler.

# src/llamafactory/train/sft/trainer.py
def compute_loss(self, model, inputs, return_outputs=False, **kwargs):

    # Check if we should use ALST loss computation
    if should_use_alst_loss(inputs, sequence_parallel_group):
        # ...
        # calculate LOSS externally here:
        loss_tuple = self.alst_loss_handler.compute_alst_loss(model, inputs, return_outputs)
        return loss_tuple

Thus

import torch
import torch.nn as nn
import torch.distributed as dist
from trl import SFTTrainer, SFTConfig, DataCollatorForCompletionOnlyLM
from deepspeed.runtime.sequence_parallel.ulysses_sp import UlyssesSPDataLoaderAdapter
from deepspeed.utils import groups as ds_groups


# We use this collator to enforce "completion only" masking manually.
# This ensures 'labels' are created correctly before DeepSpeed sees the batch.
response_template = "### Answer:" 
collator = DataCollatorForCompletionOnlyLM(
    response_template=response_template, 
    tokenizer=tokenizer
)


class DeepSpeedUlyssesTrainer(SFTTrainer):
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Explicitly tell Trainer NOT to pass loss kwargs (like 'labels') 
        # to the model. We want the model to return outputs only, so we can handle 
        # the missing labels logic ourselves in 'compute_loss'.
        self.model_accepts_loss_kwargs = False

    # Manually wrap DataLoader 
    # Logic: Transformers v4.57.3 does not natively support Ulysses. We must intercept 
    # the dataloader and wrap it with DeepSpeed's adapter to handle sequence slicing.
    def get_train_dataloader(self):
        # 1. Get the standard dataloader from parent
        dataloader = super().get_train_dataloader()
        
        # Get the Sequence Parallel Group from DeepSpeed
        # We rely on your ds_config.json having "sequence_parallel_size" set.
        sp_group = ds_groups._get_sequence_parallel_group()
        if sp_group is None:
            # Fallback if groups aren't initialized yet (rare in Trainer flow)
            return dataloader

        # Wrap with Ulysses Adapter
        # This performs the magic: slices batch['input_ids'], pops 'labels', adds 'shift_labels'
        return UlyssesSPDataLoaderAdapter(
            dataloader,
            sp_rank=dist.get_rank(sp_group),
            sp_group=sp_group,
            sp_world_size=dist.get_world_size(sp_group),
            device=self.model.device
        )

    # Manually Compute & Aggregate Loss
    # Logic: DeepSpeed popped 'labels', so model returns loss=None. We calculate it here.
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        # Forward Pass (Model returns logits, but None for loss)
        outputs = model(**inputs)
        
        #  Handle the Missing Loss
        if outputs.loss is None and "shift_labels" in inputs:
            logits = outputs.logits
            labels = inputs["shift_labels"]
            
            # Compute Local Loss (Standard CrossEntropy)
            shift_logits = logits.view(-1, self.model.config.vocab_size)
            shift_labels = labels.view(-1)
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(shift_logits, shift_labels)
            
            # Distributed Loss Aggregation
            # Since sequences are split, we must weight the loss by valid tokens on each GPU
            # to get the mathematically correct global loss for logging/gradients.
            sp_group = ds_groups._get_sequence_parallel_group()
            if sp_group is not None and dist.get_world_size(sp_group) > 1:
                # Gather losses from all ranks
                losses = dist.nn.functional.all_gather(loss, group=sp_group)
                
                # Count valid (non-padding) tokens on this rank
                valid_tokens = (labels != -100).view(-1).sum()
                valid_tokens_all = dist.nn.functional.all_gather(valid_tokens, group=sp_group)
                
                # Weighted average: sum(loss_i * valid_tokens_i) / total_valid_tokens
                total_loss = sum(losses[i] * valid_tokens_all[i] for i in range(len(losses)))
                total_count = sum(valid_tokens_all)
                
                loss = total_loss / max(total_count, 1)
        else:
            # Fallback for standard training (DDP)
            loss = outputs.loss

        return (loss, outputs) if return_outputs else loss


training_args = SFTConfig(
    output_dir="./output",
    #  Set False to avoid Trainer validation errors. 
    # We explicitly pass our own 'collator' (defined above) which handles the masking.
    completion_only_loss=False, 
    
    # Packing is incompatible with the DataCollatorForCompletionOnlyLM logic.
    packing=False,
    
    # Standard Params
    per_device_train_batch_size=1, 
    gradient_accumulation_steps=1,
    learning_rate=2e-4,
    bf16=True,
    gradient_checkpointing=True,
    logging_steps=1
)

trainer = DeepSpeedUlyssesTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    formatting_func=formatting_prompts_func, 
    data_collator=collator, # Pass our specific collator here
)

trainer.train()

Also, if you're computing logits for a huge sequence (e.g., 100k+ tokens) to calculate loss consumes massive memory (vocab_size * seq_len). Then we should use tiling to break this into chunks to save memory.

You can look at https://github.com/snowflakedb/ArcticTraining/blob/main/arctic_training/trainer/sft_trainer.py#L141C30-L141C50 for further implementation of that!

Thus, our current workflow(UlyssesTrainer) v/s standard Trainer behaviour :

stateDiagram-v2
    direction TB

    state StandardTrainer {
        [*] --> LoadBatch
        LoadBatch --> ForwardWithLabels
        ForwardWithLabels --> LossInModel
        LossInModel --> Backward
        Backward --> [*]
    }

    state UlyssesTrainer {
        [*] --> AdaptBatch
        AdaptBatch --> RemoveLabels
        RemoveLabels --> AddShiftLabels
        AddShiftLabels --> SlicePerGpu
        SlicePerGpu --> ForwardNoLabels
        ForwardNoLabels --> LogitsOnly
        LogitsOnly --> ComputeLoss
        ComputeLoss --> CrossEntropy
        CrossEntropy --> AggregateLoss
        AggregateLoss --> Backward
        Backward --> [*]
    }

Loading

[shon-otmazgin-wix]

Hi @akshatvishu thanks again.

I tried what you suggested but doing
sp_group = ds_groups._get_sequence_parallel_group() inside get_train_dataloader returns
KeyError: 'No sequence parallel group found'

maybe its too early and that's why v5.0.0rc call UlyssesSPDataLoaderAdapter inside _inner_training_loop function?

I see the mpu is set from deepspeed only when calling prepare inside _inner_training_loop

[akshatvishu]

Are you passing deepspeed_config either with Accelerate launch:

accelerate launch --config_file your_accelerate_config.yaml --deepspeed_config_file ds_config.json train.py

or
within

training_args = SFTConfig(
    # ... other args ...
    deepspeed="ds_config.json"  # Path to your file
)

and your deepspeed_config has "sequence_parallel_size": 4 # ur GPU ?