[WIP] LM Workload

[rka97]

This is for the LM workload.

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[fsschneider]

I might be wrong here, but I just randomly saw that we are using a default MLP expansion factor of $4$ (i.e.,

algorithmic-efficiency/algoperf/workloads/lm/lm_pytorch/workload.py

Line 40 in 5c85c7e

expand=4, # MLP expansion factor

).

However, it seems to me, that we are also using a gated linear unit, which uses an initial expansion that is twice as large (since it is split into the gate and the value), i.e.

algorithmic-efficiency/algoperf/workloads/lm/lm_pytorch/plainlm_model.py

Lines 28 to 30 in 5c85c7e

	self.fc1 = nn.Linear(dim, 2 * hidden_dim, bias=False)
	self.fc2 = nn.Linear(hidden_dim, dim, bias=False)
	self.glu = nn.GLU(dim=2)

I believe that the common expansion factor of $4$ is for models that use non-gated activation functions (e.g., original Transformer with ReLU or GPT-2 with GELU). To keep the number of parameters constant, people commonly divide by $2/3$ for gated versions. This keeps the total number of parameters in the MLP the same (some small approximation error due to the non-integer multiplication).

If you agree, we should either change the default to $4/3$, or adjust the computation of the hidden_dim in the MLP to account for this scaling.

@rka97 @priyakasimbeg @Niccolo-Ajroldi

[Niccolo-Ajroldi]

@fsschneider agree, we should definitely adjust the expansion factor as you suggested. This was probably an oversight, as in plainLM we rescale when using GLU.