Migrate GPT-2 to new tracer

[kwen2501]

Description

Migrated GPT-2 example to work with new tracer based pippy.

examples/hf/hf_utils.py contains utility to generate inputs for HuggingFace models.

Model architecture:

GPT2ForSequenceClassification(
  (transformer): GPT2Model(
    (wte): Embedding(50257, 768)
    (wpe): Embedding(1024, 768)
    (drop): Dropout(p=0.1, inplace=False)
    (h): ModuleList(
      (0-11): 12 x GPT2Block(
        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (attn): GPT2Attention(
          (c_attn): Conv1D()
          (c_proj): Conv1D()
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (mlp): GPT2MLP(
          (c_fc): Conv1D()
          (c_proj): Conv1D()
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
  )
  (score): Linear(in_features=768, out_features=2, bias=False)
)

Run

$ torchrun --nproc-per-node 4 pippy_gpt2.py

Output

https://gist.github.com/kwen2501/b9ed6158d8d0dc90b16824aa6abd8d72

[xw285cornell]

Does this mean the split point is right before/after this submodule (specified by FQN)? I guess it's also not super clear to me what's the meaning of SplitPoint.beginning. Also curious - tracing will "flatten" the submodule into a bunch of aten ops, so there is no longer a concept of submodule. So does this API finds the first or last node of the submodule?

[kwen2501]

BEGINNING means right before; END would mean right after.

[kwen2501]

Annotation occurs before tracing. So the module structure is still there at that time.

[H-Huang]

How come there is no corresponding END?

[H-Huang]

is input_ids in this case just a single microbatch? Or is it the entire minibatch

[kwen2501]

Entire.
When PipelineStage actually runs, it splits the batch internally before feeding to scheduler.

[H-Huang]

Looks good, just a few questions.