[Feature Request] Full training and LoRA support for Wan 2.1 and Wan 2.2 #362

Activity

1. entrpn commented on Mar 24, 2026

   

   entrpn

   on Mar 24, 2026

   Collaborator

   More actions

   @linwk20 we don't have plans at the moment to add LoRA training to Wan 2.1 or Wan 2.2 full training.

   We happily accept contributions and I can work with you towards getting those contributions into the repo.

   For LoRA training, tunix would be the preferred method.

   For Wan 2.2 training, it would be very similar to the Wan 2.1 training script but loading two models instead and switching them based on the high/low noise ratio.
2. harkanwalbedi commented on Sep 11, 2026

   

   harkanwalbedi

   on Sep 11, 2026

   More actions

   Hey @entrpn — saw your note about Tunix being the preferred path for LoRA training. Happy to
   pick this up if it's still open.

   @linwk20 — are you already working on this? Happy to split it up or hand off pieces if so.

   Spent some time reading through the LoRA path. The encouraging part is that the trainer barely
   needs to change. base_wan_trainer splits on nnx.Param, and everything after that is generic
   over whatever comes back — loss, value_and_grad, apply_gradients, partition specs, the
   checkpoint save. Filter on nnx.LoRAParam instead and you get the adapter-only TrainState,
   optimizer state and sharding without touching wan_trainer at all. So the actual work is
   adapter injection upstream of that split, plus export.

   The part I want your call on before writing code is Tunix. It delegates to
   qwix.apply_lora_to_model, which injects by tracing a dummy forward pass and expects a
   get_model_input() on the model that we don't have. It also turns remat off to make that trace
   work. Our configs ship with remat_policy: NONE so that's not biting today, but the block
   forward is wrapped in gradient_checkpoint.apply(...) inside nnx.scan, so anyone enabling it
   runs into the same thing. Either way, trace-based injection over scanned blocks doesn't drop in
   cleanly right now.

   So:

   1. Would you rather I make injection scan-aware upstream in qwix, or build the adapter locally
      in MaxDiffusion against nnx.LoRAParam and come back to Tunix once it handles stacked
      layouts? Related question — is taking on tunix + qwix as dependencies fine here? Neither is
      in the tree today.
   2. For 2.2, do you want both experts trained in one run by sampling timesteps across the
      boundary, or two separate runs producing the two files Wan2_2NNXLoraLoader already expects?

   Rough sequencing I had in mind: adapter plus shape tests, then the 2.1 trainer, then export
   round-tripping back through Wan2_1NNXLoraLoader, then the 2.2 dual-expert path last. The first
   three stand on their own if the last one slips.

   On testing, I'd put the adapter and injection tests on a 1-layer WanModel the way
   wan_transformer_test already does, so most of it stays reviewable at CI scale. End to end I can
   run base_wan_1_3b for 2.1. The 2.2 path is the only piece that really needs a slice, which is
   the other reason I'd leave it for last.

   If that direction looks right I'll start with a small PR for just the adapter layer and tests.
3. linwk20 commented on Sep 11, 2026

   

   linwk20

   on Sep 11, 2026

   Author

   More actions

   Hi @harkanwalbedi , thanks for offering to pick this up!

   I did originally intend to work on this using the TPUs I had access to through Google’s academic program.
   Unfortunately, about a week after my earlier reply, as I was preparing to start the implementation, Google revoked my TPU access and declined my request for continued TPU resources. Without that access, I haven’t been able to move forward with the implementation and testing.

   Please feel free to take this on. I really appreciate your help, and thanks to @entrpn for the guidance as well!

   👍React with 👍1Reacted by Harkanwal Bedi