Implement tensor-parallel linear layers (column-parallel then row-parallel) with correct f

Super-hard

Cluster Infrastructure: 5D Parallelism, Scheduling, and Checkpoints4 / 4

Chapter 254 / 4

Cluster Infrastructure: 5D Parallelism, Scheduling, and Checkpoints

Implement tensor-parallel linear layers (column-parallel then row-parallel) with correct f

Super-hard

Implement tensor-parallel linear layers (column-parallel then row-parallel) with correct forward/backward collectives in PyTorch.

Target interface

class _CopyToModelParallel

Complete the starter code in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.

solution.pyPython

Local draft

import torch, torch.distributed as dist
from torch.autograd import Function

class _CopyToModelParallel(Function):

    @staticmethod
    def forward(ctx, x):
        raise NotImplementedError

    @staticmethod
    def backward(ctx, grad):
        raise NotImplementedError

class _ReduceFromModelParallel(Function):

    @staticmethod
    def forward(ctx, x):
        raise NotImplementedError

    @staticmethod
    def backward(ctx, grad):
        raise NotImplementedError

class ColumnParallelLinear(torch.nn.Module):

    def __init__(self, d_in, d_out, world, rank):
        raise NotImplementedError

    def forward(self, x):
        raise NotImplementedError

class RowParallelLinear(torch.nn.Module):

    def __init__(self, d_in, d_out, world, rank):
        raise NotImplementedError

    def forward(self, x_shard):
        raise NotImplementedError

class ParallelFFN(torch.nn.Module):

    def __init__(self, d_model, d_ff, world, rank):
        raise NotImplementedError

    def forward(self, x):
        raise NotImplementedError

Local draft

⌘/Ctrl↵to submit

AI review

Ready when you are

Submit your solution for a score and feedback on correctness, edge cases, and clarity. Any correct approach is accepted.

Chapter 254 / 4

Cluster Infrastructure: 5D Parallelism, Scheduling, and Checkpoints

Implement tensor-parallel linear layers (column-parallel then row-parallel) with correct f

Super-hard

Implement tensor-parallel linear layers (column-parallel then row-parallel) with correct forward/backward collectives in PyTorch.

Target interface

class _CopyToModelParallel

Complete the starter code in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.