fix: compilation error

csrc/search/rpc.hpp

@@ -1,6 +1,7 @@
 #ifndef RPC_HPP
 #define RPC_HPP
 
+#include <cstdint>
 #include <string>
 #include <vector>
 #include <unordered_map>

realhf/base/cluster.py

@@ -8,8 +8,12 @@
 
 
 def get_user_tmp():
-    user = getpass.getuser()
-    user_tmp = os.path.join("/home", user, ".cache", "realhf")
+    home_dir = os.environ.get('HOME', '')
+    if not home_dir:
+        user = getpass.getuser()
+        user_tmp = os.path.join("/home", user, ".cache", "realhf")
+    else:
+        user_tmp = os.path.join(home_dir, ".cache", "realhf")
     os.makedirs(user_tmp, exist_ok=True)
     return user_tmp
 

realhf/impl/model/backend/megatron.py

@@ -8,27 +8,15 @@
 import torch.distributed as dist
 import transformers
 
-try:
-    from megatron.core import parallel_state
-    from megatron.core.distributed.distributed_data_parallel import (
-        DistributedDataParallel,
-    )
-    from megatron.core.distributed.param_and_grad_buffer import ParamAndGradBuffer
-    from megatron.core.optimizer import DistributedOptimizer, get_megatron_optimizer
-    from megatron.core.optimizer.clip_grads import clip_grad_norm_fp32, count_zeros_fp32
-    from megatron.core.optimizer.optimizer_config import OptimizerConfig
-    from megatron.core.transformer.transformer_config import TransformerConfig
-except (ModuleNotFoundError, ImportError):
-    # importing megatron.core in CPU container will fail due to the requirement of apex
-    # Here class types must be defined for type hinting
-    class TransformerConfig:
-        pass
-
-    class DistributedDataParallel:
-        pass
-
-    class DistributedOptimizer:
-        pass
+from megatron.core import parallel_state
+from megatron.core.distributed.distributed_data_parallel import (
+    DistributedDataParallel,
+)
+from megatron.core.distributed.param_and_grad_buffer import ParamAndGradBuffer
+from megatron.core.optimizer import DistributedOptimizer, get_megatron_optimizer
+from megatron.core.optimizer.clip_grads import clip_grad_norm_fp32, count_zeros_fp32
+from megatron.core.optimizer.optimizer_config import OptimizerConfig
+from megatron.core.transformer.transformer_config import TransformerConfig
 
 
 from realhf.api.core import model_api

realhf/impl/model/modules/attn.py

@@ -16,14 +16,11 @@
 from .mlp import LayerNormQKVLinear
 from .rotary import RotaryEmbedding
 
-try:
-    from flash_attn import (
-        flash_attn_func,
-        flash_attn_varlen_func,
-        flash_attn_with_kvcache,
-    )
-except ModuleNotFoundError:
-    pass
+from flash_attn import (
+    flash_attn_func,
+    flash_attn_varlen_func,
+    flash_attn_with_kvcache,
+)
 
 logger = logging.getLogger("Attention")