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

disable_collection

Uninstall the hooks installed during enable_collection, disabling further dump collection.

enable_collection

Install the hooks defined during setup_hooks, enabling the collection of the dumps.

flush

Write all dump buffers out to disk.

setup_hooks

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.

disable_collection()[source]#

Uninstall the hooks installed during enable_collection, disabling further dump collection.

flush()[source]#

Write all dump buffers out to disk.

previous

cerebras.modelzoo.common.dump_context

next

cerebras.modelzoo.common.half_dtype