Skip to content

ACESTEP_DTYPE #1183

Description

@MRafStudio

ACESTEP_DTYPE

You have added a value to the environment, but you are not checking in init_service_loader.py and init_service_orchestrator.py
My suggestion of a fix in init_service_loader.py:

elif device == "cuda" and not gpu_config.cuda_supports_bfloat16():
    # Check if using float32 (manual override or future auto-detection)
    if getattr(self, "dtype", None) == torch.float32:
        # float32 doesn't need eager attention workaround — SDPA is stable
        logger.info(
            "[initialize_service] float32 detected on Pre-Ampere CUDA: "
            "using SDPA (eager attention not needed for float32)."
        )
        attn_implementation = "sdpa"
    else:
        # Pre-Ampere GPUs in float16 can overflow in SDPA's fused softmax
        # with longer sequences, producing NaN/Inf latents.
        # Eager attention upcasts to float32 for softmax, avoiding overflow.
        logger.info(
            "[initialize_service] Pre-Ampere CUDA detected: using eager "
            "attention for float16 numerical stability."
        )
        attn_implementation = "eager"

and in init_service_orchestrator.py:

elif resolved_device == "cuda":
    # Check manual dtype override from environment first
    env_dtype_str = os.environ.get("ACESTEP_DTYPE", "").strip().lower()
    if env_dtype_str in ("float32", "float16", "bfloat16"):
        self.dtype = getattr(torch, env_dtype_str)
        logger.info(
            f"[initialize_service] ACESTEP_DTYPE={env_dtype_str} override: "
            f"using dtype={self.dtype}."
        )
    elif gpu_config.cuda_supports_bfloat16():
        self.dtype = torch.bfloat16
    else:
        self.dtype = torch.float16
        logger.info(
            "[initialize_service] Pre-Ampere CUDA detected: "
            "using float16 instead of bfloat16."
        )

Sincerely, RAF

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions