LLMs-from-scratch

LLMs-from-scratch copied to clipboard

i just check out the code of appendix-A/01_main-chapter-code /DDP-script.py,how about adding

from torch.profiler import profile
with profile() as prof:
    #the main function training code
if rank == 0:
    print("exporting trace")
    prof.export_chrome_trace("trace_ddp_simple.json")

than we can see the tracing profiling json file in google Chrome