core.resharding.refit

Module Contents

Functions

swap_model_weights	Orchestrate weight swap/refit.
reshard_model_weights	Reshard and copy model weights from src_model to target_model using service.

Data

RefitBackendName

API

core.resharding.refit.RefitBackendName

None

core.resharding.refit.swap_model_weights(

src_model: megatron.core.models.common.language_module.language_module.LanguageModule,

target_model: megatron.core.models.common.language_module.language_module.LanguageModule,

refit_method: Union[core.resharding.refit.RefitBackendName, core.resharding.copy_services.base.CopyService],

)

Orchestrate weight swap/refit.

• refit_method can be:

  • a string backend name (one of the supported refit backends), or

  • a CopyService instance.

core.resharding.refit.reshard_model_weights(

src_model: megatron.core.models.common.language_module.language_module.LanguageModule,

target_model: megatron.core.models.common.language_module.language_module.LanguageModule,

service: core.resharding.copy_services.base.CopyService,

)

Reshard and copy model weights from src_model to target_model using service.