[VPP] Defer embedding initialization sync until all local model chunks are built

[SonglinLife]

Describe the bug

With pipeline parallelism and virtual pipeline parallelism enabled, each pipeline rank constructs its local virtual-pipeline model chunks sequentially.

However, tied-embedding initialization synchronization is performed while an individual chunk is still being constructed.

During construction, LanguageModule.setup_embeddings_and_output_layer() prepares duplicated embedding weights and immediately performs the initial tied-embedding all_reduce。The first and last pipeline ranks reach the same embedding-group collective while constructing different virtual stages。

For example, with PP=6 and VPP=3, construction proceeds approximately as:

PP0: build VP0 -> embedding all_reduce (wait) -> build VP1 -> build VP2
PP5: build VP0 -> build VP1 -> build VP2 -> embedding all_reduce

PP0 cannot continue constructing VP1 and VP2 until PP5 finishes constructing its earlier chunks and reaches the matching collective.

Both supported VPP construction paths iterate over local VP chunks sequentially:

• training.py
• dist_utils.py

Steps/Code to reproduce bug

Use a model configuration with:

• pipeline_model_parallel_size > 1
• virtual_pipeline_model_parallel_size > 1
• Multi-Token Prediction (MTP) enabled, for example --mtp-num-layers 1

Expected behavior

Initial embedding synchronization should not block construction of the remaining local VP chunks. Each pipeline rank should be able to continue building its local virtual stages without waiting for other ranks reach the matching collective.

PP0: build VP0 -> build VP1 -> build VP2 -> embedding all_reduce
PP5: build VP0 -> build VP1 -> build VP2 -> embedding all_reduce

Proposed change

Defer tied-embedding initialization synchronization until all local VPP model chunks have been constructed, then perform the synchronization before continuing model setup.