[BUG] Sequence Parallel(Ulysses) Training Gradient Scaling Issue

When training a language model (LM) with DeepSpeed's Sequence Parallel (Ulysses), it's typical to get a cross-entropy loss for each rank. To compute the gradients accurately, as I understand it, an in-place division by tensor_to_allreduce.div_ followed by an all-reduce operation is necessary.

process_group = self.dp_process_group if process_group is None else process_group
..
tensor_to_allreduce.div_(dist.get_world_size(group=process_group) / float(self.sequence_parallel_size))

Without performing tensor_to_allreduce.div_, the gradient would be scaled by the sequence parallel size, resulting in much higher gradients than expected. In an effort to address this, I've reflected this change in this commit, but looking at the current main code, it seems like the divide is always set to false from this commit, so it appears tensor_to_allreduce.div_ might not be applied correctly.

Alternatively, would it be acceptable to just divide the loss by SEQUENCE_PARALLEL_WORLD_SIZE and then perform a backward operation?

• example code

for step, batch in enumerate(dataloader):            
    if ENABLE_DS_SEQUENCE_PARALLEL:    
        # get sub-sequence for sequence parallel
        seq_length = batch['input_ids'].size(1)                        
        assert seq_length % SEQUENCE_PARALLEL_WORLD_SIZE == 0
        sub_seq_length = seq_length // SEQUENCE_PARALLEL_WORLD_SIZE
        sub_seq_start = SEQUENCE_PARALLEL_RANK * sub_seq_length
        sub_seq_end = (SEQUENCE_PARALLEL_RANK + 1) * sub_seq_length            
        # move to device [B, T/SP]
        input_ids = batch['input_ids'][:, sub_seq_start:sub_seq_end].to(device)                    
        attention_mask = batch['attention_mask'][:, sub_seq_start:sub_seq_end].to(device)                    
        labels = batch['labels'][:, sub_seq_start:sub_seq_end].to(device)
        position_ids = torch.arange(seq_length).unsqueeze(0)
        position_ids = position_ids[:, sub_seq_start:sub_seq_end].to(device)           
    else:
        # move to device [B, T]
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)   
        position_ids = None
                    
    #forward() method
    loss = model_engine(input_ids=input_ids, attention_mask=attention_mask, labels=labels, position_ids=position_ids)

    #############################################runs backpropagation    
    model_engine.backward(loss/SEQUENCE_PARALLEL_WORLD_SIZE) #Would it be acceptable to calculate the loss in this alternative way?
    ###############################################

    #weight update
    model_engine.step()

Start by reading the referenced gradient reduction code in deepspeed/runtime/zero/stage_1_and_2.py and compare it with the behavior introduced by PR 4957. Use the example sequence-parallel training loop to verify gradient scaling across ranks, including the proposed loss division. Done means the correct scaling behavior is established and the issue's question is resolved with a reproducible validation.