[tx] Fix lora_A not reinitialized after clearing lora adapter and support seeds in LoraConfig

[raulchen]

Summary

• Fix gradient flow bug when reusing adapter slots after deletion
• Rename update_adapter_config → init_lora_adapter and clear_adapter_config → clear_lora_adapter

Problem

When an adapter slot is cleared and reused, both lora_A and lora_B are zero, resulting in zero gradients during training.

Solution

init_lora_adapter now fully initializes the adapter (lora_A with he_uniform, lora_B with zeros) instead of assuming lora_A ha
s valid random values.

Test plan

• Added test_adapter_reuse_initializes_lora_adapter to verify proper initialization on slot reuse
• Run existing LoRA training tests

[gemini-code-assist]

Code Review

This pull request effectively addresses a critical bug where LoRA adapter weights (lora_A) were not reinitialized upon reuse, causing training failures. The solution is robust: lora_A is now correctly reinitialized using he_uniform, and lora_B is explicitly zeroed out. The function update_adapter_config has been aptly renamed to reinitialize_adapter to better reflect its new, comprehensive functionality.

The necessary changes to support this fix, such as storing rngs on the model and updating the ModelForCausalLM protocol, have been implemented correctly. Additionally, the change to prevent loading rngs from checkpoints in load_safetensors is a valuable improvement for model state management.

The new test, test_adapter_reuse_reinitializes_lora_adapter, is well-constructed and thoroughly validates the fix by simulating the exact scenario of adapter reuse that was causing the issue.

Overall, this is a high-quality contribution that resolves a significant issue and improves the codebase's clarity and correctness. The changes are clean, well-tested, and follow best practices.

[raulchen]

also de-flakey this test that was also introduced in the last PR

[pcmoritz]

/gemini review

[gemini-code-assist]

Code Review

This pull request effectively resolves a critical bug where reusing LoRA adapter slots would lead to zero gradients due to uninitialized weights. The core fix, which re-initializes lora_A with he_uniform within the newly renamed init_lora_adapter function, is well-implemented and correct. The accompanying changes, such as renaming functions for clarity and ensuring the model's RNGs are accessible, are logical and well-executed. The addition of a dedicated test case, test_adapter_reuse_initializes_lora_adapter, is excellent as it directly validates the fix. I've included a few minor suggestions to improve code consistency and readability. Overall, this is a high-quality contribution.

[pcmoritz]

/gemini review

[gemini-code-assist]

Code Review

This pull request effectively addresses a critical bug where LoRA adapter weights were not re-initialized upon reuse, which could halt training. The solution of moving the weight initialization into the newly renamed init_lora_adapter function is correct and well-implemented. The renaming of functions improves code clarity, and the addition of a new test case, test_adapter_reuse_initializes_lora_adapter, is excellent for ensuring this bug does not regress. Furthermore, adding seed to LoraConfig for reproducibility is a great enhancement. I have a couple of minor suggestions to further improve code style and maintainability, but overall, this is a high-quality change.