Usage of nn.Identity

[Andrei-Aksionov]

A quick question (and here I rely on your expertise, since I don't have one regarding to using many-many GPUs).

Here:

torchtitan/torchtitan/models/common/decoder.py

Lines 124 to 138 in 3c81104

	def forward(
	self,
	tokens: torch.Tensor,
	attention_masks: AttentionMasksType | None = None,
	positions: torch.Tensor | None = None,
	):
	# passthrough for nonexistent layers, allows easy configuration of pipeline parallel stages
	h = self.tok_embeddings(tokens) if self.tok_embeddings is not None else tokens
	for layer in self.layers.values():
	h = layer(h, self.freqs_cis, attention_masks, positions)
	h = self.norm(h) if self.norm is not None else h
	output = self.output(h) if self.output is not None else h
	return output

Why not assign nn.Identity in the init method instead?

        self.norm = config.norm.build() or nn.Identity()
        self.output = config.output.build() or nn.Identity()

so the forward function will look like:

    def forward(
        self,
        tokens: torch.Tensor,
        attention_masks: AttentionMasksType | None = None,
        positions: torch.Tensor | None = None,
    ):
        # passthrough for nonexistent layers, allows easy configuration of pipeline parallel stages
        h = self.tok_embeddings(tokens) if self.tok_embeddings is not None else tokens

        for layer in self.layers.values():
            h = layer(h, self.freqs_cis, attention_masks, positions)

        h = self.norm(h)
        return self.output(h)

Yes, it looks cleaner. But this is not the main point.
With it we can drop if statements and thus reduce code branching.
In theory, without code branching:

• torch.compile should have an easier life
• easier synchronization for each GPU in case these two layers are created conditionally per-rank

---

But I have a very strong feeling that there is a clear reason why nn.Identity is not used.
I just don't see it.

[tianyu-l]

I think the reason is because distributed checkpointing and Pipeline Parallel composability.

In current way, there's no ambiguity in state dict of model across multiple PP ranks.

If using nn.Identity (e.g. HuggingFace transformers), you need something like https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/transformers_modeling_backend/model.py#L673

[Andrei-Aksionov]

Oh, I see. Didn't think about it.
Thanks for the answer.

I guess something could be done via a custom class so nothing breaks, e.g.:

class BlackHoleIdentity(nn.Identity):
    def forward(self, input, *args, **kwargs):
        return input

    def _save_to_state_dict(self, *args, **kwargs):
        pass

    def _load_from_state_dict(self, *args, **kwargs):
        pass

    def reset_parameters(self):
        pass

    @property
    def weight(self):
        return torch.empty(0)

    @property
    def bias(self):
        return torch.empty(0)

But feels like over-complication (plus might introduce silent bugs).
None sounds fine 🙂