What jumps out at me is
Before stream s, allocated:  2978816...
After stream s before graph g, allocated:  13397522944

The model evidently isn’t that big (about 3 MB) but somehow after warmup the allocated memory shoots up to 13 GB.
I’m not surprised the reserved memory gets pushed up to ~15 GB (it increases to accommodate the high-water mark of the built graph). But after backward(), the graph should be torn down, and the allocated memory should settle back to whatever the model params and optimizer states need which should be several times the model’s param set size, in other words 3 MB x N where N is a small integer.
Can you add
gc.collect()
torch.cuda.empty_cache()

just above
print('After stream s before graph g, allocated: ', torch.cuda.memory_allocated('cuda'), '; reserved: ', torch.cuda.memory_reserved('cuda'))

and report the values?

can I save the graph and run it on a different machine

No, there’s currently no way to serialize and deserialize graphs.

more examples of CUDA graphs with PyTorch

The master docs are the best resource.
https://pytorch.org/docs/master/notes/cuda.html#cuda-graphs
Some internal people are adding cuda graphs to DeepLearningExamples/PyTorch at master · NVIDIA/DeepLearningExamples · GitHub 1, but I don’t know how far that effort has progressed.