modalities.utils.profilers package

Submodules

modalities.utils.profilers.batch_generator module

class modalities.utils.profilers.batch_generator.DataTypeEnum(value)

Bases: LookupEnum

bfloat16 = torch.bfloat16

float32 = torch.float32

int64 = torch.int64

class modalities.utils.profilers.batch_generator.DatasetBatchGeneratorIF

Bases: ABC

get_dataset_batch()

Return type:

DatasetBatch

class modalities.utils.profilers.batch_generator.RandomDatasetBatchGenerator(dims, data_type, min_val, max_val)

Bases: DatasetBatchGeneratorIF

Parameters:

• dims (dict[str, int])

• data_type (DataTypeEnum)

• min_val (int)

• max_val (int)

get_dataset_batch()

Return type:

DatasetBatch

class modalities.utils.profilers.batch_generator.RandomDatasetBatchGeneratorConfig(**data)

Bases: BaseModel

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

• dims (dict[str, int])

• data_type (DataTypeEnum)

• min_val (int)

• max_val (int)

data_type: DataTypeEnum

dims: dict[str, int]

max_val: int

min_val: int

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

modalities.utils.profilers.modalities_profiler module

class modalities.utils.profilers.modalities_profiler.CustomComponentRegisterable(component_key, variant_key, custom_component, custom_config)

Bases: object

Parameters:

• component_key (str)

• variant_key (str)

• custom_component (type)

• custom_config (type)

component_key: str = <dataclasses._MISSING_TYPE object>

custom_component: type = <dataclasses._MISSING_TYPE object>

custom_config: type = <dataclasses._MISSING_TYPE object>

variant_key: str = <dataclasses._MISSING_TYPE object>

class modalities.utils.profilers.modalities_profiler.InstantiationModel(**data)

Bases: BaseModel

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

steppable_component (Annotated[SteppableComponentIF, <modalities.config.pydantic_if_types.PydanticThirdPartyTypeIF object at 0x7f3743ed2190>])

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

steppable_component: Annotated[SteppableComponentIF]

class modalities.utils.profilers.modalities_profiler.ModalitiesProfiler

Bases: object

static export_profiling_results(profiler_context_manager, trace_output_path, memory_output_path, summary_output_path, global_rank, local_rank, profiled_ranks)

Export profiling results to specified output paths if the current rank is in profiled_ranks.

Return type:

None

Parameters:

• profiler_context_manager (profile)

• trace_output_path (Path)

• memory_output_path (Path)

• summary_output_path (Path)

• global_rank (int)

• local_rank (int)

• profiled_ranks (list[int])

Args:

profiler_context_manager (profile): The profiler context manager. trace_output_path (Path): Path to save the Chrome trace. memory_output_path (Path): Path to save the memory timeline. summary_output_path (Path): Path to save the summary table. global_rank (int): The global rank of the current process. local_rank (int): The local rank of the current process. profiled_ranks (list[int]): List of ranks to profile.

static profile(steppable_component, num_total_steps, profile_context_manager, show_progress=False)

Profile a steppable component using the provided profiler context manager.

Return type:

None

Parameters:

• steppable_component (SteppableComponentIF)

• num_total_steps (int)

• profile_context_manager (profile)

• show_progress (bool)

Args:

steppable_component (SteppableComponentIF): The steppable component to profile. num_total_steps (int): Total number of steps to run. profile_context_manager (profile): The profiler context manager. show_progress (bool): Whether to show a progress bar. Defaults to False.

class modalities.utils.profilers.modalities_profiler.ModalitiesProfilerStarter

Bases: object

Starter class to run profiling either in single process or distributed mode.

static run_distributed(config_file_path, num_measurement_steps, num_wait_steps, num_warmup_steps, experiment_root_path, profiled_ranks, experiment_id=None, custom_component_registerables=None)

Run distributed profiling using the Modalities Profiler. This method is primarily intended to run large-scale profiling e.g., for model training.

Args:

