Fixed Memory Leaks: Added explicit cleanup of partial replicas if cloning fails mid-process
Better Error Handling: Exception handling in worker() now properly returns exceptions rather than raising, allowing cleanup of other threads
Thread Safety: Removed unused imports and ensured proper synchronization points for CUDA/XPU devices
Robust Splitting: Fixed split_batch to handle mixed lists of tensors and non-tensors correctly
Cleanup Safety: cleanup_parallel_model now handles cases where the model was partially initialized
VRAM Management: Added soft_empty_cache calls between cloning operations to prevent OOM when cloning to multiple devices
Management: Added auto balance to improve performance
Attribute Handling: Improved clone_module_simple to better handle missing attributes in layer reconstruction
Config Extraction: Added serialization check to extract_model_config to prevent storing non-serializable objects
FearL0rd