Liger-Kernel fails to compute loss when using train+eval and SP>1

[sfc-gh-sbekman]

note: This issue at the moment is FYI to compile my observations and what I have tried to make it work. I don't personally have a need to do eval while training so unless someone asks this might not be an issue for anybody so why bother as this is not our code failing.

---

Liger-Kernel fails to compute loss when using train+eval and returns None - removing eval (leaving pure training) removes the problem. Tested with SFT trainer and SP>1.

Last I checked its fused_linear_cross_entropy feature was computing backward in forward w/o checking if the model is running under no_grad. So I tried to disable just that kernel with:

        return AutoLigerKernelForCausalLM.from_pretrained(
            self.config.name_or_path,
[...]
            fused_linear_cross_entropy=False,
        )

but it still returned None.

So for now I have added an exception:

        if self.config.model.type == "liger":
            # letting liger do fused logits+loss calculation
            outputs = self.model(**batch, use_cache=False)
            loss = outputs.loss
            if loss is None:
                raise ValueError("Liger-Kernel failed to compute loss (returned None) - it's known to fail with eval enabled along train steps.")

I also checked shift_labels contain valid non -100 tokens.

if users want this feature we would need to escalate this with the Liger-Kernel devs.

[Tcc0403]

Hi, I would like to look into this issue. Do I just run this recipe sequence-parallelism/run-h100-sp8-llama-8b-no-ulysses-yes-liger-no-extras.yml to reproduce the issue? Or how do I enable eval mode with ArticTraining? Feel free to point me in some docs that could be helpful. I'm new to this project but eager to help!

A question about mentioned recipe: is sequence_parallel_size supposed to 8 instead of 1? given the filename is called sp8

ArcticTraining/projects/sequence-parallelism/run-h100-sp8-llama-8b-no-ulysses-yes-liger-no-extras.yml

Line 6 in 8f61c50

sequence_parallel_size: 1

[sfc-gh-sbekman]

Thank you for reporting, @Tcc0403 - the size will be fixed here #282

Let me create a recipe for the liger issue with eval - the eval was added after ALST was released.

[sfc-gh-sbekman]

To reproduce:

# run-h200-sp2-llama-8b-liger-eval.yml
type: sft
micro_batch_size: 1
exit_iteration: 3
min_iterations: 3
sequence_parallel_size: 2

deepspeed:
  zero_optimization:
    stage: 3
    offload_optimizer:
      device: cpu
  seq_parallel_communication_data_type: bf16

optimizer:
  type: cpu_adam
  learning_rate: 0.00001

model:
  type: "liger"
  #name_or_path: meta-llama/Llama-3.1-8B-Instruct
  name_or_path: TinyLlama/TinyLlama-1.1B-Chat-v1.0
  attn_implementation: sdpa

data:
  train_eval_split: [0.8, 0.2]
  sources:
    - HuggingFaceH4/ultrachat_200k:train[:100]
  cache_dir: data-cache
  dl_num_workers: 1

  max_length: 16_000

eval_interval: 1

logger:
  level: WARNING

then prepare:

git clone https://github.com/snowflakedb/ArcticTraining && cd ArcticTraining
pip install liger-kernel
pip install -e .

then run with the yaml file from the top - I changed the model to a smaller one so it's much faster.

I am using just 2 gpus and sp=2 (any gpu type should do as the model is very small).

$ arctic_training run-h200-sp2-llama-8b-liger-eval.yml --master_port 9818

I rechecked it works fine with sp=1, so it must be something that is different when sp>1.

The error to expect is:

