fix: add cleanup method for clean the processes

[zlH518]

I found that after running simple_example.py, nvidia-smi still showed that I had processes consuming memory. I suspected that the created processes were not being cleared, so I added code to clean up the processes, which solved the problem.

[zlH518]

i also find when i set the n_stage>1, there are some device error.
I fixed the device transfer logic in modeling_qwen3.py. Key improvements include: Dynamically obtaining device location: Instead of hardcoding torch.device(0) or torch.device(i+1), the target device is determined by querying the device location of the actual model parameters.
Device transfer after embedding: After performing embedding, hidden_states is immediately transferred to the device of the first layer to avoid device mismatch.
Inter-layer device transfer: When transferring between different device groups, the actual device location of the next layer is queried instead of assuming a device index.

[zheyishine]

Appreciate it!
It is better to apply your second PR to the scaffold.py file, as most of our modeling codes are auto-generated by scaffold.py instead of copy-and-revise.

[hjyai94]

After

Thanks for pointing that out! The device error happens because tie_word_embeddings is set to true in some Qwen3 config.json files. When this option is enabled, the embed_tokens and lm_head layers share the same weights.