Add TODO for supporting other kinds of parallelism

Signed-off-by: John St John <jstjohn@nvidia.com>

Author: jstjohn

sub-packages/bionemo-llm/src/bionemo/llm/utils/callbacks.py

@@ -89,6 +89,8 @@ def data_parallel_rank(self) -> int:
     @property
     def should_write_predictions(self) -> bool:
         """Returns the context parallel rank."""
+
+        # TODO: handle expert parallelism and other kinds of parallelism
         return self.save_all_model_parallel_ranks or (
             parallel_state.is_pipeline_last_stage()
             and parallel_state.get_tensor_model_parallel_rank() == 0