fix bug when using DP in trl, the batch size of input and output dism…

[kaixuanliu]

No description provided.

[kaixuanliu]

Steps to reproduce the bug:

git clone https://github.com/huggingface/trl.git
cd trl
git checkout 3ef9faf257
pip install .
export CUDA_VISIBLE_DEVICES=0,1,2
pytest -sv -rA tests/slow/test_sft_slow.py::SFTTrainerSlowTester::test_train_offloading_0_trl_internal_testing_tiny_LlamaForCausalLM_3_2

it will fail and return error:

def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        """
        Compute training loss and additionally compute token accuracies
        """
        mode = "train" if self.model.training else "eval"
        (loss, outputs) = super().compute_loss(
            model, inputs, return_outputs=True, num_items_in_batch=num_items_in_batch
        )
        if mode == "train":
            # When using padding-free, the attention_mask is not present in the inputs, instead we have cu_seq_lens_q,
            # cu_seq_lens_k, and max_length_k, max_length_q and position_ids.
            if "attention_mask" in inputs:
                num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item()
            elif "position_ids" in inputs:
                local_num_tokens = torch.tensor(inputs["position_ids"].size(1), device=inputs["position_ids"].device)
                num_tokens_in_batch = self.accelerator.gather_for_metrics(local_num_tokens).sum().item()
            else:
                raise ValueError("Expected 'attention_mask' or 'position_ids' in inputs.")
            self._total_train_tokens += num_tokens_in_batch
        self._metrics[mode]["num_tokens"] = [self._total_train_tokens]

        # Compute token accuracy if we have labels and if the model is not using Liger (no logits)
        if "labels" in inputs and not self.args.use_liger_kernel:
            shift_logits = outputs.logits[..., :-1, :].contiguous()
            shift_labels = inputs["labels"][..., 1:].contiguous()

            # Get predictions
            predictions = shift_logits.argmax(dim=-1)

            # Create mask for non-padding tokens (assuming ignore_index is -100)
            mask = shift_labels != -100

            # Calculate accuracy only on non-padding tokens
>           correct_predictions = (predictions == shift_labels) & mask
E           RuntimeError: The size of tensor a (2) must match the size of tensor b (6) at non-singleton dimension 0

It crashes as num_items_in_batch in L3837 is a 1-D tensor, and it cannot be scattered to multi-gpus successfully, hence although the input bs=6 in L3839, the output bs will be 2, and hence the test case fails.

[kaixuanliu]

@zach-huggingface, @SunMarc and @qgallouedec, pls help review

[SunMarc]

Thanks ! Can you add a test that cover this specific case ?

[kaixuanliu]

@SunMarc , Hi thx for advice. I think the existing one is OK for this case:
pytest -sv -rA tests/trainer/test_trainer.py::TrainerIntegrationTest::test_num_batches_in_training_with_gradient_accumulation
I added related assertion in latest commit. Pls help check if it is OK.

[yao-matrix]

@kaixuanliu , CI has failed cases, pls take a look

[kaixuanliu]

@yao-matrix , Updated the code and the failed case passed. I also double checked the failed case on my own machine. @SunMarc Can you help review again? thx!

[kaixuanliu]

@SunMarc Hi, this is a 2 weeks ago PR, can you help review it? Many thanks!

[kaixuanliu]

@qgallouedec ,Hi, can you help review? Thx.

[yao-matrix]

@qgallouedec, could you help review this PR?

[regisss]

actual_bs could be defined inside the IF block at line 3819 since it is only used there

[kaixuanliu]

@regisss ，thx for the review! I put actual_bs here is because inputs will be deleted in L3798. It is also OK to put the delete and free memory operation later. Have updated the code.

[qgallouedec]

Shouldn't this be fixed in trl instead?

[qgallouedec]

I'm not quite sure to understand because you use 3 devices, but don't run the test in a distributed manner?

[kaixuanliu]

@qgallouedec Hi, I think it is a corner case issue that is not properly processed in transformers. Although it can be avoided in trl level or other upper application level, it is best to handle it in their common base level(transformers). DP maybe a dated approach, but since it is not deprecated formally, it's best to fix it. WDYT?

[SunMarc]

Sorry for the delay, I have a left a few questions !

[SunMarc]

Instead of assert, let's just raise RuntimeError instead. Also the error msg don't help that much, is there an actionnable step here for the users to fix the issue ?

[SunMarc]

do we really need those checks as we didn't need until now ? I feel like the isse was if num_items_in_batch, not with the labels or loss bs actually

[SunMarc]

we are already average the loss somewhere else no ?

[SunMarc]

happy to have that but I ran the test you told me and it passed without this PR, maybe i'm doing something wrong ? pytest -sv -rA tests/trainer/test_trainer.py::TrainerIntegrationTest::test_num_batches_in_training_with_gradient_accumulation

[kaixuanliu]

Hi @SunMarc , you may need to revert trl to commit 3ef9faf257 as my comment above to reproduce. Anyway, I think it make sense as @qgallouedec mentioned that when using DP, it's best to use accelerate launch, so we can close this PR.

[qgallouedec]

If you want to use DP, you should launch the training with accelerate launch, not directly with python (or pytest).
When not using DP, make sure that only one device is visible (CUDA_VISIBLE_DEVICES=0)