Calculating FLOPs and training memory per GPU

[arghavan-kpm]

Hi,

Would you please give some instructions about the way you calculated FLOPs for a single forward pass during training and training memory per GPU? Thanks.

[devzhk]

Hi,

We used fvcore from facebook research to automatically estimate FLOPs. For the GPU memory usage, we reported the maximum GPU memory occupied by tensors with torch.cuda.max_memory_allocated.

[arghavan-kpm]

Thanks so much for your response. I tried fvcore but I keep getting this error message: "RuntimeError: Detected that you are using FX to torch.jit.trace a dynamo-optimized function. This is not supported at the moment." Could you please share the details of where do you call FlopCountAnalysis() and what do you pass as the input args?

[devzhk]

Just disable torch.compile().