fix: remove tie weight check

[RayenTian]

What does this PR do ?

This pull request removes the tie weight check in the DTensor Worker, addressing a previous limitation with training models that utilize tie_word_embeddings when tensor parallel size (tp_size) > 1.

Issues

closes #684

Test Result

For remove NRL_SKIP_TIED_WEIGHT_CHECK env variable

Experiments were run with meta-llama/Llama-3.2-1B, Qwen/Qwen2.5-1.5B-Instruct and google/gemma-2-2b-it. With the truncate rate set within a reasonable range, both tp_size=1 and tp_size=2 produced comparable reward results, indicating that training proceeds as expected. This demonstrates consistency in model behavior across different tensor parallel configurations, confirming correct handling of tied word embeddings during training.

Main setup:

• policy.generation.temperature=1.0
• policy.max_total_sequence_length=2048
• cluster.gpus_per_node=8

Llama-3.2-1B

Qwen2.5-1.5B-Instruct

google/gemma-2-2b-it

Test with:

• policy.dynamic_batching.enabled=true
• policy.sequence_packing.enabled=false
• policy.train_global_batch_size=128

For remove model.config.tie_word_embeddings in nemo_rl/models/dtensor/parallelize.py

Experiments were conducted using both the Qwen/Qwen2.5-1.5B-Instruct and Qwen/Qwen2.5-7B-Instruct models, representing the original tied and untied weight configurations, respectively. I disabled Qwen's optimized parallel plan in PARALLIZE_FUNCTIONS, forcing both models to utilize the HP plan for testing.

• The results indicate that, after removing the model.config.tie_word_embeddings check, the HP plan and the optimized plan produce identical outcomes.

Qwen/Qwen2.5-1.5B-Instruct

• This model uses tied word embeddings by default.
• Current results are from a short test of 10 steps.

Qwen/Qwen2.5-7B-Instruct

• This model uses untied word embeddings by default.
• Current results are from a short test of 10 steps.

More Custom Test

Some additional tests were also conducted here.
For models like Qwen/Qwen2.5-1.5B-Instruct, which by default use tied word embeddings, we forcibly disabled the tie_word_embedding setting. The results showed significant anomalies in both the token_mult_prob_error and the reward. We suspect this may be because vLLM still keeps the embeddings tied on its side, so when you call update_weights, there might be some issues (in practice, vLLM may be using either the embed or the lm_head weight, rather than truly untying them as in training). This causes the token_mult_prob_error to behave abnormally.

For models like Qwen/Qwen2.5-7B-Instruct, which by default do not use tied embeddings, we forcibly enabled tie_word_embedding. The results showed large differences in reward, but the token_mult_prob_error remained stable. This might be because, originally, the embed and lm_head weights were different due to being untied; forcing a tie essentially removes the original lm_head weights. When we call update_weights to vLLM, its lm_head is also overwritten with the embed weights, so the token_mult_prob_error remains normal. However, since we forcibly replaced the original weights, the reward becomes abnormal.

Additional Notes on Gemma Model Testing

• gemma-3-1b-it: This model sets kv_head=1, which is currently not compatible with tensor parallelism (TP=2). As a result, it could not be tested in a TP=2 configuration.

[yuki-97]

Thanks @RayenTian for verifying and removing these things!

Can you also check if we can also remove model.config.tie_word_embeddings in HF TP plan in nemo_rl/models/dtensor/parallelize.py?

[yuki-97]

One more thing, can you also search issues/227 in the repo and remove (or update) them?

[SahilJain314]

can you please attach the token_mult_prob_error plots to the description? (we can't always see the errors when we just look at convergence plots)

[RayenTian]

can you please attach the token_mult_prob_error plots to the description? (we can't always see the errors when we just look at convergence plots)

Thank you for the suggestion. I have added the plots as requested.

[SahilJain314]

thanks for rebasing, lgtm now.