Repository: hpcaitech/ColossalAI
File: examples/language/bert/finetune.py (Python)

Code before:
        for batch in dataloader:
            batch = move_to_cuda(batch)
            labels = batch["labels"]
            batch_size = batch["input_ids"].shape[0]
            if hasattr(booster.plugin, "stage_manager") and booster.plugin.stage_manager is not None:
                pg_mesh = booster.plugin.pg_mesh
                pp_group = booster.plugin.pp_group
                current_pp_group_ranks = pg_mesh.get_ranks_in_group(pp_group)
                current_rank = dist.get_rank()

                batch = iter([batch])

                outputs = booster.execute_pipeline(batch,
                                                   model,
                                                   criterion,
                                                   optimizer,
                                                   return_loss=True,
                                                   return_outputs=True)

                if booster.plugin.stage_manager.is_last_stage():
                    val_loss = outputs["loss"]

                    #TODO get merged output
                    #logits = outputs["outputs"].logits
                    logits = outputs["outputs"][0].logits
                    logits = logits.repeat((2, 1))
                    ####

                    accum_loss.add_(val_loss)

                    if num_labels > 1:
                        preds = torch.argmax(logits, axis=1)
                    elif num_labels == 1:
                        preds = logits.squeeze()

                    dist.broadcast(preds, src=current_rank, group=pp_group)
                    dist.broadcast(val_loss, src=current_rank, group=pp_group)

                    metric.add_batch(predictions=preds, references=labels)
                elif current_rank in current_pp_group_ranks:
                    val_loss = torch.empty((1,), device=get_current_device())
                    preds = torch.empty((batch_size,), dtype=torch.int64, device=get_current_device())

                    dist.broadcast(preds, src=current_pp_group_ranks[-1], group=pp_group)
                    dist.broadcast(val_loss, src=current_pp_group_ranks[-1], group=pp_group)

                    accum_loss.add_(val_loss)
                    metric.add_batch(predictions=preds, references=labels)

            else:
                batch = move_to_cuda(batch)

Review comment (suggestion):
looks like this line results in "ValueError: Mismatch in the number of predictions (64) and references (32)". removing the repeat works

Code after:
        for batch in dataloader:
            batch = move_to_cuda(batch)
            labels = batch["labels"]
            batch_size = batch["input_ids"].shape[0]
            if hasattr(booster.plugin, "stage_manager") and booster.plugin.stage_manager is not None:
                pg_mesh = booster.plugin.pg_mesh
                pp_group = booster.plugin.pp_group
                current_pp_group_ranks = pg_mesh.get_ranks_in_group(pp_group)
                current_rank = dist.get_rank()
                #TODO pass dataloader to execute_pipeline directly
                batch = iter([batch])
                outputs = booster.execute_pipeline(batch,
                                                   model,
                                                   criterion,
                                                   optimizer,
                                                   return_loss=True,
                                                   return_outputs=True)

                if booster.plugin.stage_manager.is_last_stage():
                    val_loss = outputs["loss"]

                    logits = outputs["outputs"]["logits"]

                    accum_loss.add_(val_loss)

                    if num_labels > 1:
                        preds = torch.argmax(logits, axis=1)
                    elif num_labels == 1:
                        preds = logits.squeeze()

                    dist.broadcast(preds, src=current_rank, group=pp_group)
                    dist.broadcast(val_loss, src=current_rank, group=pp_group)

                    metric.add_batch(predictions=preds, references=labels)
                elif current_rank in current_pp_group_ranks:
                    val_loss = torch.empty((1,), device=get_current_device())
                    preds = torch.empty((batch_size,), dtype=torch.int64, device=get_current_device())

                    dist.broadcast(preds, src=current_pp_group_ranks[-1], group=pp_group)
                    dist.broadcast(val_loss, src=current_pp_group_ranks[-1], group=pp_group)

                    accum_loss.add_(val_loss)
                    metric.add_batch(predictions=preds, references=labels)

            else:
                batch = move_to_cuda(batch)
                outputs = model(**batch)
                val_loss, logits = outputs[:2]
                accum_loss.add_(val_loss)

                if num_labels > 1: