`nemo_automodel.components.distributed.parallel_styles`#

Module Contents#

Classes#

`ColwiseParallelLora`
`RowwiseParallelLora`

Functions#

translate_to_lora

API#

class nemo_automodel.components.distributed.parallel_styles.ColwiseParallelLora#

Bases: torch.distributed.tensor.parallel.ColwiseParallel

_partition_linear_fn(name, module, device_mesh)#

class nemo_automodel.components.distributed.parallel_styles.RowwiseParallelLora#

Bases: torch.distributed.tensor.parallel.RowwiseParallel

_partition_linear_fn(name, module, device_mesh)#

nemo_automodel.components.distributed.parallel_styles.translate_to_lora(plan)#