class PackedDataset(Dataset):
    """
    Performs greedy sample packing on a provided dataset.

    (...)

    A packed sample is made up of individual smaller sequence length samples jammed together
    within ``max_seq_len``. For example, if max_seq_len is 6 and there are varied
    length samples::

        tokens = [
            [S1, S1, S1, S2, S2, pad],
            [S3, S3, S4, S4, pad, pad],
            ...,
        ]

    To prevent cross-contamination, the following mask would be returned for the
    first pack in the example::

        mask = [
            [1, 0, 0, 0, 0, 0],
            [1, 1, 0, 0, 0, 0],
            [1, 1, 1, 0, 0, 0],
            [0, 0, 0, 1, 0, 0],
            [0, 0, 0, 1, 1, 0],
            [0, 0, 0, 0, 0, 1],
        ]

    The position ids would be::

        input_pos = [
            [0, 1, 2, 0, 1, 2],
            [0, 1, 0, 1, 2, 3],
            ...,
        ]

# https://github.com/NVIDIA/Megatron-LM/blob/f2f0f7bfd88fcb1243df55275988d6af52daea35/megatron/core/model_parallel_config.py#L385-L389
class ModelParallelConfig:
    ###################
    # Pipeline Parallel
    ###################
    variable_seq_lengths: bool = False
    """Support for variable sequence lengths across microbatches. Setting this communicates the size
        of tensors during pipeline parallelism communication, because of this extra overhead it
        should only be set if the sequence length varies by microbatch within a global batch.
    """

# https://github.com/NVIDIA/Megatron-LM/blob/f2f0f7bfd88fcb1243df55275988d6af52daea35/megatron/core/pipeline_parallel/schedules.py#L116-L121
def get_forward_backward_func(...):
    """
    seq_length (int, required): Sequence length of the current global batch. If this is a dual-stack
        transformer, this is the encoder's sequence length. This is ignored if variable_seq_lengths
        in the config is True. Otherwise, each microbatch in the current global batch size must use
        this sequence length.

    micro_batch_size (int, required): The number of sequences in a microbatch.

    """

# https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/pipelining/schedules.py#L505-L527
class _PipelineSchedule(ABC):
    @abstractmethod
    def step(
        self,
        *args,
        target=None,
        losses: list | None = None,
        return_outputs=True,
        loss_kwargs: dict[str, Any] | None = None,
        **kwargs,
    ):
        """
        Run one iteration of the pipeline schedule with *whole-batch* input.
        Will chunk the input into microbatches automatically, and go through the
        microbatches according to the schedule implementation.

        args: positional arguments to the model (as in non-pipeline case).
        kwargs: keyword arguments to the model (as in non-pipeline case).
        target: target for the loss function.
        losses: a list to store the losses for each microbatch.
        return_outputs: whether to return the outputs from the last stage.
        loss_kwargs: extra keyword arguments forwarded to the loss function.
        """

# https://github.com/NVIDIA/Megatron-LM/blob/f2f0f7bfd88fcb1243df55275988d6af52daea35/megatron/core/pipeline_parallel/p2p_communication.py
# Simplified
class P2PCommunicator:
    def _communicate_shapes(self, ...):
        if is_sender:
            send_prev_shape_tensor = torch.tensor(tensor_send_prev.size(), ...)
            send_next_shape_tensor = torch.tensor(tensor_send_next.size(), ...)
            send_prev_op = P2POp(isend, send_prev_shape_tensor, self.prev_rank, self.pp_group)
            send_next_op = P2POp(isend, send_next_shape_tensor, self.next_rank, self.pp_group)
            ops = [send_prev_op, send_next_op]
        else:
            recv_prev_shape_tensor = torch.empty((3,), ...)
            recv_next_shape_tensor = torch.empty((3,), ...)
            recv_prev_op = P2POp(irecv, recv_prev_shape_tensor, self.prev_rank, self.pp_group)
            recv_next_op = P2POp(irecv, recv_next_shape_tensor, self.next_rank, self.pp_group)
            ops = [recv_prev_op, recv_next_op]

        work_list = torch.distributed.batch_isend_irecv(ops)  # submit to cuda stream
        for work in work_list:
            work.wait()  # cuda stream wait (non blocking)

        if is_sender:
            return [0, 0, 0], [0, 0, 0]
        recv_prev_shape = recv_prev_shape_tensor.tolist()  # block on D2H
        recv_next_shape = recv_next_shape_tensor.tolist()  # block on D2H
        return recv_prev_shape, recv_next_shape

    def _communicate(self, ...):
        # shape
        if config.variable_seq_lengths:
            recv_prev_shape, recv_next_shape = self._communicate_shapes(...)
        else:
            recv_prev_shape, recv_next_shape = tensor_shape, tensor_shape

        # payload
        if is_sender:
            send_prev_op = P2POp(isend, tensor_send_prev, self.prev_rank, self.pp_group)
            send_next_op = P2POp(isend, tensor_send_next, self.next_rank, self.pp_group)
            ops = [send_prev_op, send_next_op]
        else:
            tensor_recv_prev = torch.empty(recv_prev_shape, ...)
            tensor_recv_next = torch.empty(recv_next_shape, ...)
            recv_prev_op = P2POp(irecv, tensor_recv_prev, self.prev_rank, self.pp_group)
            recv_next_op = P2POp(irecv, tensor_recv_next, self.next_rank, self.pp_group)
            ops = [recv_prev_op, recv_next_op]
        work_list = torch.distributed.batch_isend_irecv(ops)
        if is_sender:
            return None, None, work_list
        return tensor_recv_prev, tensor_recv_next, work_list

# https://github.com/pytorch/torchtitan/blob/4c6481182c1c3d7815e3a119a271eebead94c71a/torchtitan/trainer.py
# Over-simplified
class Trainer:
    def __init__(self, ...):
        model = model_spec.model(config.model)  # on meta device

        # Fork: how model inits
        if parallel_dims.pp_enabled:
            pp_schedule, model_parts, pp_has_first_stage, pp_has_last_stage = \
                model_spec.pipelining_fn(model, ...)
            del model
            for m in model_parts:
                m.to_empty(device)
                cast(BaseModel, m).init_weights()
                m.train()
            ensure_pp_loss_visible(parallel_dims, pp_schedule_name)
        else:
            pp_schedule, pp_has_first_stage, pp_has_last_stage = None, None, None
            model = model_spec.parallelize_fn(model, ...)
            model.to_empty(device)
            cast(BaseModel, model).init_weights()
            model.train()
            model_parts = [model]
        
        # Fork: where lm_head lives
        if isinstance(loss_fn, ChunkedLossWrapper):
            if parallel_dims.pp_enabled:
                if pp_has_last_stage:  # lm_head in PP last stage
                    loss_fn.set_lm_head(model_parts[-1].lm_head)
                    model_parts[-1]._skip_lm_head = True
                else:
                    pass  # non-last stage: no lm_head
            else:
                # lm_head in the only model part
                assert len(model_parts) == 1
                loss_fn.set_lm_head(model_parts[0].lm_head) 
                model_parts[0]._skip_lm_head = True
                

    def forward_backward_step(self, ...):
        # Fork: model forward + backward vs schedule step
        if parallel_dims.pp_enabled:
            with train_context():
                # Fork: stage-dependent args passed to schedule step
                if pp_has_last_stage:
                    targets, losses = labels, []
                else:
                    targets, losses = None, None

                if pp_has_first_stage:
                    pp_schedule.step(inputs, target=targets, losses=losses, ...)
                else:
                    pp_schedule.step(        target=targets, losses=losses, ...)

            # Fork: stage-dependent loss computation
            if pp_has_last_stage:
                assert losses is not None
                loss = sum(stack(losses)).to(device)
            else:
                loss = tensor([-1.0], device=device)  # fake value

        else:
            assert len(model_parts) == 1
            with train_context():
                pred = model_parts[0](inputs, ...)
                loss, _ = loss_fn(pred, labels, ...)
                loss.backward()
        return loss


class Validator:
    def validate(self, model_parts, ...)
        # Fork patterns similar to forward_backward_step

        if parallel_dims.pp_enabled:
            with train_context():
                if pp_has_first_stage:
                    pp_schedule.step(inputs, target=None,   losses=None,   ...)
                    loss = tensor([-1.0], device=device)
                elif pp_has_last_stage:
                    losses = []
                    pp_schedule.step(        target=labels, losses=losses, ...)
                    loss = sum(stack(losses)).to(device)
                else:
                    pp_schedule.step(        target=None,   losses=None,   ...)
                    loss = tensor([-1.0], device=device)
        else:
            ...

# https://github.com/pytorch/torchtitan/blob/4c6481182c1c3d7815e3a119a271eebead94c71a/torchtitan/distributed/pipeline_parallel.py#L428
# simplified

def _split_module(whole_model: nn.Module, modules_to_keep: set[str]) -> nn.Module:
    model = copy.deepcopy(whole_model)
    for name, m in model.named_children():
        if isinstance(m, (nn.ModuleDict, nn.ModuleList)):
            layers_to_keep: set[str] = ...
            if layers_to_keep:
                # Keep only specified layers
                if isinstance(m, nn.ModuleDict):
                    for layer_name in list(m.keys()):
                        if layer_name not in layers_to_keep:
                            del m[layer_name]
                elif isinstance(m, nn.ModuleList):
                    indices_to_keep: list[int] = ...
                    new_layers = nn.ModuleList([
                        l for i, l in enumerate(m) if i in indices_to_keep])
                    setattr(model, name, new_layers)
            else:
                # No layers from this structure needed, set to empty structure
                if isinstance(m, nn.ModuleDict):
                    setattr(model, name, nn.ModuleDict())
                elif isinstance(m, ModuleList):
                    setattr(model, name, nn.ModuleList())
        elif name not in modules_to_keep:
            # Replace with None
            setattr(model, name, None)
    return model

# https://github.com/pytorch/torchtitan/blob/4c6481182c1c3d7815e3a119a271eebead94c71a/torchtitan/models/common/decoder.py#L262
# simplified
class Decoder(BaseModel):
    """Base class for autoregressive decoder-only language models."""

    def forward(self, tokens: torch.Tensor, ...):
        # Note: `tokens` is int token ids in the first stage,
        # but becomes hidden states in later stages.
        if self.tok_embeddings is not None:
            h = self.tok_embeddings(tokens)
        else:
            h = tokens

        # Note: all stages happen to have an iterable `layers`.
        for layer in self.layers.values():
            h = layer(h, ...)

        # Note: only last stage has `norm`.
        if self.norm is not None:
            h = self.norm(h)
        
        # Note: only last stage has `lm_head`
        if self.lm_head is not None:
            output = self.lm_head(h)
        else:
            output = h

        # Note: `output` is hidden states in earlier stages,
        # but becomes logits in the last stage.
        return output

class ModelOutput(Protocol):
    @property
    def loss(self) -> Tensor: ...

class Model[InT, BatchT, OutT: ModelOutput](ABC):
    @abstractmethod
    def prepare_batch(self, inputs: InT) -> BatchT: ...

    @abstractmethod
    def forward(self, batch: BatchT) -> OutT: ...

class Trainer[InT, BatchT, OutT: ModelOutput]:
    loader: Iterator[InT]
    model: Model[InT, BatchT, OutT]

    def step(self) -> OutT:
        inputs = next(self.loader)
        batch = self.model.prepare_batch(inputs)
        out = self.model.forward(batch)
        out.loss.backward()
        return out