[BUG] How to drop some batches entirely to avoid calculating backpropagation while still updating the model for the rest

Describe the bug
Basically, when using deepspeed inside the transformer trainer.py file, I would like to drop certain batches when they satisfy certain conditions, and avoid the calculation of gradients and backpropagations to save time and memory. However, currently the code would just lag on self.accelerator.backward(loss), (for both skipped and unskipped samples) which never finishes.

To Reproduce

We modify this function to satisfy our goal, highlighted by ####. Code segments without '####' is the original code, which runs fine on its own.

(this code is basically modified from transformer.trainer, inside trainer.py)

def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:
model.train()
inputs = self._prepare_inputs(inputs)

    if is_sagemaker_mp_enabled():
        loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
        return loss_mb.reduce_mean().detach().to(self.args.device)

    with self.compute_loss_context_manager():
        loss, all_skips = self.compute_loss(model, inputs) 
        #### SKIP LOGIC ####
        if meets_skip_condition:
           dummy_tensor = torch.tensor(0.0, device=self.accelerator.device, requires_grad=True)
           dummy_loss = dummy_tensor * 0.0 
           loss = dummy_loss
           model.zero_grad()
           self.accelerator.backward(loss)
           return loss
        #### END SKIP LOGIC ####

    if self.args.n_gpu > 1:
        loss = loss.mean()  # mean() to average on multi-gpu parallel training

    if self.use_apex:
        with amp.scale_loss(loss, self.optimizer) as scaled_loss:
            scaled_loss.backward()
    else:
        self.accelerator.backward(loss)

    print(f"Finished self.accelerator.backward(loss) at rank {dist.get_rank()}")

Expected behavior
I would like the function to continue executing, and still compute gradients for samples that are not dropped.

System info (please complete the following information):

• OS: Linux
• GPU count and types one machine with x8 A100
• Python version 3.10

Launcher context
Are you launching your experiment with the deepspeed launcher, MPI, or something else?
the command is:
deepspeed llava/train/train_mem.py
--deepspeed /home/lanxy/LLaVA/scripts/zero2_mod.json
--model_name_or_path lmsys/vicuna-7b-v1.5
--version v1
--data_path /home/lanxy/Dataset/LLaVA-Finetune/llava_v1_5_mix665k.json
--image_folder /home/lanxy/Dataset/LLaVA-Finetune
--vision_tower openai/clip-vit-large-patch14-336
--mm_projector_type mlp2x_gelu
--pretrain_mm_mlp_adapter /home/lanxy/llava-v1.5-7b-pretrain/mm_projector.bin
--mm_vision_select_layer -2
--mm_use_im_start_end False
--mm_use_im_patch_token False
--image_aspect_ratio pad
--group_by_modality_length True
--bf16 True
--output_dir ./checkpoints/llava-v1.5-7b${RUN_NUM}
--num_train_epochs 1
--per_device_train_batch_size 1
--per_device_eval_batch_size 1
--gradient_accumulation_steps 1
--save_strategy "steps"
--save_steps 10000
--save_total_limit 1
--learning_rate 2e-5
--weight_decay 0.
--warmup_ratio 0.03
--lr_scheduler_type "cosine"
--logging_steps 1
--tf32 True
--model_max_length 2048
--gradient_checkpointing True
--dataloader_num_workers 4
--lazy_preprocess True
--report_to none
--run_name llava-v1.5-7b${RUN_NUM}

Docker context
NA

Additional context
NA

Start by reproducing the reported behavior in transformer trainer.py using the shown training_step modification and the provided DeepSpeed launcher configuration. Focus on the path through self.accelerator.backward(loss) for skipped and unskipped batches. Done means skipped batches do not hang while gradients still update the remaining samples.