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