Sequence Parallelism, memory usage question

[EthanChen1234]

the paper, "Sequence Parallelism: Long Sequence Training from System Perspective"

for the tensor parallelism, the activations and trainable weight consume memory.
the activations, contain 4BLH/N and BLH. the trainable weight should be 4H^2/N + 4H^2/N = 8H^2/N.

I'm confused, the 32H^2/N memory use.

[EthanChen1234]

I'm still waiting.

[FrankLeeeee]

Hi, it 32H^2/N refers to the space occupied by M2 which is the weight of the linear layer. It is 32 because we count in the gradients, momentum, variance assuming using Adam optimizer instead of the weight alone.

[0oshowero0]

Hi, I'm also confused about this. In my calcualtion (according to ZeRO), the memory usage for M2 should be:
model parameters: $\frac{4H^2}{N}$
gradients: $\frac{4H^2}{N}$
optimizer states: $\frac{4H^2}{N} * 3$ (parameters, momentum, variance).

So it should be $\frac{20H^2}{N}$ for the 1st/2st linear and $\frac{40H^2}{N}$ together.