memprofiler

Classes

MemProfiler

MemProfiler is a powerful memory profiling utility that synthesizes numerous complementary profiling methods.

MemProfiler

A powerful memory profiling utility that synthesizes numerous complementary profiling methods.

class fts_examples.profiling.memprofiler.MemProfiler(*args, **kwargs)

MemProfiler is a powerful memory profiling utility that synthesizes numerous complementary profiling methods.

The following profiling utilities are integrated and simultaneously configured:

• FSDP2MemTracker

• cuda memory snapshot and allocator history tracking

• host-level memory tracking

• custom memory hooks (e.g. for activation checkpoint memory tracking via saved_tensors_hooks etc.)

The MemProfiler is a powerful memory profiling utility that synthesizes numerous complementary profiling methods.

The following profiling utilities are integrated and simultaneously configured:

• FSDP2MemTracker

• cuda memory snapshot and allocator history tracking

• host-level memory tracking

• custom memory hooks (e.g. for activation checkpoint memory tracking via saved_tensors_hooks etc.)

See this example for usage guidance.

Parameters:

• *args – Ignored positional arguments.

• **kwargs – Ignored keyword arguments.

memory_stats

A dictionary of dictionaries containing memory statistics.

Type:

DefaultDict

fsdp_memory_stats

A dictionary containing FSDP memory statistics.

Type:

Dict

rank

The rank of the current process.

Type:

int

module

The module being profiled.

Type:

Any

memprof_log_dir

The directory where memory profiles are saved.

Type:

str

fsdp_mem_tracker

The FSDP memory tracker.

Type:

FSDPMemTracker

saved_tensors_funcs

A list of functions used to track saved tensors.

Type:

list[Callable]

_state

The internal state of the MemProfiler.

Type:

MemProfInternalState

add_memprofiler_hooks()

Adds all hooks registered for memory tracking.

Currently, hooks are registered for the pre_forward and post_forward points. Hooks are added to the modules of the provided module and are responsible for tracking memory usage. Supported hooks are registered based on the configuration of memory_hooks in memprofiler_cfg.

After registering all hooks, this method calls exec_reset_state_hooks() to reset the state of all hooks.

Return type:

None

connect(obj_ref)

Connects the MemProfiler to the given target module.

Parameters:

obj_ref (Any) – The module to be profiled.

Return type:

None

cuda_allocator_history_snap(snap_key)

Dumps a snapshot of the CUDA memory allocator history.

The snapshot is saved to a file named cuda_alloc_rank_<snap_key>.pickle in the directory specified by memprof_log_dir.

Parameters:

snap_key (tuple) – A tuple used to identify the snapshot.

Return type:

None

done(iter_idx)

Checks if the profiling is done.

Parameters:

iter_idx (int) – The current iteration index.

Returns:

Whether the profiling is done.

Return type:

bool

dump_memory_stats()

Dumps the collected memory statistics to a pickle file and/or a yaml file.

If dump_memorystats_pickle is True, the memory statistics will be saved to a pickle file at memprof_log_dir / f”rank_{self.rank}_memory_stats.pickle”.

If dump_memorystats_yaml is True, the memory statistics will be saved to a yaml file at memprof_log_dir / f”rank_{self.rank}_memory_stats.yaml”.

If track_fsdp_mem is True, the collected FSDP memory statistics will be saved to a yaml file at memprof_log_dir / f”rank_{self.rank}_fsdp_memory_stats.yaml”.

Return type:

None

exec_reset_state_hooks()

Executes all hooks registered for resetting memory tracking state.

These hooks are responsible for resetting state (e.g. saved tensor sizes, RSS, etc.) tracked by the hooks registered in add_memprofiler_hooks().

Return type:

None

gen_snap_keys(fn_name, iter_ctx, iter_idx=None)

Generates the MemProfiler snapshot key for a given function and iteration context.

Parameters:

• fn_name (str) – The name of the function to generate a snapshot key for.

• iter_ctx (str) – The iteration context of the function to generate a snapshot key for.

• iter_idx (int | None) – The iteration index to use in the snapshot key. If None, the current iteration index will be used.

Return type:

tuple[int, int, tuple]

Returns:

A tuple of (iteration index, snapshot key). The snapshot key is a tuple of (rank, function name, iteration index, iteration context).

Note

The snapshot key format is rank.fn_name.iter_idx.iter_ctx.

init_memprof_log_dir()

Initializes the directory where MemProfiler will save profiling artifacts.

The directory is determined by the following logic: :rtype: None

1. If memprofiler_cfg.save_dir is set, use that directory.

2. Otherwise, use a directory named memprofiler in the current Lightning log_dir.

maybe_init_fsdp_mem_tracker()

Initializes FSDP2 memory tracker if track_fsdp_mem is True.

Return type:

None

