[Wav2Vec2ForCTC] multi-GPU training's backward hangs
Author: LEECHOONGHO
Pull request: no
Labels: name: bug

### System Info

ubuntu@gpu-1:~$ transformers-cli env

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

- `transformers` version: 4.16.2
- Platform: Linux-4.15.0-143-generic-x86_64-with-glibc2.10
- Python version: 3.8.5
- PyTorch version (GPU?): 1.11.0+cu102 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@patrickvonplaten

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Execution Code : 
`$ python train.py task=wav2vec2_train task.available_gpu_ids=\'4,5,6,7\' task.project_name=wav2vec2_asr hydra.job.chdir=False`

```
### train.py

from utils.task_manager import TaskManager

@hydra.main(version_base=None, config_path=os.path.join("configs"), config_name="train", )
def hydra_main(configs: DictConfig) -> None:
    
    number_of_available_gpus = len(eval(configs.task.available_gpu_ids)) if not isinstance(eval(configs.task.available_gpu_ids), int) else 1
    
    if number_of_available_gpus > 1:
        os.environ['MASTER_ADDR'] = configs.task.master_addr
        os.environ['MASTER_PORT'] = configs.task.master_port
        os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
        os.environ["CUDA_VISIBLE_DEVICES"] = configs.task.available_gpu_ids
    
    torch.manual_seed(int(configs.task.seed))
    
    configs.task.world_size = number_of_available_gpus
    
    task_manager = TaskManager.train(configs.task)



if __name__ == '__main__':
    hydra_train_init()
    hydra_main()
```

