search

foundation-model-stack / fms-fsdp

🚀 Efficiently (pre)training foundation models with native PyTorch features, including FSDP for training and SDPA implementation of Flash attention v2.
https://pytorch.org/docs/stable/fsdp.html
Apache License 2.0
116 stars 18 forks source link

add FLOP counter #29

Open lchu-ibm opened 4 months ago

lchu-ibm commented 4 months ago

add a flop counter to the code with a bool flag.

it is already available in the flop_counter branch but will require some extra work to prettify it and integrate it with a flag.

low priority. to be visited in the future.