config_file_path (Path): Path to the configuration file. num_measurement_steps (int): Number of measurement steps for profiling. num_wait_steps (int): Number of wait steps before profiling starts. num_warmup_steps (int): Number of warmup steps before measurement starts. experiment_root_path (Path): Root path to store experiment results. profiled_ranks (list[int]): List of ranks to profile. experiment_id (str, optional): Experiment ID. If None, it will be generated. Defaults to None. custom_component_registerables (list[CustomComponentRegisterable], optional): List of custom

components to register. Defaults to None.

Parameters:

• config_file_path (Path)

• num_measurement_steps (int)

• num_wait_steps (int)

• num_warmup_steps (int)

• experiment_root_path (Path)

• profiled_ranks (list[int])

• experiment_id (str | None)

• custom_component_registerables (list[CustomComponentRegisterable] | None)

static run_single_process(config_file_path, num_measurement_steps, num_wait_steps, num_warmup_steps, experiment_root_path, experiment_id=None, custom_component_registerables=None)

Run single process profiling using the Modalities Profiler. This method is primarily intended for quick profiling experiments (e.g., single modules such as custom modules vs Pytorch equivalents) on a single GPU. This type of profiling can be seen as a middle ground between distributed profiling and single kernel profiling (e.g., using Nsight Compute).

Args:

config_file_path (Path): Path to the configuration file. num_measurement_steps (int): Number of measurement steps for profiling. num_wait_steps (int): Number of wait steps before profiling starts. num_warmup_steps (int): Number of warmup steps before measurement starts. experiment_root_path (Path): Root path to store experiment results. experiment_id (str, optional): Experiment ID. If None, it will be generated. custom_component_registerables (list[CustomComponentRegisterable], optional): List of custom

components to register. Defaults to None.

Parameters:

• config_file_path (Path)

• num_measurement_steps (int)

• num_wait_steps (int)

• num_warmup_steps (int)

• experiment_root_path (Path)

• experiment_id (str | None)

• custom_component_registerables (list[CustomComponentRegisterable] | None)

modalities.utils.profilers.steppable_component_configs module

class modalities.utils.profilers.steppable_component_configs.SteppableForwardPassConfig(**data)

Bases: BaseModel

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

• model (Annotated[Module, <modalities.config.pydantic_if_types.PydanticThirdPartyTypeIF object at 0x7f3743ed0e90>])

• dataset_batch_generator (Annotated[DatasetBatchGeneratorIF, <modalities.config.pydantic_if_types.PydanticThirdPartyTypeIF object at 0x7f3743ed1e90>])

• loss_fn (Annotated[Loss, <modalities.config.pydantic_if_types.PydanticThirdPartyTypeIF object at 0x7f3743ed1790>] | None)

dataset_batch_generator: Annotated[DatasetBatchGeneratorIF]

loss_fn: Optional[Annotated[Loss]]

model: Annotated[Module]

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

modalities.utils.profilers.steppable_components module

class modalities.utils.profilers.steppable_components.SteppableForwardPass(model, dataset_batch_generator, loss_fn=None)

Bases: SteppableComponentIF

A steppable component that performs a forward pass on the model using batches from the dataset batch generator. Optionally computes the loss if a loss function is provided. The component is used for profiling.

Initializes the SteppableForwardPass component.

Args:

model (nn.Module): The model to perform the forward pass on. dataset_batch_generator (DatasetBatchGeneratorIF): The dataset batch generator to provide input batches. loss_fn (Loss, optional): The loss function to compute the loss. Defaults to None.

Parameters:

• model (Module)

• dataset_batch_generator (DatasetBatchGeneratorIF)

• loss_fn (Loss | None)

step()

Performs a single step of the forward pass and computes the loss if a loss function is provided.

Return type:

None

modalities.utils.profilers.steppable_components_if module

class modalities.utils.profilers.steppable_components_if.SteppableComponentIF

Bases: object

step()

Return type:

None

Module contents