LMFuser 0.0.1__py3-none-any.whl

lmfuser/utils.py

@@ -0,0 +1,249 @@
+import os
+from typing import Any, overload, Union, Optional, Literal
+from random import Random
+
+import torch
+from torch import distributed as dist
+from torch.utils.data import get_worker_info
+from torch import Tensor
+import atexit
+
+from typing import TypeVar
+import random
+import logging
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar('T')
+
+def get_world_size() -> int:
+    if 'WORLD_SIZE' in os.environ:
+        return int(os.environ['WORLD_SIZE'])
+    else:
+        return 1
+
+def get_global_rank() -> int:
+    if 'RANK' in os.environ:
+        return int(os.environ['RANK'])
+    else:
+        return 0
+
+def get_local_rank() -> int:
+    if "LOCAL_RANK" in os.environ:
+        return int(os.environ["LOCAL_RANK"])
+    else:
+        # Provide a default or handle the case of non-distributed training
+        return 0
+
+def dist_init() -> None:
+    if dist.is_initialized():
+        return
+
+    device_type = get_default_device_type()
+
+    if device_type == 'cuda':
+        dist.init_process_group(backend='nccl')
+        torch.cuda.set_device(get_local_rank())
+    elif device_type == 'npu':
+        dist.init_process_group(backend='hccl')
+    else:
+        dist.init_process_group(backend='mpi')
+
+    atexit.register(dist.destroy_process_group)
+
+def weighted_random_choice(
+    elements: list[T], 
+    probabilities: list[float],
+    rand: Optional[Random] = None
+) -> T:
+    """
+    Choose a random element from a list based on specified probabilities.
+
+    Args:
+        elements (List[T]): A list of elements to choose from.
+        probabilities (List[float]): A list of probabilities associated with each element. Must sum to 1.
+        rand (Optional[Random]): The random number generator.
+
+    Returns:
+        T: A randomly chosen element from the list.
+    """
+    if len(elements) != len(probabilities):
+        raise ValueError("Elements and probabilities must have the same length.")
+    if not all(0 <= p <= 1 for p in probabilities):
+        raise ValueError("Probabilities must be non-negative and non-greater than 1.")
+    if not abs(sum(probabilities) - 1.0) < 1e-6:
+        raise ValueError("Probabilities must sum to 1.")
+
+    if rand is None:
+        index = random.choices(list(range(len(elements))), probabilities)[0]
+    else:
+        index = rand.choices(list(range(len(elements))), probabilities)[0]
+
+    return elements[index]
+
+def partition_list(lst: list[T], num_shards: int, index: int) -> list[T]:
+    # Ensure the number of shards is positive and index is valid
+    if num_shards <= 0 or index >= num_shards or index < 0:
+        raise ValueError("Invalid number of shards or index, "
+                         f"number of shards: {num_shards};"
+                         f"index: {index}.")
+
+    # Calculate the size of each shard
+    shard_size = len(lst) // num_shards
+    remainder = len(lst) % num_shards
+    
+    # Calculate the start and end indices for the partition
+    start = index * shard_size + min(index, remainder)
+    end = (index + 1) * shard_size + min(index + 1, remainder)
+    
+    return lst[start:end]    
+
+DEVICE_TYPE: Optional[Literal['cuda', 'npu', 'cpu']] = None
+def get_default_device_type() -> Literal['cuda', 'npu', 'cpu']:
+    global DEVICE_TYPE
+    if DEVICE_TYPE is not None:
+        return DEVICE_TYPE
+
+    device = os.environ.get('HURRICANE_DEVICE')
+    if device is not None:
+        DEVICE_TYPE = device # type: ignore
+        return DEVICE_TYPE # type: ignore
+
+    if torch.cuda.is_available():
+        DEVICE_TYPE = 'cuda'
+        return DEVICE_TYPE
+
+    try:
+        __import__('torch_npu')
+        if torch_npu.npu.is_available(): # type: ignore
+            DEVICE_TYPE = 'npu'
+            return DEVICE_TYPE
+    except ImportError:
+        ...
+
+    DEVICE_TYPE = 'cpu'
+
+    return DEVICE_TYPE
+
+DEVICE: Optional[str] = None
+def get_default_device() -> str | int:
+    """
+    Get the default device for the current process.
+    """
+    global DEVICE
+    if DEVICE is not None:
+        return DEVICE
+
+    device_type = get_default_device_type()
+    if device_type == 'cpu':
+        return -1
+    if device_type == 'cuda':
+        return get_local_rank()
+
+    return f'{device_type}:{get_local_rank()}'
+
+@overload
+def dist_avg(value: Tensor) -> Tensor: ...
+@overload
+def dist_avg(value: int) -> float: ...
+@overload
+def dist_avg(value: float) -> float: ...
+def dist_avg(value: Union[torch.Tensor, int, float]) -> Union[torch.Tensor, float]:
+    if not dist.is_initialized():
+        return value
+    
+    if isinstance(value, torch.Tensor):
+        return_tensor = True
+        value = value.to(get_default_device())
+    else:
+        return_tensor = False
+        value = torch.tensor(value, device=get_default_device(), dtype=torch.float32)
+        
+    dist.all_reduce(value, dist.ReduceOp.SUM)
+    dist.barrier()
+
+    if return_tensor:
+        return value / dist.get_world_size()
+
+    return value.item() / dist.get_world_size()
+
+
+def gather_object(local_object: T) -> list[T]:
+    if not dist.is_initialized():
+        return [local_object]
+    world_size = dist.get_world_size()
+
+    gathered = [None for _ in range(world_size)]
+    dist.all_gather_object(gathered, local_object)
+
+    if isinstance(gathered[0], list):
+        results = []
+        for l in gathered:
+            results += l # type: ignore
+        gathered = results
+    
+    return gathered # type: ignore
+
+def tensor_all_gather(tensor: Tensor) -> Tensor:
+    if get_world_size() <= 1:
+        return tensor
+    results = [torch.empty_like(tensor, device=get_default_device()) for _ in range(dist.get_world_size())]
+
+    dist.all_gather(results, tensor.contiguous().to(get_default_device()))
+
+    return torch.cat(results, dim=0).to(tensor.device)
+
+def batch_all_gather(batch: dict[str, Any]) -> dict[str, Any]:
+    """把各个rank的同一批batch数据汇总到rank0，方便计算metric或者刷库等等。
+    不是分布式的话原样返回
+
+    Args:
+        batch (Dict[str, Union[Tensor, List[Any]]]): 一个batch的数据，注意一个key对应一个list
+
+    Returns:
+        Dict[Dict[str, Union[Tensor, List[Any]]]]: 汇总后的batch
+    """
+    logger.info(f'Begin all gather on rank {get_global_rank()}')
+    if not dist.is_initialized():
+        return batch
+
+    gathered = {}
+    for k, v_list in batch.items():
+        logger.info(f'Begin gather key {k} on rank {get_global_rank()}')
+        if isinstance(v_list, list):
+            logger.info(f'Key {k} with length {len(v_list)} on rank {get_global_rank()}')
+        logger.info(f'after barrier on rank {get_global_rank()}')
+        if isinstance(v_list, Tensor):
+            gathered[k] = tensor_all_gather(v_list)
+        else:
+            if isinstance(v_list[0], Tensor):
+                v_tensor = torch.cat(v_list, dim=0)
+                gathered[k] = tensor_all_gather(v_tensor)
+            else:
+                gathered[k] = gather_object(v_list)
+
+    return gathered
+
+def cal_acc_num(batch_size: int, sub_batch_size: int, world_size: int) -> int:
+    acc_num = batch_size // sub_batch_size // world_size
+    an = batch_size / sub_batch_size / world_size
+    if abs(acc_num - an) > 1e-8:
+        raise ValueError(
+            f'Batchsize无法被平分！请根据显卡数合理设置“BatchSize”与“SubBatchSize”'
+        )
+    return acc_num
+
+
+class MethodOverideChecker:
+    def is_overridden(self, method_name: str) -> bool:
+        """
+        Check if the method `method_name` is overridden in this instance's class
+        compared to the Parent class.
+        """
+        cls = self.__class__
+        for parent in self.__class__.__bases__:
+            if not hasattr(parent, method_name):
+                continue
+            return getattr(cls, method_name, None) is not getattr(parent, method_name, None)
+
+        raise ValueError(f'No such method {method_name} in parent classes')

