[fix] resolve SP/CP gradient issue in fla and qwen35 ckpt conversion - #852
Conversation
- Add fallback config loading when transformers doesn't recognize the model_type (reads config.json directly and constructs a namespace) - Fix all-gather backward pass for duplicated computation (use custom autograd function that returns local gradient instead of reduce-scatter) - Add tensor_parallel_output_grad=False for sequence parallel gather - Compute layer_types from full_attention_interval when config doesn't expose it directly - Patch tied lm_head in checkpoint converter when lm_head.weight is absent from HF weights - Proactively clear memory when free GPU memory is low before distributed calls Made-with: Cursor
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces several enhancements for handling Hugging Face models and memory management. Key changes include a fallback mechanism for loading HF configurations with unknown model types, a custom all-gather operation to prevent gradient inflation during duplicated computations, and logic to handle tied lm_head weights during model conversion. Feedback highlights a significant performance concern regarding memory checks placed in the hot path of distributed calls, as well as opportunities to reduce code duplication by refactoring shared configuration loading and layer type computation logic into utility modules.
| mem_info = available_memory() | ||
| if mem_info["free_GB"] < 3: | ||
| clear_memory() |
There was a problem hiding this comment.
Calling available_memory() and potentially clear_memory() before every distributed call is extremely expensive and will severely degrade training performance. available_memory queries CUDA memory info via mem_get_info, and clear_memory performs a torch.cuda.synchronize(). Since this context manager wraps almost all collective operations (as seen in the monkey-patching logic on lines 59-81), this adds significant overhead and unnecessary synchronization points to the hot path. This check should be removed from the try block. It is already present in the except block (line 283), which is the appropriate place for diagnostic memory clearing after a failure.
| if not hasattr(text_config, "layer_types"): | ||
| interval = getattr(text_config, "full_attention_interval", 4) | ||
| n = text_config.num_hidden_layers | ||
| text_config.layer_types = [ | ||
| "full_attention" if (i + 1) % interval == 0 else "linear_attention" for i in range(n) | ||
| ] |
There was a problem hiding this comment.
The logic for computing layer_types based on full_attention_interval is duplicated here and in miles_plugins/models/qwen3_next.py. This logic should be encapsulated in a shared helper function (possibly within a shared HF config loader) to ensure consistency and reduce code duplication across model implementations.
| def _load_hf_config_with_fallback(checkpoint_path: str): | ||
| """Load HF config with fallback for model types unknown to transformers.""" | ||
| try: | ||
| from transformers import AutoConfig | ||
|
|
||
| return AutoConfig.from_pretrained(checkpoint_path, trust_remote_code=True) | ||
| except (ValueError, KeyError): | ||
| config_path = os.path.join(checkpoint_path, "config.json") | ||
| with open(config_path) as f: | ||
| config_dict = json.load(f) | ||
|
|
||
| dtype_map = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32} | ||
|
|
||
| def fix_dtype(d): | ||
| if "torch_dtype" in d: | ||
| d["torch_dtype"] = dtype_map.get(d["torch_dtype"], d["torch_dtype"]) | ||
| if "dtype" in d: | ||
| d["dtype"] = dtype_map.get(d["dtype"], d["dtype"]) | ||
|
|
||
| fix_dtype(config_dict) | ||
| ns = type("HFConfig", (), config_dict)() | ||
| if "text_config" in config_dict: | ||
| fix_dtype(config_dict["text_config"]) | ||
| ns.text_config = type("TextConfig", (), config_dict["text_config"])() | ||
| return ns |
There was a problem hiding this comment.
The function _load_hf_config_with_fallback is nearly identical to _load_hf_config introduced in miles_plugins/models/hf_attention.py. To improve maintainability and adhere to DRY principles, this logic should be moved to a shared utility module and imported in both locations.
guapisolo
changed the title
guapisolo
changed the title
|
@guapisolo LGTM functionally! Should we refactor shared utilities (get_text_config, ensure_layer_types, get_hybrid_attention_spec) from qwen3_5.py and qwen3_next.py into hybrid_utils.py to deduplicate the common hybrid attention specs? If so, I can pitch it later, merge this PR first! |
2 participants
Thanks to this PR THUDM/slime#1748.
Also this PR handle mbridge ckpt conversion issue for Qwen3.5 0.6B, 2B and 4B, which tied embeddings.
Qwen3.5-4B:
Qwen3.5-9B: