Support complicated use cases with TiedLayerSpec

[limjcst]

I want to reuse a composed module in the pipeline. For example, the following MyModule has a member linear, which is also a module.

class MyModule(torch.nn.Module):
    def __init__(self, n_in: int, n_out: int):
        super().__init__()
        self.linear = torch.nn.Linear(n_in, n_out)
        self.layer_norm = torch.nn.LayerNorm(n_out)

    def forward(self, data: torch.Tensor) -> torch.Tensor:
        hidden = self.linear(data)
        hidden = self.layer_norm(hidden)
        return hidden

MyModule.linear.weight should be synchronized among related ranks. As a result, I add linear.weight to TiedLayerSpec.tied_weight_attr.
BTW, I generate the whole tied_weight_attr by the following instruction.

tied_weight_attr = [name for name, p in layer.named_parameters() if p.numel() > 1]

However, the builtin getattr used by PipelineModule fails to find a nested attribute like linear.weight.
Hence, this PR first extends the builtin getattr to a recursive version PipelineModule._recursive_getattr, accessing each attribute segment one by one.

Meanwhile, the order of tied weights matters in synchronization. This PR suggests to sort tie_keys in PipelineModule._index_tied_modules to avoid hanging.

[tohtana]

Thank you @limjcst for the contribution! This is a significant improvement.

[limjcst]

nv-accelerate-v100 failed, raising "invalid command 'bdist_wheel'". However, this job succeeded in another run.

Note that the failed job used "cached wheel-0.46.1-py3-none-any.whl.metadata".
@agronholm suggested setuptools >= 70.1 for wheel==0.46.0. Is this suggestion meaningful for DeepSpeed?

pypa/wheel#660 (comment)

[loadams]

nv-accelerate-v100 failed, raising "invalid command 'bdist_wheel'". However, this job succeeded in another run.

Note that the failed job used "cached wheel-0.46.1-py3-none-any.whl.metadata". @agronholm suggested setuptools >= 70.1 for wheel==0.46.0. Is this suggestion meaningful for DeepSpeed?

pypa/wheel#660 (comment)

Thanks @limjcst - I saw this failure on another PR and will take a look and merge the fixes into your PR when ready.

[loadams]

nv-accelerate-v100 failed, raising "invalid command 'bdist_wheel'". However, this job succeeded in another run.
Note that the failed job used "cached wheel-0.46.1-py3-none-any.whl.metadata". @agronholm suggested setuptools >= 70.1 for wheel==0.46.0. Is this suggestion meaningful for DeepSpeed?
pypa/wheel#660 (comment)

Thanks @limjcst - I saw this failure on another PR and will take a look and merge the fixes into your PR when ready.

@limjcst - it looks like the wheel team has yanked the problemabic wheel version so tests should be passing again.

[agronholm]

Nevertheless, upgrading to setuptools >= 70.1 should prevent any future issues. I noticed that this project lacks a standard pyproject.toml which would be able to specify build requirements.

[loadams]

Nevertheless, upgrading to setuptools >= 70.1 should prevent any future issues. I noticed that this project lacks a standard pyproject.toml which would be able to specify build requirements.

@agronholm - yes, we have a PR for one here which we will prioritize merging as we know this is needed.