Function train_shared_loss

tests/unit/runtime/test_multiple_models.py:21–46

(num_models, config_dict, dtype)

def train_shared_loss(num_models, config_dict, dtype):
    hidden_dim = 64

    models = [create_model(config_dict) for _ in range(num_models)]
    data_loader = random_dataloader(model=models[0],
                                    total_samples=4,
                                    hidden_dim=hidden_dim,
                                    device=models[0].device,
                                    dtype=dtype)
    dist.barrier()
    for _, batch in enumerate(data_loader):
        losses = [m.module(batch[0], batch[1]) for m in models]
        loss = sum(l / (i + 1) for i, l in enumerate(losses))
        loss.backward()

        for m in models:
            m._backward_epilogue()

        for m in models:
            m.step()

        for m in models:
            m.optimizer.zero_grad()

    for m in models:
        m.destroy()


def train_independent_loss(num_models, config_dict, dtype):