cerebras.modelzoo.common.dump_context.DumpContext#
- class cerebras.modelzoo.common.dump_context.DumpContext(outdir, model, buffer_steps=None)[source]#
Bases:
contextlib.ContextDecorator
A debug utility context manager. When provided with a torch.nn.Module, the resulting context manager can be entered to enable dumping of all module forward and backward outputs to a npz, for comparing numerics between implementations.
Sets up global module hoooks to either dump intermediate activations on CPU/GPU or name the traced tensors for correlating with debug dumps on CS2.
The recursive name of the torch.nn.Module is memoized, and the output of FWD and BWD of each module is saved as keys in a .npz file.
- Parameters
outdir (str) – Where to output dumps_{i}.npz
model (torch.nn.Module) – root module to name its children
buffer_steps (int) – If given, flush to a new .npz file after this many steps
Methods
Uninstall the hooks installed during enable_collection, disabling further dump collection.
Install the hooks defined during setup_hooks, enabling the collection of the dumps.
Write all dump buffers out to disk.
Define hooking functions on the given torch.nn.Module, but don't install them.
- setup_hooks(model)[source]#
Define hooking functions on the given torch.nn.Module, but don’t install them.
- Parameters
model – torch.nn.Module that serves as the root for recursive names
- enable_collection()[source]#
Install the hooks defined during setup_hooks, enabling the collection of the dumps.