```
### utils.taskmanager.py

class TaskManager:
    def __init__(self, config, device=0):
        self.config = config
        self.device = device
        
    @classmethod
    def train(cls, config):
        
        cls = TaskManager(config)
        
        if config.world_size > 1:

            mp.spawn(
                cls._train,
                args=(
                    config,
                ),
                nprocs=min(torch.cuda.device_count(), config.world_size),
                join=True,
            )
        else:
            cls._train(cls.device, config)
            
        return cls
  def _train(self, rank, config):
        if self.config.world_size > 1:
            torch.cuda.set_device(rank)
            dist.init_process_group("nccl", rank=rank, world_size=self.config.world_size)
        
        model, optimizer, scheduler, scaler, processor, train_dataloader, eval_dataloader = self._prepare_training(rank)
        
        step = 1
        
        print('Start Training')
        
        for epoch in range(config.num_train_epochs):

            train_dataloader.batch_sampler.set_epoch(epoch)

            for i, inputs in enumerate(train_dataloader):                                 

                inputs = {k: v.cuda(rank) if torch.is_tensor(v) else v for k, v in inputs.items()}
                with autocast(enabled=config.use_fp16):
                    outputs = model(**inputs)
                                
                scaler.scale(outputs.loss).backward()
                
                if (i+1) % config.gradient_accumulation_step == 0 or i == len(train_dataloader):
                                
                    scaler.unscale_(optimizer)

                    scheduler.step()

                    scaler.step(optimizer)

                    scaler.update()

                    optimizer.zero_grad()
                    model.zero_grad()
                    
                    step += 1
                                
                    if rank == 0:
                        self._eval(rank, config, model, eval_dataloader, step)
                        self._log_train_metrics({'Train loss' : outputs.loss}, step)
                        self._save_model(model, step)

    def _log_train_metrics(self, metrics, step):
        if step % self.config.logging_steps == 0:
            self._log_metrics(metrics, step)
        
    def _eval(self, rank, config, model, eval_dataloader, step):
        if step % self.config.eval_steps == 0:
            model.eval()
            cers = 0.
            loss = 0.
            with torch.no_grad():
                for inputs in eval_dataloader:                                 
                    inputs = {k: v.cuda(rank) if torch.is_tensor(v) else v for k, v in inputs.items()}
                    outputs = model(**inputs)
                    
                    cers += self.compute_metrics(outputs.logits, inputs['labels'])['wer']
                    loss += outputs.loss
                
                avg_cer = cers / len(eval_dataloader)
                avg_loss = loss / len(eval_dataloader)
                
                metrics = {'Eval avg_loss' : avg_loss, 'Eval avg_cer' : avg_cer}
                
                self._log_metrics(metrics, step)
                
            model.train()        
        
        
    def _log_metrics(self, metrics, step):
        wandb.log(data=metrics, step=step, commit=False)
        
    def _save_model(self, model, step):
        if step % self.config.save_steps == 0:
            model.save_pretrained(os.path.join(self.config.saved_model_path, self.config.project_name , '_' + str(step)), save_config=True)

    def _prepare_training(self, rank):
        
        if rank == 0:
            self._get_vocab_dict()
        
        processor = self._get_processor()
        
        self.processor = processor
        
        train_dataloader, eval_dataloader = self._get_dataloader(processor, rank)
        
        model = self._get_model(processor)
        
        model.cuda(rank)
        
        
        if self.config.world_size > 1:
            model = DDP(model, device_ids=[rank], output_device=rank, find_unused_parameters=True)
        
        model.train()        
        
        optimizer = transformers.AdamW(
            params=model.parameters(),
            lr=self.config.learning_rate,
            betas=eval(self.config.betas),
            eps=self.config.eps,
            weight_decay=self.config.weight_decay,
        )

        scheduler = transformers.get_scheduler(
            name='constant_with_warmup', 
            optimizer=optimizer, 
            num_warmup_steps=self.config.warmup_steps, 
            num_training_steps=len(train_dataloader) * self.config.num_train_epochs
        )
        
        scaler = GradScaler(enabled=self.config.use_fp16)
        
        if rank == 0:
            self._set_wandb(model)
            
        self._get_wer_metric()
        
        return model, optimizer, scheduler, scaler, processor, train_dataloader, eval_dataloader
    
    
    def _set_wandb(self, model):
        
        wandb.init(project=self.config.project_name, entity="xinapse")
        
        wandb.config = model.config
        
                
    def _get_vocab_dict(self):
        
        text_manager = TextManager(self.config)
                
        if not os.path.exists(self.config.vocab_dict_path):
            print(f'Vocab file not found in {self.config.vocab_dict_path}. make vocab dictionary from texts in {os.path.join(self.config.dataset_path, self.config.manifest_file_name)}')
            text_manager.make_vocab()
        
    def _get_processor(self):
        
        tokenizer = Wav2Vec2CTCTokenizer(
            self.config.vocab_dict_path, 
            unk_token=self.config.unk_token, 
            pad_token=self.config.pad_token, 
            word_delimiter_token=self.config.word_delimiter_token
        )
        
        feature_extractor = Wav2Vec2FeatureExtractor(
            feature_size=1, 
            sampling_rate=self.config.sampling_rate, 
            padding_value=0.0, 
            do_normalize=True, 
            return_attention_mask=False
        )
        processor = Wav2Vec2Processor(
            feature_extractor=feature_extractor, 
            tokenizer=tokenizer
        )
        
        return processor

    def _load_data(self):
        
        df = pd.read_pickle(
            os.path.join(self.config.dataset_path, self.config.manifest_file_name)
        )

        train_df, eval_df = train_test_split(
            df, test_size=self.config.num_eval_data
        )

        train_df.reset_index(drop=True, inplace=True)
        eval_df.reset_index(drop=True, inplace=True)
        
        return train_df, eval_df
        
    def _get_dataset(self, processor):
        
        train_df, eval_df = self._load_data()
        
        train_dataset = STTDataset(self.config, processor, train_df)
        eval_dataset = STTDataset(self.config, processor, eval_df)
        
        return train_dataset, eval_dataset
        
    def _get_collator(self, processor):
        
        data_collator = DataCollatorCTCWithPadding(processor=processor, padding=True)
        
        return data_collator
    
    def _get_dataloader(self, processor, rank):
        
        train_dataset, eval_dataset = self._get_dataset(processor)
        data_collator = self._get_collator(processor)
        
        train_sampler = DistributedBucketSampler(
            train_dataset,
            self.config.batch_size,
            eval(self.config.bucket_audio_lengths),
            num_replicas=self.config.world_size,
            rank=rank,
            shuffle=True
        )

        train_dataloader = torch.utils.data.DataLoader(
            train_dataset,
            shuffle=(train_sampler is None),
            num_workers=self.config.num_worker, 
            collate_fn=data_collator,
            pin_memory=True,
            batch_sampler=train_sampler
        )

        if rank == 0:
            eval_dataloader = torch.utils.data.DataLoader(
                eval_dataset,
                shuffle=False,
                num_workers=self.config.num_worker,
                batch_size=self.config.batch_size,
                drop_last=False,
                collate_fn=data_collator
            )
        else:
            eval_dataloader = None
            
        return train_dataloader, eval_dataloader
        
    def _get_model(self, processor):
        
        model = Wav2Vec2ForCTC.from_pretrained(
            self.config.pretrained_model_path,
            local_files_only=True,
            attention_dropout=0.1,
            hidden_dropout=0.1,
            feat_proj_dropout=0.0,
            mask_time_prob=0.05,
            layerdrop=0.1,
            ctc_loss_reduction="mean", 
            pad_token_id=processor.tokenizer.pad_token_id,
            vocab_size=len(processor.tokenizer)
        )
        
        model.freeze_feature_encoder()
        model.config.ctc_zero_infinity = True
        model.config.apply_spec_augment = True
                
        return model
```

### Expected behavior

I wrote my own training script to train wav2vec2 asr model with my custom dataset.



Executing train script in single-GPU (setting number of CUDA_VISIBLE_DEVICES=0) works fine.
But When I use multi-GPU with DDP, training progress hangs at `scaler.scale(outputs.loss).backward()` and GPU usage freezes like below

Thanks in advance,
Have a great day.

![image](https://user-images.githubusercontent.com/44384060/177467387-f899b125-9c63-4fcc-a3fa-73671c1a07dc.png)


Comments:
I don't think we support training with `hydra` in transformers - could you maybe try to use our official examples for Wav2Vec2: https://github.com/huggingface/transformers/blob/main/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py  instead.

Thank you for your reply.
Problem was in wandb's system monitoring process.

When I set wandb.init not to monitor system and set layer_drop=0.0, training works fine.

But the training loss freezes(?) after lr scheduler's warmup step(500).

Thank you. 

Ah yeah, the reason might be that you're using `gradient_checkpointing` and `layer_drop` at the same time. I actually don't recommend using `layer_drop` for Wav2Vec2 training