static memprofilable(func)

Decorate a function to enable use with MemProfiler.

This decorator is used to enable profiling a method with MemProfiler. If the method is decorated with this decorator and the object has an enabled MemProfiler instance, the method will be profiled using the MemProfiler instance.

Parameters:

func – The function to decorate.

Returns:

The decorated function.

memprofile_ctx(fn_name)

Context manager to orchestrate profiling for a given function.

This context manager is responsible for calling the relevant subordinate context managers that manage profiling for a given function. The context managers invoked are:

1. memprofile_meta_ctx: Sets metadata for the profiling invocation.

2. memprofile_fsdp_ctx: Collects FSDP memory statistics.

3. memprofile_snap_ctx: Collects memory snapshots at the start and end of a method.

Parameters:

• memprofiler (MemProfiler) – The MemProfiler instance.

• fn_name (str) – The name of the function being profiled.

Yields:

None

memprofile_fsdp_ctx(fn_name, track_inputs_target=None)

Sets the FSDP memory tracker context manager if fsdp_mem_tracker_enabled is True.

This context manager takes care of calling the update_collect_state and reset_mod_stats methods of the MemProfiler instance as well as the track_inputs and reset methods of the FSDPMemTracker instance.

Parameters:

• memprofiler (MemProfiler) – The MemProfiler instance.

• fn_name (str) – The name of the function being profiled.

• track_inputs_target (Tuple | None) – The FSDP inputs to track. If None, no inputs will be tracked.

Yields:

None

memprofile_meta_ctx(fn_name)

Context manager to set metadata for MemProfiler invocations.

This context manager always sets the iteration state for the given function name to enable orchestration of relevant profiling invocations.

Parameters:

• memprofiler (MemProfiler) – The MemProfiler instance.

• fn_name (str) – The name of the function being profiled.

Yields:

None

memprofile_snap_ctx(fn_name)

Context manager for taking memory snapshots at the start and end of a method.

This context manager takes care of calling the snap method of the MemProfiler instance at the start and end of the context. It also handles the case where we are only collecting FSDP memory statistics.

Parameters:

• memprofiler (MemProfiler) – The MemProfiler instance.

• fn_name (str) – The name of the function being profiled.

Yields:

None

remove_memprofiler_hooks()

Removes all hooks added by MemProfiler.

Return type:

None

save_fsdp_mem_reports()

Saves the collected FSDP memory statistics to a yaml file.

For each function name and iteration in the collected FSDP memory statistics, a log file is created at memprof_log_dir / f”fsdp_mem_rank_{self.rank}_{fn_name}_{iter}.log” containing the FSDP memory statistics in both tabular and modulewise formats.

The modulewise format is displayed up to the specified fsdp_mem_track_module_depth, and the units for the memory statistics are specified by fsdp_mem_tracker_units.

If fsdp_mem_tracker_tabulate is True, the FSDP memory statistics will be displayed in a tabular format. Otherwise, a plain text format will be used.

Return type:

None

Returns:

None

snap(fn_name, iter_ctx, reset_mem_hooks=False)

Collects a memory snapshot for the given function and iteration context.

If the current iteration index is greater than the number of warmup iterations, then this function will check if the iteration index is within the profiling range. If it is, then a memory snapshot will be collected. If not, the memory profiler state for the given function and iteration context will be torn down.

Parameters:

• fn_name (str) – The name of the function to collect a memory snapshot for.

• iter_ctx (str) – The iteration context of the function to collect a memory snapshot for.

• reset_mem_hooks (bool, optional) – Whether to reset the memory hooks after collecting the snapshot. Defaults to False.

Return type:

None

teardown_prof(fn_name, iter_ctx)

Tears down profiling state for a function.

This method is responsible for:

1. Disabling profiling for the given function and iteration context.

2. Marking the function as done if it is not being profiled in either the “start” or “end” iteration contexts.

3. Removing all hooks registered by MemProfiler if the function is marked as done and retain_hooks_for_funcs is True.

Parameters:

• fn_name (str) – The name of the function to tear down profiling state for.

• iter_ctx (str) – The iteration context of the function to tear down profiling state for.

Return type:

None

update_collect_state(fn_name, iter_ctx, iter_idx=None)

Updates the MemProfiler state for a given function and iteration context.

Parameters:

• fn_name (str) – The name of the function to update the state for.

• iter_ctx (str) – The iteration context of the function to update the state for.

• iter_idx (int, optional) – The iteration index to use when generating the snapshot key. If None is provided, the current iteration index will be used.

Return type:

None

property fsdp_mem_tracker_root_module

If track_fsdp_mem is enabled, this is the root FSDP module used for FSDP2 memory tracking.

The root module must have fully_shard applied for FSDP2 memory tracking. If track_fsdp_mem is disabled, this is None.