modalities.models.parallelism package

Submodules

modalities.models.parallelism.pipeline_parallelism module

modalities.models.parallelism.pipeline_parallelism_configs module

modalities.models.parallelism.stages_generator module

class modalities.models.parallelism.stages_generator.GPT2LLMStagesGenerator(num_model_layers, input_layer_equivalence=1, output_layer_equivalence=1)[source]

Bases: StagesGenerator

Parameters:

  • num_model_layers (int)

  • input_layer_equivalence (int)

  • output_layer_equivalence (int)

class modalities.models.parallelism.stages_generator.StagesGenerator(num_model_layers, input_layer_equivalence=1, output_layer_equivalence=1)[source]

Bases: ABC

Parameters:

  • num_model_layers (int)

  • input_layer_equivalence (int)

  • output_layer_equivalence (int)

get_stages(num_layers_per_stage, pp_dims)[source]

Generate FQNs for each stage in a model.

Return type:

list[list[str]]

Parameters:

  • num_layers_per_stage (int)

  • pp_dims (int)

Args:

num_layers_per_stage (int): Number of layers per stage. pp_dims (int): Number of pipeline parallel dimensions.

Returns:

list[list[str]]: A list containing FQNs for each stage.

modalities.models.parallelism.stages_generator_configs module

class modalities.models.parallelism.stages_generator_configs.FQNsPerStageGeneratorConfig(**data)[source]

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.

model_config: ClassVar[ConfigDict] = {}

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

class modalities.models.parallelism.stages_generator_configs.GPT2LLMStagesGeneratorConfig(**data)[source]

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:

  • num_model_layers (Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Strict(strict=True), Ge(ge=1)])])

  • input_layer_equivalence (Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Strict(strict=True), Ge(ge=1)])])

  • output_layer_equivalence (Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Strict(strict=True), Ge(ge=1)])])

input_layer_equivalence: Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Strict(strict=True), Ge(ge=1)])]
model_config: ClassVar[ConfigDict] = {}

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

num_model_layers: Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Strict(strict=True), Ge(ge=1)])]
output_layer_equivalence: Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Strict(strict=True), Ge(ge=1)])]

Module contents