lmfuser-0.0.1.dist-info/METADATA

@@ -0,0 +1,30 @@
+Metadata-Version: 2.4
+Name: LMFuser
+Version: 0.0.1
+Summary: The LMFuser training framework.
+Project-URL: Homepage, https://github.com/TYTTYTTYT/LMFuser
+Project-URL: Documentation, https://github.com/TYTTYTTYT/LMFuser
+Project-URL: Repository, https://github.com/TYTTYTTYT/LMFuser
+Project-URL: Bug Tracker, https://github.com/TYTTYTTYT/LMFuser
+Project-URL: Changelog, https://github.com/TYTTYTTYT/LMFuser
+Author-email: Yintao Tai <tai.yintao@gmail.com>
+Maintainer-email: Yintao Tai <tai.yintao@gmail.com>
+License-File: LICENSE
+Keywords: machine learning,pytorch
+Classifier: Intended Audience :: Information Technology
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Utilities
+Requires-Python: >=3.11
+Requires-Dist: hyperargs>=0.1.3
+Requires-Dist: lmfuser-data
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: pyarrow
+Requires-Dist: torch
+Requires-Dist: tqdm
+Requires-Dist: wandb
+Description-Content-Type: text/markdown
+
+# LMFuser
+A robust, multi-task PyTorch framework for training versatile Language Models at scale.

