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 Nsys Profiling from NeMo 1.0 to NeMo 2.0#

In NeMo 2.0, the way to configure Nsys profiling has changed from a YAML configuration to using a dedicated callback. This guide will help you migrate your Nsys profiling setup.

NeMo 1.0 (Previous Release)#

In NeMo 1.0, Nsys profiling was configured in the YAML configuration file.

model:
    nsys_profile:
        enabled: False
        start_step: 10  # Global batch to start profiling
        end_step: 10 # Global batch to end profiling
        ranks: [0] # Global rank IDs to profile
        gen_shape: False # Generate model and kernel details including input shapes

NeMo 2.0 (New Release)#

In NeMo 2.0, Nsys profiling is configured using the NsysCallback class. Here’s how to set it up:

from nemo import lightning as nl
from nemo.lightning.pytorch.callbacks import NsysCallback

trainer = nl.Trainer(
    ...
    callbacks=[NsysCallback(
        enabled=False,
        start_step=10,
        end_step=10,
        ranks=[0],
        gen_shape=False
    )]
)

Migration Steps#

  1. Remove the nsys_profile section from your YAML config file.

  2. Add the following import to your Python script:

    from nemo.lightning.pytorch.callbacks import NsysCallback
    
  3. When creating your Trainer, add NsysCallback to the callbacks list:

    trainer = nl.Trainer(
        ...
        callbacks=[NsysCallback(
            enabled=False,
            start_step=10,
            end_step=10,
            ranks=[0],
            gen_shape=False
        )]
    )
    
  4. Adjust the parameters in NsysCallback to match your previous YAML configuration.