[rank1]:   File "/code/users/stas/github/sf/arctictraining-main/arctic_training/trainer/sft_trainer.py", line 92, in loss
[rank1]:     raise ValueError(
[rank1]: ValueError: Liger-Kernel failed to compute loss (returned None) - it's known to fail with eval enabled along train steps when SP>1

To help you navigate to the most likely places where the problem comes from:

1. The first suspect code is:

ArcticTraining/arctic_training/trainer/sft_trainer.py

Lines 48 to 95 in 8f61c50

	if self.config.sequence_parallel_size == 1:
	# if model.type=liger is configured - this will use a much more efficient fused
	# logits+loss liger kernel - using significantly less gpu memory and a bit faster
	# compute (liger fused logits+loss kernel does not repeat forward during backward)
	outputs = self.model(**batch, use_cache=False)
	loss = outputs.loss
	return loss
	# Ulysses SP expectations:
	# 1. batch has `labels`` replaced with `shift_labels`` (which are already preshifted in
	# DataLoader)
	# 2. this rank deals with a seqlen dimension shard so once the loss is calculated it needs
	# to do a differentiable weighted loss average to get the grads right
	if "labels" in batch:
	raise ValueError(
	"found labels in batch - they shouldn't be there, instead shift_labels should be there - check"
	" that UlyssesSPDataLoaderAdapter has been applied to the original DataLoader object"
	)
	if "shift_labels" not in batch:
	raise ValueError(
	"shift_labels are missing from the batch - check that UlyssesSPDataLoaderAdapter has been"
	" applied to the original DataLoader object"
	)
	shift_labels = batch["shift_labels"]
	# We have 2 implementation of efficient tiled logits+loss computation.
	# 1. Liger fused cross-entropy is the fastest/most memory efficient way - liger-kernel
	# doesn't recompute forward inside backward, instead it computes the gradients in the
	# forward path.
	# 2. But liger kernel isn't implemented for all HF Transformers models, so then we fall
	# back onto our tiled logits+loss compute implementation that is almost as efficient
	# memory-wise, but which has more compute overhead before backward re-runs forward. The
	# total memory usage is very similar, but cuda cache flushes earlier if pushing close to
	# OOM than liger.
	if self.config.model.type == "liger":
	# letting liger do fused logits+loss calculation
	outputs = self.model(**batch, use_cache=False)
	loss = outputs.loss
	if loss is None:
	# XXX: not sure why this happens with SP>1 and eval-enabled, I checked shift_labels contain valid non -100 tokens - disabling fused_linear_cross_entropy=False in AutoLigerKernelForCausalLM.from_pretrained doesn't help. all works when eval is off.
	raise ValueError(
	"Liger-Kernel failed to compute loss (returned None) - it's known to fail with eval enabled along"
	" train steps when SP>1."
	)

but there is no difference between:

ArcticTraining/arctic_training/trainer/sft_trainer.py

Lines 48 to 54 in 8f61c50

	if self.config.sequence_parallel_size == 1:
	# if model.type=liger is configured - this will use a much more efficient fused
	# logits+loss liger kernel - using significantly less gpu memory and a bit faster
	# compute (liger fused logits+loss kernel does not repeat forward during backward)
	outputs = self.model(**batch, use_cache=False)
	loss = outputs.loss
	return loss

and:

ArcticTraining/arctic_training/trainer/sft_trainer.py

Lines 86 to 89 in 8f61c50

	# letting liger do fused logits+loss calculation
	outputs = self.model(**batch, use_cache=False)
	loss = outputs.loss

the only reason I have a special branch for sp>1+liger is that I need to aggregate loss across ranks, but otherwise the code is exactly the same.

Practically the sequence liger-kernel sees is divided by SP size. so often there could be a problem with a batch containing just -100s - but I checked that wasn't the case.

2. The second is the DL wrapper

ArcticTraining/arctic_training/trainer/trainer.py

Lines 246 to 262 in 8f61c50

	# wrap the DL with Ulysses one
	self.train_dataloader = UlyssesSPDataLoaderAdapter(
	self.train_dataloader,
	sp_rank=self.sp_rank,
	sp_group=self.sp_group,
	sp_world_size=self.sp_world_size,
	device=self.device,
	)
	if self.eval_dataloader is not None:
	self.eval_dataloader = UlyssesSPDataLoaderAdapter(
	self.eval_dataloader,
	sp_rank=self.sp_rank,
	sp_group=self.sp_group,
	sp_world_size=self.sp_world_size,
	device=self.device,
	)

The data is sharded and computed a chunk on each rank.

Perhaps I missed something else, but training works just fine, somehow eval breaks this. You can comment out eval_interval: 1 and it should work again.

Thank you for looking into it, @Tcc0403

[Tcc0403]

I found the issue!

In liger flce patch, there are two branches in model.forward

1. If the model is in training mode, liger calls fused linear cross entropy to skip logit materialization and only calculates loss.
2. If the model is in eval mode, liger actually calls the original forward method,
   CrossEntropy(self.lm_head(hidden_states, labels), to track full logits and loss.

So, the root cause of the problem has nothing to do with the custom ops, it's that liger doesn't check the shift_labels in the eval branch, leaving loss=None without performing the loss calculation when we only passes shift_labels in eval mode.

else:
        logits = self.lm_head(kept_hidden_states)
- 			if labels is not None:
+			if labels is not None or shift_labels is not None:
            loss = self.loss_function(
                logits=logits,
                labels=labels,
+				shift_labels=shift_labels,
                vocab_size=self.config.vocab_size,
                **kwargs,
            )

With above changes, we can run the provided recipe without any errors!

Parameter Offload - Persistent parameters statistics: param_count = 45, numel = 92160
Train iter: 1/3  33% | loss: 7.4589 | iter time: 5.583s | iter tflops: 12.2 | lr: 9.167E-06 | seqlen: 8.96K | step time: 5.355s | step tflops: 12.8 | epoch: 0 | MA 2.02 GB | Max_MA 3.73 GB | NV 8.04 GB
Eval Batches: 100%|████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00,  4.82batch/s]
 Eval iter: 1/3  33% | loss: 7.3806 | lr: 9.167E-06 | seqlen: 8.96K | epoch: 0 | MA 3.74 GB | Max_MA 5.25 GB | NV 10.31 GB
Train iter: 2/3  66% | loss: 6.7772 | iter time: 2.415s | iter tflops: 40.8 | lr: 8.333E-06 | seqlen: 11.52K | step time: 1.276s | step tflops: 77.2 | epoch: 0 | MA 2.02 GB | Max_MA 4.93 GB | NV 10.31 GB
Eval Batches: 100%|████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 10.06batch/s]
 Eval iter: 2/3  66% | loss: 6.6968 | lr: 8.333E-06 | seqlen: 11.52K | epoch: 0 | MA 3.74 GB | Max_MA 5.25 GB | NV 10.31 GB
Train iter: 3/3 100% | loss: 7.2859 | iter time: 1.918s | iter tflops: 43.2 | lr: 7.500E-06 | seqlen: 10.24K | step time: 1.189s | step tflops: 69.7 | epoch: 0 | MA 2.02 GB | Max_MA 4.87 GB | NV 10.31 GB
Eval Batches: 100%|████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 10.13batch/s]
 Eval iter: 3/3 100% | loss: 6.4933 | lr: 7.500E-06 | seqlen: 10.24K | epoch: 0 | MA 3.74 GB | Max_MA 5.25 GB | NV 10.31 GB

[sfc-gh-sbekman]

Aha! Thank you for discovering the issue and the fix, @Tcc0403!

[sfc-gh-sbekman]

The fix has been merged linkedin/Liger-Kernel#888 - waiting for liger-kernel==0.6.3 release to resolve this issue.