The steps of "Training complete"

[AIxyz]

I use different num GPUs to SFT, but when I get "Training complete" for (2GPUs, 4GPUs, 8GPUs), the progress bar hasn't finished loading of steps. @exlaw

# 1 GPU
python3 train.py --config configs/example.yaml

# 2 GPUs
accelerate launch --multi_gpu --num_processes 2 --mixed_precision bf16 train.py --config configs/example.yaml

# 4 GPUs
accelerate launch --multi_gpu --num_processes 4 --mixed_precision bf16 train.py --config configs/example.yaml

# 8 GPUs
accelerate launch --multi_gpu --num_processes 8 --mixed_precision bf16 train.py --config configs/example.yaml

[AIxyz]

In finetune/src/trainer.py, I print len(self.train_dataloader) before and after self.accelerator.prepare, it changes the len(self.train_dataloader)

I guess the "total" of tqdm should be len(self.train_dataloader) after self.accelerator.prepare

    def train(self):
        """Main training loop."""
        logger.info(f"Starting training: {len(self.train_dataloader)} batches per GPU, {self.num_training_steps} total update steps")
        
        progress_bar = tqdm(total=self.num_training_steps, disable=not self.accelerator.is_local_main_process)
        
        for epoch in range(self.config.num_train_epochs):
            self.model.train()
            
            if self.train_sampler is not None:
                self.train_sampler.set_epoch(epoch)
            
            for batch in self.train_dataloader:
                with self.accelerator.accumulate(self.model):
                    loss, logs = self.train_step(batch)
                    # ……………………………………………………………………………………………………………………
        
        progress_bar.close()
        self._save_checkpoint(final=True)
        if self.wandb:
            self.wandb.finish()
        logger.info("Training complete!")