[core][distributed] simplify code to support pipeline parallel

[youkaichao]

to minimize the line of code change for a model to support pipeline parallel.

in the model weight loading part, just add:

                if name not in params_dict:
                    # in pipeline parallelism, we may have layers that are not
                    # present on this rank
                    continue

(the benefit is we don't need to introduce extra indentation of the try-except code)

in the layer construction part, just add:

        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda: LlamaDecoderLayer(config=config,
                                      cache_config=cache_config,
                                      quant_config=quant_config))

hopefully, with these change, code for a model to support pp will be easier, and the review can also be easier.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only trigger fastcheck CI to run, which consists only a small and essential subset of tests to quickly catch errors with the flexibility to run extra individual tests on top (you can do this by unblocking test steps in the Buildkite run).

Full CI run is still required to merge this PR so once the PR is ready to go, please make sure to run it. If you need all test signals in between PR commits, you can trigger full CI as well.

To run full CI, you can do one of these:

• Comment /ready on the PR
• Add ready label to the PR
• Enable auto-merge.

🚀

[DarkLight1337]

Perhaps it would be better to move the make_layers function into the utils file inside models/?

[robertgshaw2-redhat]

This could silence some hard to track down bugs when loading more complex state dicts (in quantized case)

Could we try to check if the name corresponds to a layer not on the device b/c of PP?

[youkaichao]

this is quite difficult, as load_weights does not have layer information.

[youkaichao]

come up with a workaround in 61fa242. PTAL!

[youkaichao]

Perhaps it would be better to move the make_layers function into the utils file inside models/?

fixed in 347399e

[andoorve]

I noticed CUDA out of memory error on the basic correctness tests here. Is it reproducible locally? I think this change shouldn't cause that so it's possibly a flaky test?

[youkaichao]

/ready

[youkaichao]

@andoorve finally figured it out, it is because lru cache stores a reference of the model, and then fails the gc system :(

[youkaichao]

merge first to unblock the following models support.

[yudian0504]

:-) "xx.11".startswith("xx.1")

[youkaichao]

sorry for the bug, and thanks for pointing it out so quickly! please take a look at #6446 .