Sequence parallelism is not supported for NemotronHForCausalLM

[aayush-sarvam]

Is there a way use sequence parallelism for NemotronHForCausalLM model? I am getting:

  File "/data/TTS/aayush/nemo-rl/nemo_rl/models/policy/dtensor_policy_worker.py", line 352, in __init__
    self.model = _parallelize_model(
                 ^^^^^^^^^^^^^^^^^^^
  File "/data/TTS/aayush/nemo-rl/nemo_rl/models/dtensor/parallelize.py", line 529, in _parallelize_model
    return _parallelize_nm5_h(
           ^^^^^^^^^^^^^^^^^^^
  File "/data/TTS/aayush/nemo-rl/nemo_rl/models/dtensor/parallelize.py", line 415, in _parallelize_nm5_h
    assert not sequence_parallel, (
           ^^^^^^^^^^^^^^^^^^^^^
AssertionError: Sequence parallelism is not supported for NemotronHForCausalLM

[guyueh1]

@joyang-nv can you confirm if this is a current limitation of dtensor?

[joyang-nv]

@aayush-sarvam , can you let us know model name? I think you are using dtensor policy v1. We are focusing on v2 (which is built on top of NeMo Automodel) now. But we need model name to verify on v2.