modalities.models.parallelism package

Submodules

modalities.models.parallelism.pipeline_parallelism module

class modalities.models.parallelism.pipeline_parallelism.ComponentSelectorFromPipeline

Bases: object

static select(pipeline, selection_type)

Selects a component from the pipeline based on the selection type.

Return type:

Any

Parameters:

• pipeline (Pipeline)

• selection_type (PipelineSelectionTypes)

class modalities.models.parallelism.pipeline_parallelism.Pipeline(pp_stage, model_part, pp_schedule=None)

Bases: object

Parameters:

• pp_stage (PipelineStage)

• model_part (Module)

• pp_schedule (PipelineScheduleSingle | None)

property is_first_pp_stage: bool

property is_last_pp_stage: bool

property model_part: Module

property pp_schedule: PipelineScheduleSingle | None

property pp_stage: PipelineStage

class modalities.models.parallelism.pipeline_parallelism.PipelineFactory

Bases: object

Pipeline factory class to create pipelined models.

static get_pipeline(pp_stage, model_part, pp_schedule=None)

Return type:

Pipeline

Parameters:

• pp_stage (PipelineStage)

• model_part (NNModel)

• pp_schedule (PipelineScheduleSingle | None)

static get_scheduled_pipeline(loss_fn, pp_schedule_name, batch_size, microbatch_size, pp_degree, pipeline)

Return type:

Pipeline

Parameters:

• loss_fn (Loss)

• pp_schedule_name (str)

• batch_size (int)

• microbatch_size (int)

• pp_degree (int)

• pipeline (Pipeline)

static get_staged_pipeline(whole_model, stages_generator, device_mesh, local_rank, pp_schedule_name, num_layers_per_stage)

Return type:

Pipeline

Parameters:

• whole_model (NNModel)

• stages_generator (StagesGenerator)

• device_mesh (DeviceMesh)

• local_rank (int)

• pp_schedule_name (str)

• num_layers_per_stage (int)

class modalities.models.parallelism.pipeline_parallelism.PipelineSelectionTypes(value)

Bases: Enum

Enum for pipeline selection types.

MODEL_PART = 'MODEL_PART'

PP_SCHEDULE = 'PP_SCHEDULE'

PP_STAGE = 'PP_STAGE'

modalities.models.parallelism.pipeline_parallelism_configs module

class modalities.models.parallelism.pipeline_parallelism_configs.ComponentSelectorFromPipelineConfig(**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:

• pipeline (Annotated[Pipeline, <modalities.config.pydantic_if_types.PydanticThirdPartyTypeIF object at 0x7f3743ed2010>])

• selection_type (PipelineSelectionTypes)

model_config: ClassVar[ConfigDict] = {}

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

pipeline: Annotated[Pipeline]

selection_type: PipelineSelectionTypes

class modalities.models.parallelism.pipeline_parallelism_configs.FQNsPerStageGeneratorConfig(**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.

model_config: ClassVar[ConfigDict] = {}

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

class modalities.models.parallelism.pipeline_parallelism_configs.PipelineConfig(**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:

• pp_stage (Annotated[PipelineStage, <modalities.config.pydantic_if_types.PydanticThirdPartyTypeIF object at 0x7f3743ed20d0>])

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

• pp_schedule (Annotated[Pipeline, <modalities.config.pydantic_if_types.PydanticThirdPartyTypeIF object at 0x7f3743ed2010>] | None)

model_config: ClassVar[ConfigDict] = {}

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

model_part: Annotated[Module]

pp_schedule: Optional[Annotated[Pipeline]]

pp_stage: Annotated[PipelineStage]

class modalities.models.parallelism.pipeline_parallelism_configs.ScheduledPipelineConfig(**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:

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

• pp_schedule_name (str)

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

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

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

• pipeline (Annotated[Pipeline, <modalities.config.pydantic_if_types.PydanticThirdPartyTypeIF object at 0x7f3743ed2010>])

batch_size: Annotated[int]

loss_fn: Annotated[Loss]

microbatch_size: Annotated[int]

model_config: ClassVar[ConfigDict] = {}

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

pipeline: Annotated[Pipeline]

pp_degree: Annotated[int]

pp_schedule_name: str

class modalities.models.parallelism.pipeline_parallelism_configs.StagedPipelineConfig(**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:

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

• stages_generator (Annotated[StagesGenerator, <modalities.config.pydantic_if_types.PydanticThirdPartyTypeIF object at 0x7f3743ed1f50>])

• device_mesh (Annotated[DeviceMesh, <modalities.config.pydantic_if_types.PydanticThirdPartyTypeIF object at 0x7f3743ed1c10>])

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

• pp_schedule_name (str)

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

device_mesh: Annotated[DeviceMesh]

local_rank: Annotated[int]

model_config: ClassVar[ConfigDict] = {}

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

num_layers_per_stage: Annotated[int]

pp_schedule_name: str

stages_generator: Annotated[StagesGenerator]

whole_model: Annotated[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)

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)

Bases: ABC

Parameters:

• num_model_layers (int)

• input_layer_equivalence (int)

• output_layer_equivalence (int)

get_stages(num_layers_per_stage, pp_dims)

Generate FQNs for each stage in a GPT-2 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)

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)

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]

model_config: ClassVar[ConfigDict] = {}

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

num_model_layers: Annotated[int]

output_layer_equivalence: Annotated[int]

Module contents