lmfuser-0.0.1.dist-info/RECORD

@@ -0,0 +1,13 @@
+lmfuser/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+lmfuser/model_loader.py,sha256=Y3Rm_hjdDYvKZuvAnblJ7URcDozOwe-QqBniyXrlzLs,1188
+lmfuser/optimizers.py,sha256=9Mp4-_WuKdEcSweSumm0eGgL8l2nUUlLbcgMPPEbYPU,5476
+lmfuser/schedulers.py,sha256=g5Can4dOOeO9YiK0NLiDZt1dAmGRXN6sFKKIcTFTGR8,8035
+lmfuser/task.py,sha256=3OBLuDnwlqe-ctWqh9WBzS0bbxRoj5gx_PibHBkWhHU,9825
+lmfuser/utils.py,sha256=lIXp1CAtTKqoCMdcU8-bSZxS36AMmBZ8a8j_xDHzoI4,7829
+lmfuser/runners/__init__.py,sha256=azfFi3bGPMDG2SEMFiIxKJgA-JGS6vuRlipkJcB7O90,51
+lmfuser/runners/ddp_runner.py,sha256=y3sIkEtzqrbVQzMy1O5iZbprGTagYCTM98uEiO7vvMU,23995
+lmfuser/runners/runner.py,sha256=L6R1Vp62mpx_73kxAByQrNfIXc-kKv6Vq5eT8kF-fhM,1297
+lmfuser-0.0.1.dist-info/METADATA,sha256=t1ielj85-3NNopKxPvj8G2ovAIBW_Q6c75kRZjLjVPQ,1104
+lmfuser-0.0.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+lmfuser-0.0.1.dist-info/licenses/LICENSE,sha256=4cDSUqshGiBlu4cAQNp8fFvQvy0N5xHX8VdJ0pzbPNU,1067
+lmfuser-0.0.1.dist-info/RECORD,,

lmfuser-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

lmfuser-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Yintao Tai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.