Important

You are viewing the NeMo 2.0 documentation. This release introduces significant changes to the API and a new library, NeMo Run. We are currently porting all features from NeMo 1.0 to 2.0. For documentation on previous versions or features not yet available in 2.0, please refer to the NeMo 24.07 documentation.

Migrate Parallelism Configurations from NeMo 1.0 to NeMo 2.0#

In NeMo 2.0, model parallelisms are configured via the MegatronStrategy.

NeMo 1.0 (Previous Release)#

In NeMo 1.0, all model parallelisms were set under the model section in the YAML configuration file.

model:
  tensor_model_parallel_size: 2
  pipeline_model_parallel_size: 1
  virtual_pipeline_model_parallel_size: null
  context_parallel_size: 2
  sequence_parallel: True
  ...

NeMo 2.0 (New Release)#

In NeMo 2.0, these settings are controlled from the MegatronStrategy.

from nemo import lightning as nl

strategy = nl.MegatronStrategy(
    tensor_model_parallel_size=2,
    pipeline_model_parallel_size=1,
    virtual_pipeline_model_parallel_size=None,
    context_parallel_size=2,
    sequence_parallel=True,
)

trainer = nl.Trainer(
   strategy=strategy,
   ...
)

Migrate Parallelism Configurations#

  1. Locate and remove all parallelism configurations in your NeMo 1.0 YAML config file.

  2. Add the following import to your Python script:

    from nemo import lightning as nl

  3. Create a MegatronStrategy and provide the parallelism configuration:

    strategy = nl.MegatronStrategy(
    tensor_model_parallel_size=2,
    pipeline_model_parallel_size=1,
    virtual_pipeline_model_parallel_size=None,
    context_parallel_size=2,
    sequence_parallel=True,
)

  4. Add the strategy to your Trainer (see Trainer migration guide):

    trainer = nl.Trainer(
    ...
    strategy=strategy,
    ...
)
On this page